Whitelist what pickle is serializing in context (#11688)

* Whitelist what pickle is serializing in context

* Indent

Author: lmazuel

sdk/core/azure-core/azure/core/pipeline/__init__.py

@@ -63,6 +63,9 @@ class PipelineContext(dict):
     :param transport: The HTTP transport type.
     :param kwargs: Developer-defined keyword arguments.
     """
+    _PICKLE_CONTEXT = {
+        'deserialized_data'
+    }
 
     def __init__(self, transport, **kwargs):  # pylint: disable=super-init-not-called
         self.transport = transport
@@ -75,6 +78,21 @@ def __getstate__(self):
         del state['transport']
         return state
 
+    def __reduce__(self):
+        reduced = super(PipelineContext, self).__reduce__()
+        saved_context = {}
+        for key, value in self.items():
+            if key in self._PICKLE_CONTEXT:
+                saved_context[key] = value
+        # 1 is for from __reduce__ spec of pickle (generic args for recreation)
+        # 2 is how dict is implementing __reduce__ (dict specific)
+        # tuple are read-only, we use a list in the meantime
+        reduced = list(reduced)
+        dict_reduced_result = list(reduced[1])
+        dict_reduced_result[2] = saved_context
+        reduced[1] = tuple(dict_reduced_result)
+        return tuple(reduced)
+
     def __setstate__(self, state):
         self.__dict__.update(state)
         # Re-create the unpickable entries

sdk/core/azure-core/tests/test_universal_pipeline.py

@@ -25,6 +25,7 @@
 #
 #--------------------------------------------------------------------------
 import logging
+import pickle
 try:
     from unittest import mock
 except ImportError:
@@ -56,6 +57,37 @@
     HTTPPolicy,
 )
 
+def test_pipeline_context():
+    kwargs={
+        'stream':True,
+        'cont_token':"bla"
+    }
+    context = PipelineContext('transport', **kwargs)
+    context['foo'] = 'bar'
+    context['xyz'] = '123'
+    context['deserialized_data'] = 'marvelous'
+
+    assert context['foo'] == 'bar'
+    assert context.options == kwargs
+
+    with pytest.raises(TypeError):
+        context.clear()
+
+    with pytest.raises(TypeError):
+        context.update({})
+
+    assert context.pop('foo') == 'bar'
+    assert 'foo' not in context
+
+    serialized = pickle.dumps(context)
+
+    revived_context = pickle.loads(serialized)
+    assert revived_context.options == kwargs
+    assert revived_context.transport is None
+    assert 'deserialized_data' in revived_context
+    assert len(revived_context) == 1
+
+
 def test_request_history():
     class Non_deep_copiable(object):
         def __deepcopy__(self, memodict={}):