Only cache desirable data in http cache

Author: rayluo

msal/throttled_http_client.py

@@ -3,6 +3,8 @@
 
 from .individual_cache import _IndividualCache as IndividualCache
 from .individual_cache import _ExpiringMapping as ExpiringMapping
+from .oauth2cli.http import Response
+from .exceptions import MsalServiceError
 
 
 # https://datatracker.ietf.org/doc/html/rfc8628#section-3.4
@@ -67,13 +69,30 @@ def _hash(raw):
         return sha256(repr(raw).encode("utf-8")).hexdigest()
 
 
+class NormalizedResponse(Response):
+    """A http response with the shape defined in Response,
+    but contains only the data we will store in cache.
+    """
+    def __init__(self, raw_response):
+        super().__init__()
+        self.status_code = raw_response.status_code
+        self.text = raw_response.text
+        self.headers = raw_response.headers
+
+    ## Note: Don't use the following line,
+    ## because when being pickled, it will indirectly pickle the whole raw_response
+    # self.raise_for_status = raw_response.raise_for_status
+    def raise_for_status(self):
+        if self.status_code >= 400:
+            raise MsalServiceError("HTTP Error: {}".format(self.status_code))
+
+
 class ThrottledHttpClient(ThrottledHttpClientBase):
+    """A throttled http client wrapper that is tailored for MSAL."""
     def __init__(self, http_client, *, default_throttle_time=None, **kwargs):
+        """Decorate self.post() and self.get() dynamically"""
         super(ThrottledHttpClient, self).__init__(http_client, **kwargs)
-
-        _post = http_client.post  # We'll patch _post, and keep original post() intact
-
-        _post = IndividualCache(
+        self.post = IndividualCache(
             # Internal specs requires throttling on at least token endpoint,
             # here we have a generic patch for POST on all endpoints.
             mapping=self._expiring_mapping,
@@ -91,9 +110,9 @@ def __init__(self, http_client, *, default_throttle_time=None, **kwargs):
                                 _extract_data(kwargs, "username")))),  # "account" of ROPC
                     ),
             expires_in=RetryAfterParser(default_throttle_time or 5).parse,
-            )(_post)
+            )(self.post)
 
-        _post = IndividualCache(  # It covers the "UI required cache"
+        self.post = IndividualCache(  # It covers the "UI required cache"
             mapping=self._expiring_mapping,
             key_maker=lambda func, args, kwargs: "POST {} hash={} 400".format(
                 args[0],  # It is the url, typically containing authority and tenant
@@ -128,9 +147,7 @@ def __init__(self, http_client, *, default_throttle_time=None, **kwargs):
                 and "retry-after" not in set(  # Leave it to the Retry-After decorator
                     h.lower() for h in getattr(result, "headers", {}).keys())
                 else 0,
-            )(_post)
-
-        self.post = _post
+            )(self.post)
 
         self.get = IndividualCache(  # Typically those discovery GETs
             mapping=self._expiring_mapping,
@@ -140,9 +157,10 @@ def __init__(self, http_client, *, default_throttle_time=None, **kwargs):
                 ),
             expires_in=lambda result=None, **ignored:
                 3600*24 if 200 <= result.status_code < 300 else 0,
-            )(http_client.get)
+            )(self.get)
 
-    # The following 2 methods have been defined dynamically by __init__()
-    #def post(self, *args, **kwargs): pass
-    #def get(self, *args, **kwargs): pass
+    def post(self, *args, **kwargs):
+        return NormalizedResponse(super(ThrottledHttpClient, self).post(*args, **kwargs))
 
+    def get(self, *args, **kwargs):
+        return NormalizedResponse(super(ThrottledHttpClient, self).get(*args, **kwargs))

tests/test_throttled_http_client.py

@@ -1,8 +1,11 @@
 # Test cases for https://identitydivision.visualstudio.com/devex/_git/AuthLibrariesApiReview?version=GBdev&path=%2FService%20protection%2FIntial%20set%20of%20protection%20measures.md&_a=preview&anchor=common-test-cases
+import pickle
 from time import sleep
 from random import random
 import logging
-from msal.throttled_http_client import ThrottledHttpClient
+
+from msal.throttled_http_client import ThrottledHttpClient, NormalizedResponse
+
 from tests import unittest
 from tests.http_client import MinimalResponse
 
@@ -12,16 +15,19 @@
 
 
 class DummyHttpClient(object):
+    SIGNATURE = str(random()).encode("utf-8")
     def __init__(self, status_code=None, response_headers=None):
         self._status_code = status_code
         self._response_headers = response_headers
 
     def _build_dummy_response(self):
-        return MinimalResponse(
+        response = MinimalResponse(
             status_code=self._status_code,
             headers=self._response_headers,
             text=random(),  # So that we'd know whether a new response is received
             )
+        response.undesirable_data = self.__class__.SIGNATURE  # To be tested for its absence in pickled response
+        return response
 
     def post(self, url, params=None, data=None, headers=None, **kwargs):
         return self._build_dummy_response()
@@ -39,6 +45,12 @@ class CloseMethodCalled(Exception):
 
 class TestHttpDecoration(unittest.TestCase):
 
+    def assertValidResponse(self, response):
+        self.assertIsInstance(response, NormalizedResponse)
+        self.assertNotIn(
+            DummyHttpClient.SIGNATURE, pickle.dumps(response),
+            "A pickled response should not contain undesirable data")
+
     def test_throttled_http_client_should_not_alter_original_http_client(self):
         original_http_client = DummyHttpClient()
         original_get = original_http_client.get
@@ -112,13 +124,19 @@ def test_one_invalid_grant_should_block_a_similar_request(self):
         http_client = DummyHttpClient(
             status_code=400)  # It covers invalid_grant and interaction_required
         http_client = ThrottledHttpClient(http_client, http_cache=http_cache)
+
         resp1 = http_client.post("https://example.com", data={"claims": "foo"})
         logger.debug(http_cache)
+        self.assertValidResponse(resp1)
         resp1_again = http_client.post("https://example.com", data={"claims": "foo"})
+        self.assertValidResponse(resp1_again)
         self.assertEqual(resp1.text, resp1_again.text, "Should return a cached response")
+
         resp2 = http_client.post("https://example.com", data={"claims": "bar"})
+        self.assertValidResponse(resp2)
         self.assertNotEqual(resp1.text, resp2.text, "Should return a new response")
         resp2_again = http_client.post("https://example.com", data={"claims": "bar"})
+        self.assertValidResponse(resp2_again)
         self.assertEqual(resp2.text, resp2_again.text, "Should return a cached response")
 
     def test_one_foci_app_recovering_from_invalid_grant_should_also_unblock_another(self):