From ee4545c2a8b01a167791035dad63bb80b118fafd Mon Sep 17 00:00:00 2001
From: Laurent Mazuel <laurent.mazuel@gmail.com>
Date: Thu, 28 May 2020 13:44:04 -0700
Subject: [PATCH 1/2] Whitelist what pickle is serializing in context

---
 .../azure/core/pipeline/__init__.py           | 18 +++++++++++
 .../tests/test_universal_pipeline.py          | 32 +++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/sdk/core/azure-core/azure/core/pipeline/__init__.py b/sdk/core/azure-core/azure/core/pipeline/__init__.py
index cec8606569ae..39aa2ce86912 100644
--- a/sdk/core/azure-core/azure/core/pipeline/__init__.py
+++ b/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
diff --git a/sdk/core/azure-core/tests/test_universal_pipeline.py b/sdk/core/azure-core/tests/test_universal_pipeline.py
index 51f735c2da09..40223f06965f 100644
--- a/sdk/core/azure-core/tests/test_universal_pipeline.py
+++ b/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={}):

From 0a9c30af11d0d0a89d60ba9634716d9368246629 Mon Sep 17 00:00:00 2001
From: Laurent Mazuel <laurent.mazuel@gmail.com>
Date: Thu, 28 May 2020 13:45:24 -0700
Subject: [PATCH 2/2] Indent

---
 sdk/core/azure-core/azure/core/pipeline/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sdk/core/azure-core/azure/core/pipeline/__init__.py b/sdk/core/azure-core/azure/core/pipeline/__init__.py
index 39aa2ce86912..baeea6db1c7a 100644
--- a/sdk/core/azure-core/azure/core/pipeline/__init__.py
+++ b/sdk/core/azure-core/azure/core/pipeline/__init__.py
@@ -91,7 +91,7 @@ def __reduce__(self):
         dict_reduced_result = list(reduced[1])
         dict_reduced_result[2] = saved_context
         reduced[1] = tuple(dict_reduced_result)
-       return tuple(reduced)
+        return tuple(reduced)
 
     def __setstate__(self, state):
         self.__dict__.update(state)