Merge pull request #170 from guewen/8.0-unpickle-restrict

[8.0] Prevent to unpickle globals which are not jobs

connector/jobrunner/runner.py

@@ -233,10 +233,11 @@ def _initialize(self):
             cr.execute("LISTEN connector")
 
     def select_jobs(self, where, args):
-        query = "SELECT %s, uuid, id as seq, date_created, priority, eta, state " \
-                "FROM queue_job WHERE %s" % \
-                ('channel' if self.has_channel else 'NULL',
-                 where)
+        query = ("SELECT %s, uuid, id as seq, date_created, "
+                 "priority, eta, state "
+                 "FROM queue_job WHERE %s" %
+                 ('channel' if self.has_channel else 'NULL',
+                  where))
         with closing(self.conn.cursor()) as cr:
             cr.execute(query, args)
             return list(cr.fetchall())

connector/queue/job.py

@@ -25,7 +25,8 @@
 import uuid
 import sys
 from datetime import datetime, timedelta, MINYEAR
-from pickle import loads, dumps, UnpicklingError
+from cPickle import dumps, UnpicklingError, Unpickler
+from cStringIO import StringIO
 
 import openerp
 from openerp.tools.translate import _
@@ -54,6 +55,31 @@
 _logger = logging.getLogger(__name__)
 
 
+_UNPICKLE_WHITELIST = set()
+
+
+def whitelist_unpickle_global(fn_or_class):
+    """ Allow a function or class to be used in jobs
+
+    By default, the only types allowed to be used in job arguments are:
+
+    * the builtins: str/unicode, int/long, float, bool, tuple, list, dict, None
+    * the pre-registered: datetime.datetime datetime.timedelta
+
+    If you need to use an argument in a job which is not in this whitelist,
+    you can add it by using::
+
+        whitelist_unpickle_global(fn_or_class_to_register)
+
+    """
+    _UNPICKLE_WHITELIST.add(fn_or_class)
+
+
+# register common types that might be used in job arguments
+whitelist_unpickle_global(datetime)
+whitelist_unpickle_global(timedelta)
+
+
 def _unpickle(pickled):
     """ Unpickles a string and catch all types of errors it can throw,
     to raise only NotReadableJobError in case of error.
@@ -63,9 +89,26 @@ def _unpickle(pickled):
     `loads()` may raises many types of exceptions (AttributeError,
     IndexError, TypeError, KeyError, ...). They are all catched and
     raised as `NotReadableJobError`).
+
+    Pickle could be exploited by an attacker who would write a value in a job
+    that would run arbitrary code when unpickled. This is why we set a custom
+    ``find_global`` method on the ``Unpickler``, only jobs and a whitelist of
+    classes/functions are allowed to be unpickled (plus the builtins types).
     """
+    def restricted_find_global(mod_name, fn_name):
+        __import__(mod_name)
+        mod = sys.modules[mod_name]
+        fn = getattr(mod, fn_name)
+        if not (fn in JOB_REGISTRY or fn in _UNPICKLE_WHITELIST):
+            raise UnpicklingError(
+                '{}.{} is not allowed in jobs'.format(mod_name, fn_name)
+            )
+        return fn
+
+    unpickler = Unpickler(StringIO(pickled))
+    unpickler.find_global = restricted_find_global
     try:
-        unpickled = loads(pickled)
+        unpickled = unpickler.load()
     except (StandardError, UnpicklingError):
         raise NotReadableJobError('Could not unpickle.', pickled)
     return unpickled

connector/tests/test_job.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 
+import cPickle
 import mock
 import unittest2
 from datetime import datetime, timedelta
@@ -51,6 +52,15 @@ def retryable_error_task(session):
     raise RetryableJobError('Must be retried later')
 
 
+def pickle_forbidden_function(session):
+    pass
+
+
+@job
+def pickle_allowed_function(session):
+    pass
+
+
 class TestJobs(unittest2.TestCase):
     """ Test Job """
 
@@ -310,6 +320,23 @@ def test_unpickle(self):
                          'a small cucumber preserved in vinegar, '
                          'brine, or a similar solution.')
 
+    def test_unpickle_unsafe(self):
+        """ unpickling function not decorated by @job is forbidden """
+        pickled = cPickle.dumps(pickle_forbidden_function)
+        with self.assertRaises(NotReadableJobError):
+            _unpickle(pickled)
+
+    def test_unpickle_safe(self):
+        """ unpickling function decorated by @job is allowed """
+        pickled = cPickle.dumps(pickle_allowed_function)
+        self.assertEqual(_unpickle(pickled), pickle_allowed_function)
+
+    def test_unpickle_whitelist(self):
+        """ unpickling function/class that is in the whitelist is allowed """
+        arg = datetime(2016, 2, 10)
+        pickled = cPickle.dumps(arg)
+        self.assertEqual(_unpickle(pickled), arg)
+
     def test_unpickle_not_readable(self):
         with self.assertRaises(NotReadableJobError):
             self.assertEqual(_unpickle('cucumber'))