[Container Registry] Acr token cache (Azure#18107)

Author: seankane-msft

sdk/containerregistry/azure-containerregistry/azure/containerregistry/_exchange_client.py

@@ -4,6 +4,7 @@
 # Licensed under the MIT License.
 # ------------------------------------
 import re
+import time
 from typing import TYPE_CHECKING
 
 from azure.core.pipeline.policies import SansIOHTTPPolicy
@@ -47,29 +48,38 @@ def __init__(self, endpoint, credential, **kwargs):
         if not endpoint.startswith("https://") and not endpoint.startswith("http://"):
             endpoint = "https://" + endpoint
         self._endpoint = endpoint
-        self._credential_scopes = "https://management.core.windows.net/.default"
+        self.credential_scope = "https://management.core.windows.net/.default"
         self._client = ContainerRegistry(
             credential=credential,
             url=endpoint,
             sdk_moniker=USER_AGENT,
             authentication_policy=ExchangeClientAuthenticationPolicy(),
-            credential_scopes=kwargs.pop("credential_scopes", self._credential_scopes),
+            credential_scopes=kwargs.pop("credential_scopes", self.credential_scope),
             **kwargs
         )
         self._credential = credential
+        self._refresh_token = None
+        self._last_refresh_time = 0
 
     def get_acr_access_token(self, challenge, **kwargs):
         # type: (str, Dict[str, Any]) -> str
         parsed_challenge = _parse_challenge(challenge)
-        refresh_token = self.exchange_aad_token_for_refresh_token(service=parsed_challenge["service"])
+        refresh_token = self.get_refresh_token(parsed_challenge["service"], **kwargs)
         return self.exchange_refresh_token_for_access_token(
             refresh_token, service=parsed_challenge["service"], scope=parsed_challenge["scope"], **kwargs
         )
 
+    def get_refresh_token(self, service, **kwargs):
+        # type: (str, **Any) -> str
+        if not self._refresh_token or time.time() - self._last_refresh_time > 300:
+            self._refresh_token = self.exchange_aad_token_for_refresh_token(service, **kwargs)
+            self._last_refresh_time = time.time()
+        return self._refresh_token
+
     def exchange_aad_token_for_refresh_token(self, service=None, **kwargs):
         # type: (str, Dict[str, Any]) -> str
         refresh_token = self._client.authentication.exchange_aad_access_token_for_acr_refresh_token(
-            service=service, access_token=self._credential.get_token(self._credential_scopes).token, **kwargs
+            service=service, access_token=self._credential.get_token(self.credential_scope).token, **kwargs
         )
         return refresh_token.refresh_token
 

sdk/containerregistry/azure-containerregistry/azure/containerregistry/aio/_async_exchange_client.py

@@ -4,29 +4,27 @@
 # Licensed under the MIT License.
 # ------------------------------------
 import re
-from typing import TYPE_CHECKING
+import time
+from typing import TYPE_CHECKING, Dict, List, Any
 
+from azure.core.pipeline import PipelineRequest, PipelineResponse
 from azure.core.pipeline.policies import SansIOHTTPPolicy
 
 from .._generated.aio import ContainerRegistry
 from .._helpers import _parse_challenge
 from .._user_agent import USER_AGENT
 
 if TYPE_CHECKING:
-    from typing import Dict, List, Any
-    from azure.core.credentials import TokenCredential
-    from azure.core.pipeline import PipelineRequest, PipelineResponse
+    from azure.core.credentials_async import AsyncTokenCredential
 
 
 class ExchangeClientAuthenticationPolicy(SansIOHTTPPolicy):
     """Authentication policy for exchange client that does not modify the request"""
 
-    def on_request(self, request):
-        # type: (PipelineRequest) -> None
+    def on_request(self, request: PipelineRequest) -> None:
         pass
 
-    def on_response(self, request, response):
-        # type: (PipelineRequest, PipelineResponse) -> None
+    def on_response(self, request: PipelineRequest, response: PipelineResponse) -> None:
         pass
 
 
@@ -42,40 +40,48 @@ class ACRExchangeClient(object):
     BEARER = "Bearer"
     AUTHENTICATION_CHALLENGE_PARAMS_PATTERN = re.compile('(?:(\\w+)="([^""]*)")+')
 
-    def __init__(self, endpoint, credential, **kwargs):
-        # type: (str, TokenCredential, Dict[str, Any]) -> None
+    def __init__(
+        self, endpoint: str, credential: "AsyncTokencredential", **kwargs: Dict[str, Any]
+    ) -> None:
         if not endpoint.startswith("https://") and not endpoint.startswith("http://"):
             endpoint = "https://" + endpoint
         self._endpoint = endpoint
-        self._credential_scopes = "https://management.core.windows.net/.default"
+        self._credential_scope = "https://management.core.windows.net/.default"
         self._client = ContainerRegistry(
             credential=credential,
             url=endpoint,
             sdk_moniker=USER_AGENT,
             authentication_policy=ExchangeClientAuthenticationPolicy(),
-            credential_scopes=kwargs.pop("credential_scopes", self._credential_scopes),
+            credential_scopes=kwargs.pop("credential_scopes", self._credential_scope),
             **kwargs
         )
         self._credential = credential
+        self._refresh_token = None
+        self._last_refresh_time = None
 
-    async def get_acr_access_token(self, challenge, **kwargs):
-        # type: (str) -> str
+    async def get_acr_access_token(self, challenge: str, **kwargs: Dict[str, Any]) -> str:
         parsed_challenge = _parse_challenge(challenge)
-        refresh_token = await self.exchange_aad_token_for_refresh_token(service=parsed_challenge["service"])
+        refresh_token = await self.get_refresh_token(parsed_challenge["service"], **kwargs)
         return await self.exchange_refresh_token_for_access_token(
             refresh_token, service=parsed_challenge["service"], scope=parsed_challenge["scope"], **kwargs
         )
 
-    async def exchange_aad_token_for_refresh_token(self, service=None, **kwargs):
-        # type: (str, Dict[str, Any]) -> str
-        token = await self._credential.get_token(self._credential_scopes)
+    async def get_refresh_token(self, service: str, **kwargs: Dict[str, Any]) -> str:
+        if not self._refresh_token or time.time() - self._last_refresh_time > 300:
+            self._refresh_token = await self.exchange_aad_token_for_refresh_token(service, **kwargs)
+            self._last_refresh_time = time.time()
+        return self._refresh_token
+
+    async def exchange_aad_token_for_refresh_token(self, service: str = None, **kwargs: Dict[str, Any]) -> str:
+        token = await self._credential.get_token(self._credential_scope)
         refresh_token = await self._client.authentication.exchange_aad_access_token_for_acr_refresh_token(
             service, token.token, **kwargs
         )
         return refresh_token.refresh_token
 
-    async def exchange_refresh_token_for_access_token(self, refresh_token, service=None, scope=None, **kwargs):
-        # type: (str, str, str, Dict[str, Any]) -> str
+    async def exchange_refresh_token_for_access_token(
+        self, refresh_token: str, service: str = None, scope: str = None, **kwargs: Dict[str, Any]
+    ) -> str:
         access_token = await self._client.authentication.exchange_acr_refresh_token_for_acr_access_token(
             service, scope, refresh_token, **kwargs
         )
@@ -88,8 +94,7 @@ async def __aenter__(self):
     async def __aexit__(self, *args):
         self._client.__aexit__(*args)
 
-    async def close(self):
-        # type: () -> None
+    async def close(self) -> None:
         """Close sockets opened by the client.
         Calling this method is unnecessary when using the client as a context manager.
         """