[Storage] Support OAuth for import/export managed disks (Azure#22984)

Author: jalauzon-msft

sdk/storage/azure-storage-blob/CHANGELOG.md

@@ -1,5 +1,14 @@
 # Release History
 
+## 12.10.0b4 (2022-02-24)
+
+### Features Added
+- Updated clients to support both SAS and OAuth together.
+- Updated OAuth implementation to use the AAD scope returned in a Bearer challenge.
+
+### Bugs Fixed
+- Addressed a few `mypy` typing hint errors.
+
 ## 12.10.0b3 (2022-02-08)
 
 This version and all future versions will require Python 3.6+. Python 2.7 is no longer supported.

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/authentication.py

@@ -5,6 +5,7 @@
 # --------------------------------------------------------------------------
 
 import logging
+import re
 import sys
 
 try:
@@ -140,3 +141,38 @@ def on_request(self, request):
 
         self._add_authorization_header(request, string_to_sign)
         #logger.debug("String_to_sign=%s", string_to_sign)
+
+
+class StorageHttpChallenge(object):
+    def __init__(self, challenge):
+        """ Parses an HTTP WWW-Authentication Bearer challenge from the Storage service. """
+        if not challenge:
+            raise ValueError("Challenge cannot be empty")
+
+        self._parameters = {}
+        self.scheme, trimmed_challenge = challenge.strip().split(" ", 1)
+
+        # name=value pairs either comma or space separated with values possibly being
+        # enclosed in quotes
+        for item in re.split('[, ]', trimmed_challenge):
+            comps = item.split("=")
+            if len(comps) == 2:
+                key = comps[0].strip(' "')
+                value = comps[1].strip(' "')
+                if key:
+                    self._parameters[key] = value
+
+        # Extract and verify required parameters
+        self.authorization_uri = self._parameters.get('authorization_uri')
+        if not self.authorization_uri:
+            raise ValueError("Authorization Uri not found")
+
+        self.resource_id = self._parameters.get('resource_id')
+        if not self.resource_id:
+            raise ValueError("Resource id not found")
+
+        uri_path = urlparse(self.authorization_uri).path.lstrip("/")
+        self.tenant_id = uri_path.split("/")[0]
+
+    def get_value(self, key):
+        return self._parameters.get(key)

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py

@@ -25,30 +25,30 @@
 from azure.core.pipeline import Pipeline
 from azure.core.pipeline.transport import RequestsTransport, HttpTransport
 from azure.core.pipeline.policies import (
-    RedirectPolicy,
+    AzureSasCredentialPolicy,
     ContentDecodePolicy,
-    BearerTokenCredentialPolicy,
-    ProxyPolicy,
     DistributedTracingPolicy,
     HttpLoggingPolicy,
+    RedirectPolicy,
+    ProxyPolicy,
     UserAgentPolicy,
-    AzureSasCredentialPolicy
 )
 
-from .constants import STORAGE_OAUTH_SCOPE, SERVICE_HOST_BASE, CONNECTION_TIMEOUT, READ_TIMEOUT
+from .constants import CONNECTION_TIMEOUT, READ_TIMEOUT, SERVICE_HOST_BASE
 from .models import LocationMode
 from .authentication import SharedKeyCredentialPolicy
 from .shared_access_signature import QueryStringConstants
 from .request_handlers import serialize_batch_body, _get_batch_request_delimiter
 from .policies import (
-    StorageHeadersPolicy,
+    ExponentialRetry,
+    StorageBearerTokenCredentialPolicy,
     StorageContentValidation,
+    StorageHeadersPolicy,
+    StorageHosts,
+    StorageLoggingPolicy,
     StorageRequestHook,
     StorageResponseHook,
-    StorageLoggingPolicy,
-    StorageHosts,
     QueueMessagePolicy,
-    ExponentialRetry,
 )
 from .._version import VERSION
 from .response_handlers import process_storage_error, PartialBatchErrorException
@@ -208,18 +208,18 @@ def _format_query_string(self, sas_token, credential, snapshot=None, share_snaps
         if sas_token and isinstance(credential, AzureSasCredential):
             raise ValueError(
                 "You cannot use AzureSasCredential when the resource URI also contains a Shared Access Signature.")
-        if sas_token and not credential:
-            query_str += sas_token
-        elif is_credential_sastoken(credential):
+        if is_credential_sastoken(credential):
             query_str += credential.lstrip("?")
             credential = None
+        elif sas_token:
+            query_str += sas_token
         return query_str.rstrip("?&"), credential
 
     def _create_pipeline(self, credential, **kwargs):
         # type: (Any, **Any) -> Tuple[Configuration, Pipeline]
         self._credential_policy = None
         if hasattr(credential, "get_token"):
-            self._credential_policy = BearerTokenCredentialPolicy(credential, STORAGE_OAUTH_SCOPE)
+            self._credential_policy = StorageBearerTokenCredentialPolicy(credential)
         elif isinstance(credential, SharedKeyCredentialPolicy):
             self._credential_policy = credential
         elif isinstance(credential, AzureSasCredential):

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client_async.py

@@ -15,26 +15,25 @@
 from azure.core.async_paging import AsyncList
 from azure.core.exceptions import HttpResponseError
 from azure.core.pipeline.policies import (
-    ContentDecodePolicy,
-    AsyncBearerTokenCredentialPolicy,
     AsyncRedirectPolicy,
+    AzureSasCredentialPolicy,
+    ContentDecodePolicy,
     DistributedTracingPolicy,
     HttpLoggingPolicy,
-    AzureSasCredentialPolicy,
 )
 from azure.core.pipeline.transport import AsyncHttpTransport
 
-from .constants import STORAGE_OAUTH_SCOPE, CONNECTION_TIMEOUT, READ_TIMEOUT
+from .constants import CONNECTION_TIMEOUT, READ_TIMEOUT
 from .authentication import SharedKeyCredentialPolicy
 from .base_client import create_configuration
 from .policies import (
     StorageContentValidation,
-    StorageRequestHook,
-    StorageHosts,
     StorageHeadersPolicy,
+    StorageHosts,
+    StorageRequestHook,
     QueueMessagePolicy
 )
-from .policies_async import AsyncStorageResponseHook
+from .policies_async import AsyncStorageBearerTokenCredentialPolicy, AsyncStorageResponseHook
 
 from .response_handlers import process_storage_error, PartialBatchErrorException
 
@@ -70,7 +69,7 @@ def _create_pipeline(self, credential, **kwargs):
         # type: (Any, **Any) -> Tuple[Configuration, Pipeline]
         self._credential_policy = None
         if hasattr(credential, 'get_token'):
-            self._credential_policy = AsyncBearerTokenCredentialPolicy(credential, STORAGE_OAUTH_SCOPE)
+            self._credential_policy = AsyncStorageBearerTokenCredentialPolicy(credential)
         elif isinstance(credential, SharedKeyCredentialPolicy):
             self._credential_policy = credential
         elif isinstance(credential, AzureSasCredential):

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/constants.py

@@ -22,6 +22,7 @@
     # 4000MB (max block size)/ 50KB/s (an arbitrarily chosen minimum upload speed)
     READ_TIMEOUT = 80000
 
+DEFAULT_OAUTH_SCOPE = "/.default"
 STORAGE_OAUTH_SCOPE = "https://storage.azure.com/.default"
 
 SERVICE_HOST_BASE = 'core.windows.net'

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/policies.py

@@ -30,14 +30,17 @@
     )
 
 from azure.core.pipeline.policies import (
+    BearerTokenCredentialPolicy,
     HeadersPolicy,
-    SansIOHTTPPolicy,
-    NetworkTraceLoggingPolicy,
     HTTPPolicy,
-    RequestHistory
+    NetworkTraceLoggingPolicy,
+    RequestHistory,
+    SansIOHTTPPolicy,
 )
 from azure.core.exceptions import AzureError, ServiceRequestError, ServiceResponseError
 
+from .authentication import StorageHttpChallenge
+from .constants import DEFAULT_OAUTH_SCOPE, STORAGE_OAUTH_SCOPE
 from .models import LocationMode
 
 try:
@@ -46,6 +49,7 @@
     _unicode_type = str
 
 if TYPE_CHECKING:
+    from azure.core.credentials import TokenCredential
     from azure.core.pipeline import PipelineRequest, PipelineResponse
 
 
@@ -292,26 +296,36 @@ def __init__(self, **kwargs):  # pylint: disable=unused-argument
 
     def send(self, request):
         # type: (PipelineRequest) -> PipelineResponse
-        data_stream_total = request.context.get('data_stream_total') or \
-            request.context.options.pop('data_stream_total', None)
-        download_stream_current = request.context.get('download_stream_current') or \
-            request.context.options.pop('download_stream_current', None)
-        upload_stream_current = request.context.get('upload_stream_current') or \
-            request.context.options.pop('upload_stream_current', None)
+        # Values could be 0
+        data_stream_total = request.context.get('data_stream_total')
+        if data_stream_total is None:
+            data_stream_total = request.context.options.pop('data_stream_total', None)
+        download_stream_current = request.context.get('download_stream_current')
+        if download_stream_current is None:
+            download_stream_current = request.context.options.pop('download_stream_current', None)
+        upload_stream_current = request.context.get('upload_stream_current')
+        if upload_stream_current is None:
+            upload_stream_current = request.context.options.pop('upload_stream_current', None)
+
         response_callback = request.context.get('response_callback') or \
             request.context.options.pop('raw_response_hook', self._response_callback)
 
         response = self.next.send(request)
+
         will_retry = is_retry(response, request.context.options.get('mode'))
-        if not will_retry and download_stream_current is not None:
+        # Auth error could come from Bearer challenge, in which case this request will be made again
+        is_auth_error = response.http_response.status_code == 401
+        should_update_counts = not (will_retry or is_auth_error)
+
+        if should_update_counts and download_stream_current is not None:
             download_stream_current += int(response.http_response.headers.get('Content-Length', 0))
             if data_stream_total is None:
                 content_range = response.http_response.headers.get('Content-Range')
                 if content_range:
                     data_stream_total = int(content_range.split(' ', 1)[1].split('/', 1)[1])
                 else:
                     data_stream_total = download_stream_current
-        elif not will_retry and upload_stream_current is not None:
+        elif should_update_counts and upload_stream_current is not None:
             upload_stream_current += int(response.http_request.headers.get('Content-Length', 0))
         for pipeline_obj in [request, response]:
             pipeline_obj.context['data_stream_total'] = data_stream_total
@@ -620,3 +634,24 @@ def get_backoff_time(self, settings):
             if self.backoff > self.random_jitter_range else 0
         random_range_end = self.backoff + self.random_jitter_range
         return random_generator.uniform(random_range_start, random_range_end)
+
+
+class StorageBearerTokenCredentialPolicy(BearerTokenCredentialPolicy):
+    """ Custom Bearer token credential policy for following Storage Bearer challenges """
+
+    def __init__(self, credential, **kwargs):
+        # type: (TokenCredential, **Any) -> None
+        super(StorageBearerTokenCredentialPolicy, self).__init__(credential, STORAGE_OAUTH_SCOPE, **kwargs)
+
+    def on_challenge(self, request, response):
+        # type: (PipelineRequest, PipelineResponse) -> bool
+        try:
+            auth_header = response.http_response.headers.get("WWW-Authenticate")
+            challenge = StorageHttpChallenge(auth_header)
+        except ValueError:
+            return False
+
+        scope = challenge.resource_id + DEFAULT_OAUTH_SCOPE
+        self.authorize_request(request, scope, tenant_id=challenge.tenant_id)
+
+        return True

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/policies_async.py

@@ -10,12 +10,15 @@
 import logging
 from typing import Any, TYPE_CHECKING
 
-from azure.core.pipeline.policies import AsyncHTTPPolicy
+from azure.core.pipeline.policies import AsyncBearerTokenCredentialPolicy, AsyncHTTPPolicy
 from azure.core.exceptions import AzureError
 
+from .authentication import StorageHttpChallenge
+from .constants import DEFAULT_OAUTH_SCOPE, STORAGE_OAUTH_SCOPE
 from .policies import is_retry, StorageRetryPolicy
 
 if TYPE_CHECKING:
+    from azure.core.credentials_async import AsyncTokenCredential
     from azure.core.pipeline import PipelineRequest, PipelineResponse
 
 
@@ -44,28 +47,37 @@ def __init__(self, **kwargs):  # pylint: disable=unused-argument
 
     async def send(self, request):
         # type: (PipelineRequest) -> PipelineResponse
-        data_stream_total = request.context.get('data_stream_total') or \
-            request.context.options.pop('data_stream_total', None)
-        download_stream_current = request.context.get('download_stream_current') or \
-            request.context.options.pop('download_stream_current', None)
-        upload_stream_current = request.context.get('upload_stream_current') or \
-            request.context.options.pop('upload_stream_current', None)
+        # Values could be 0
+        data_stream_total = request.context.get('data_stream_total')
+        if data_stream_total is None:
+            data_stream_total = request.context.options.pop('data_stream_total', None)
+        download_stream_current = request.context.get('download_stream_current')
+        if download_stream_current is None:
+            download_stream_current = request.context.options.pop('download_stream_current', None)
+        upload_stream_current = request.context.get('upload_stream_current')
+        if upload_stream_current is None:
+            upload_stream_current = request.context.options.pop('upload_stream_current', None)
+
         response_callback = request.context.get('response_callback') or \
             request.context.options.pop('raw_response_hook', self._response_callback)
 
         response = await self.next.send(request)
         await response.http_response.load_body()
 
         will_retry = is_retry(response, request.context.options.get('mode'))
-        if not will_retry and download_stream_current is not None:
+        # Auth error could come from Bearer challenge, in which case this request will be made again
+        is_auth_error = response.http_response.status_code == 401
+        should_update_counts = not (will_retry or is_auth_error)
+
+        if should_update_counts and download_stream_current is not None:
             download_stream_current += int(response.http_response.headers.get('Content-Length', 0))
             if data_stream_total is None:
                 content_range = response.http_response.headers.get('Content-Range')
                 if content_range:
                     data_stream_total = int(content_range.split(' ', 1)[1].split('/', 1)[1])
                 else:
                     data_stream_total = download_stream_current
-        elif not will_retry and upload_stream_current is not None:
+        elif should_update_counts and upload_stream_current is not None:
             upload_stream_current += int(response.http_request.headers.get('Content-Length', 0))
         for pipeline_obj in [request, response]:
             pipeline_obj.context['data_stream_total'] = data_stream_total
@@ -218,3 +230,24 @@ def get_backoff_time(self, settings):
             if self.backoff > self.random_jitter_range else 0
         random_range_end = self.backoff + self.random_jitter_range
         return random_generator.uniform(random_range_start, random_range_end)
+
+
+class AsyncStorageBearerTokenCredentialPolicy(AsyncBearerTokenCredentialPolicy):
+    """ Custom Bearer token credential policy for following Storage Bearer challenges """
+
+    def __init__(self, credential, **kwargs):
+        # type: (AsyncTokenCredential, **Any) -> None
+        super(AsyncStorageBearerTokenCredentialPolicy, self).__init__(credential, STORAGE_OAUTH_SCOPE, **kwargs)
+
+    async def on_challenge(self, request, response):
+        # type: (PipelineRequest, PipelineResponse) -> bool
+        try:
+            auth_header = response.http_response.headers.get("WWW-Authenticate")
+            challenge = StorageHttpChallenge(auth_header)
+        except ValueError:
+            return False
+
+        scope = challenge.resource_id + DEFAULT_OAUTH_SCOPE
+        await self.authorize_request(request, scope, tenant_id=challenge.tenant_id)
+
+        return True

sdk/storage/azure-storage-blob/azure/storage/blob/_version.py

@@ -4,4 +4,4 @@
 # license information.
 # --------------------------------------------------------------------------
 
-VERSION = "12.10.0b3"
+VERSION = "12.10.0b4"