Handle App Service MSI response expires_on values (Azure#5972)

chlowell · web-flow · commit 08a2262cf63b · 2019-06-21T10:38:21.000-07:00
diff --git a/sdk/identity/azure-identity/azure/identity/_authn_client.py b/sdk/identity/azure-identity/azure/identity/_authn_client.py
@@ -3,11 +3,12 @@
 # Licensed under the MIT License. See LICENSE.txt in the project root for
 # license information.
 # -------------------------------------------------------------------------
+import calendar
 import time
 
 from azure.core import Configuration, HttpRequest
 from azure.core.credentials import AccessToken
-from azure.core.pipeline import Pipeline, PipelineRequest
+from azure.core.pipeline import Pipeline
 from azure.core.pipeline.policies import ContentDecodePolicy, NetworkTraceLoggingPolicy, RetryPolicy
 from azure.core.pipeline.transport import HttpTransport, RequestsTransport
 from msal import TokenCache
@@ -19,6 +20,7 @@
 except ImportError:
     TYPE_CHECKING = False
 if TYPE_CHECKING:
+    from time import struct_time
     from typing import Any, Dict, Iterable, Mapping, Optional
     from azure.core.pipeline import PipelineResponse
     from azure.core.pipeline.policies import HTTPPolicy
@@ -46,39 +48,61 @@ def get_cached_token(self, scopes):
         return None
 
     def _deserialize_and_cache_token(self, response, scopes, request_time):
-        # type: (PipelineResponse, Iterable[str], int) -> str
+        # type: (PipelineResponse, Iterable[str], int) -> AccessToken
         try:
-            if "deserialized_data" in response.context:
-                payload = response.context["deserialized_data"]
-            else:
-                payload = response.http_response.text()
+            # ContentDecodePolicy sets this, and should have raised if it couldn't deserialize the response
+            payload = response.context["deserialized_data"]
             token = payload["access_token"]
 
-            # these values are strings in IMDS responses but msal.TokenCache requires they be integers
+            # these values are strings in some responses but msal.TokenCache requires they be integers
             # https://github.com/AzureAD/microsoft-authentication-library-for-python/pull/55
-            expires_in = int(payload.get("expires_in", 0))
-            if expires_in != 0:
-                payload["expires_in"] = expires_in
+            if "expires_in" in payload:
+                payload["expires_in"] = int(payload["expires_in"])
             if "ext_expires_in" in payload:
                 payload["ext_expires_in"] = int(payload["ext_expires_in"])
 
+            # this will raise if the payload has neither expires_on nor expires_in
+            # (which is fine because that's unexpected, especially considering the payload contains access_token)
+            expires_on = payload.get("expires_on") or payload["expires_in"] + request_time
+
+            # ensure expires_on is an int
+            try:
+                expires_on = int(expires_on)
+            except ValueError:
+                # probably an App Service MSI response, convert it to epoch seconds
+                try:
+                    t = self._parse_app_service_expires_on(expires_on)
+                    expires_on = calendar.timegm(t)
+                except ValueError:
+                    # have a token but don't know when it expires -> treat it as single-use
+                    expires_on = request_time
+
+            # now we have an int expires_on, ensure the cache entry gets it
+            payload["expires_on"] = expires_on
+
             self._cache.add({"response": payload, "scope": scopes})
 
-            # AccessToken contains the token's expires_on time. There are four cases for setting it:
-            # 1. response has expires_on -> AccessToken uses it
-            # 2. response has expires_on and expires_in -> AccessToken uses expires_on
-            # 3. response has only expires_in -> AccessToken uses expires_in + time of request
-            # 4. response has neither expires_on or expires_in -> AccessToken sets expires_on = 0
-            #    (not expecting this case; if it occurs, the token is effectively single-use)
-            expires_on = payload.get("expires_on", 0)
-            return AccessToken(token, expires_on or expires_in + request_time)
+            return AccessToken(token, expires_on)
         except KeyError:
             if "access_token" in payload:
                 payload["access_token"] = "****"
-            raise AuthenticationError("Unexpected authentication response: {}".format(payload))
+            raise AuthenticationError("Unexpected response: {}".format(payload))
         except Exception as ex:
             raise AuthenticationError("Authentication failed: {}".format(str(ex)))
 
+    @staticmethod
+    def _parse_app_service_expires_on(expires_on):
+        # type: (str) -> struct_time
+        """
+        Parse expires_on from an App Service MSI response (e.g. "06/19/2019 23:42:01 +00:00") to struct_time.
+        Expects the time is given in UTC (i.e. has offset +00:00).
+        """
+        if not expires_on.endswith(" +00:00"):
+            raise ValueError("'{}' doesn't match expected format".format(expires_on))
+
+        # parse the string minus the timezone offset
+        return time.strptime(expires_on[: -len(" +00:00")], "%m/%d/%Y %H:%M:%S")
+
     # TODO: public, factor out of request_token
     def _prepare_request(self, method="POST", headers=None, form_data=None, params=None):
         # type: (Optional[str], Optional[Mapping[str, str]], Optional[Mapping[str, str]], Optional[Dict[str, str]]) -> HttpRequest
diff --git a/sdk/identity/azure-identity/tests/test_authn_client.py b/sdk/identity/azure-identity/tests/test_authn_client.py
@@ -19,9 +19,11 @@ def test_authn_client_deserialization():
     # using the synchronous AuthnClient to drive this test but the functionality tested is
     # in the sans I/O AuthnClientBase, shared with AsyncAuthnClient
     now = 6
-    expires_in = 1800
+    expires_in = 59 - now
     expires_on = now + expires_in
-    expected_token = "token"
+    access_token = "***"
+    expected_access_token = AccessToken(access_token, expires_on)
+    scope = "scope"
 
     mock_response = Mock(
         headers={"content-type": "application/json"}, status_code=200, content_type=["application/json"]
@@ -30,29 +32,44 @@ def test_authn_client_deserialization():
 
     # response with expires_on only
     mock_response.text = lambda: json.dumps(
-        {"access_token": expected_token, "expires_on": expires_on, "token_type": "Bearer"}
+        {"access_token": access_token, "expires_on": expires_on, "token_type": "Bearer", "resource": scope}
     )
-    token = AuthnClient("http://foo", transport=Mock(send=mock_send)).request_token("scope")
-    assert token.token == expected_token
-    assert token.expires_on == expires_on
+    token = AuthnClient("http://foo", transport=Mock(send=mock_send)).request_token(scope)
+    assert token == expected_access_token
 
-    # response with expires_in and expires_on
+    # response with expires_on only and it's a datetime string (App Service MSI)
     mock_response.text = lambda: json.dumps(
-        {"access_token": expected_token, "expires_in": expires_in, "expires_on": expires_on, "token_type": "Bearer"}
+        {
+            "access_token": access_token,
+            "expires_on": "01/01/1970 00:00:{} +00:00".format(now + expires_in),
+            "token_type": "Bearer",
+            "resource": scope,
+        }
     )
-    token = AuthnClient("http://foo", transport=Mock(send=mock_send)).request_token("scope")
-    assert token.token == expected_token
-    assert token.expires_on == expires_on
+    token = AuthnClient("http://foo", transport=Mock(send=mock_send)).request_token(scope)
+    assert token == expected_access_token
 
-    # response with expires_in only
+    # response with string expires_in and expires_on (IMDS, Cloud Shell)
     mock_response.text = lambda: json.dumps(
-        {"access_token": expected_token, "expires_in": expires_in, "token_type": "Bearer"}
+        {
+            "access_token": access_token,
+            "expires_in": str(expires_in),
+            "expires_on": str(expires_on),
+            "token_type": "Bearer",
+            "resource": scope,
+        }
+    )
+    token = AuthnClient("http://foo", transport=Mock(send=mock_send)).request_token(scope)
+    assert token == expected_access_token
+
+    # response with int expires_in (AAD)
+    mock_response.text = lambda: json.dumps(
+        {"access_token": access_token, "expires_in": expires_in, "token_type": "Bearer", "ext_expires_in": expires_in}
     )
     with patch("azure.identity._authn_client.time.time") as mock_time:
         mock_time.return_value = now
-        token = AuthnClient("http://foo", transport=Mock(send=mock_send)).request_token("scope")
-        assert token.token == expected_token
-        assert token.expires_on == expires_on
+        token = AuthnClient("http://foo", transport=Mock(send=mock_send)).request_token(scope)
+        assert token == expected_access_token
 
 
 def test_caching_when_only_expires_in_set():
