Fixed Arbitrary Code Execution by pickle.load

flink-python/pyflink/fn_execution/operation_utils.py

@@ -17,6 +17,10 @@
 ################################################################################
 import datetime
 from enum import Enum
+import pickle
+import io
+import builtins
+
 
 from typing import Any, Tuple, Dict, List
 
@@ -32,6 +36,29 @@
 _constant_num = 0
 
 
+safe_builtins = {
+    'range',
+    'complex',
+    'set',
+    'frozenset',
+    'slice',
+}
+
+class RestrictedUnpickler(pickle.Unpickler):
+
+    def find_class(self, module, name):
+        """Only allow safe classes from builtins"""
+        if module == "builtins" and name in safe_builtins:
+            return getattr(builtins, name)
+        """Forbid everything else"""
+        raise pickle.UnpicklingError("global '%s.%s' is forbidden" %
+                                     (module, name))
+
+def restricted_loads(s):
+    """Helper function analogous to pickle.loads()"""
+    return RestrictedUnpickler(io.BytesIO(s)).load()
+
+
 def wrap_pandas_result(it):
     import pandas as pd
     arrays = []
@@ -99,7 +126,7 @@ def _extract_input(args) -> Tuple[str, Dict, List]:
 
     variable_dict = {}
     user_defined_funcs = []
-
+    restricted_loads(user_defined_function_proto.payload)
     user_defined_func = pickle.loads(user_defined_function_proto.payload)
     if pandas_udaf:
         user_defined_func = PandasAggregateFunctionWrapper(user_defined_func)
@@ -218,6 +245,7 @@ def load_aggregate_function(payload):
         cls = getattr(functions, built_in_function_class_name)
         return cls()
     else:
+        restricted_loads(payload)
         return pickle.loads(payload)
 
 
@@ -232,6 +260,7 @@ def extract_data_stream_stateless_function(udf_proto):
     UserDefinedDataStreamFunction = flink_fn_execution_pb2.UserDefinedDataStreamFunction
     func = None
 
+    restricted_loads(udf_proto.payload)
     user_defined_func = pickle.loads(udf_proto.payload)
     if func_type == UserDefinedDataStreamFunction.MAP:
         func = user_defined_func.map
@@ -284,6 +313,7 @@ def wrapped_func(value):
 
 
 def extract_process_function(user_defined_function_proto, ctx):
+    restricted_loads(user_defined_function_proto.payload)
     process_function = pickle.loads(user_defined_function_proto.payload)
     process_element = process_function.process_element
 
@@ -299,6 +329,7 @@ def wrapped_process_function(value):
 
 def extract_keyed_process_function(user_defined_function_proto, ctx, on_timer_ctx,
                                    collector, keyed_state_backend):
+    restricted_loads(user_defined_function_proto.payload)
     process_function = pickle.loads(user_defined_function_proto.payload)
     process_element = process_function.process_element
     on_timer = process_function.on_timer