Skip to content

Commit e8eaf04

Browse files
jalauzon-msftjalauzon-msft
authored
[Storage] Modify UserAgent for requests using client-side encryption (Azure#31333)
1 parent 02a2df4 commit e8eaf04

File tree

12 files changed

+271
-20
lines changed

12 files changed

+271
-20
lines changed

‎sdk/storage/azure-storage-blob/assets.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,5 @@
22
"AssetsRepo": "Azure/azure-sdk-assets",
33
"AssetsRepoPrefixPath": "python",
44
"TagPrefix": "python/storage/azure-storage-blob",
5-
"Tag": "python/storage/azure-storage-blob_a27e1690db"
5+
"Tag": "python/storage/azure-storage-blob_2fc7da368a"
66
}

‎sdk/storage/azure-storage-blob/azure/storage/blob/_blob_client.py‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,9 @@
2525
from ._shared.uploads import IterStreamer
2626
from ._shared.uploads_async import AsyncIterStreamer
2727
from ._shared.request_handlers import (
28-
add_metadata_headers, get_length, read_length,
28+
add_metadata_headers,
29+
get_length,
30+
read_length,
2931
validate_and_format_range_headers)
3032
from ._shared.response_handlers import return_response_headers, process_storage_error, return_headers_and_deserialized
3133
from ._generated import AzureBlobStorage
@@ -55,7 +57,7 @@
5557
deserialize_pipeline_response_into_cls
5658
)
5759
from ._download import StorageStreamDownloader
58-
from ._encryption import StorageEncryptionMixin
60+
from ._encryption import modify_user_agent_for_encryption, StorageEncryptionMixin
5961
from ._lease import BlobLeaseClient
6062
from ._models import BlobType, BlobBlock, BlobProperties, BlobQueryError, QuickQueryDialect, \
6163
DelimitedJsonDialect, DelimitedTextDialect, PageRangePaged, PageRange
@@ -424,6 +426,13 @@ def _upload_blob_options( # pylint:disable=too-many-statements
424426
kwargs['blob_settings'] = self._config
425427
kwargs['max_concurrency'] = max_concurrency
426428
kwargs['encryption_options'] = encryption_options
429+
# Add feature flag to user agent for encryption
430+
if self.key_encryption_key:
431+
modify_user_agent_for_encryption(
432+
self._config.user_agent_policy.user_agent,
433+
self._sdk_moniker,
434+
self.encryption_version,
435+
kwargs)
427436

428437
if blob_type == BlobType.BlockBlob:
429438
kwargs['client'] = self._client.block_blob
@@ -747,7 +756,7 @@ def upload_blob(
747756

748757
def _download_blob_options(self, offset=None, length=None, encoding=None, **kwargs):
749758
# type: (Optional[int], Optional[int], Optional[str], **Any) -> Dict[str, Any]
750-
if self.require_encryption and not self.key_encryption_key:
759+
if self.require_encryption and not (self.key_encryption_key or self.key_resolver_function):
751760
raise ValueError("Encryption required but no key was provided.")
752761
if length is not None and offset is None:
753762
raise ValueError("Offset value must not be None if length is set.")
@@ -767,6 +776,14 @@ def _download_blob_options(self, offset=None, length=None, encoding=None, **kwar
767776
cpk_info = CpkInfo(encryption_key=cpk.key_value, encryption_key_sha256=cpk.key_hash,
768777
encryption_algorithm=cpk.algorithm)
769778

779+
# Add feature flag to user agent for encryption
780+
if self.key_encryption_key or self.key_resolver_function:
781+
modify_user_agent_for_encryption(
782+
self._config.user_agent_policy.user_agent,
783+
self._sdk_moniker,
784+
self.encryption_version,
785+
kwargs)
786+
770787
options = {
771788
'clients': self._client,
772789
'config': self._config,

‎sdk/storage/azure-storage-blob/azure/storage/blob/_encryption.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
# pylint: disable=too-many-lines
12
# -------------------------------------------------------------------------
23
# Copyright (c) Microsoft Corporation. All rights reserved.
34
# Licensed under the MIT License. See License.txt in the project root for
@@ -270,6 +271,30 @@ def is_encryption_v2(encryption_data: Optional[_EncryptionData]) -> bool:
270271
return encryption_data and encryption_data.encryption_agent.protocol == _ENCRYPTION_PROTOCOL_V2
271272

272273

274+
def modify_user_agent_for_encryption(
275+
user_agent: str,
276+
moniker: str,
277+
encryption_version: str,
278+
request_options: Dict[str, Any]
279+
) -> None:
280+
"""
281+
Modifies the request options to contain a user agent string updated with encryption information.
282+
Adds azstorage-clientsideencryption/<version> immediately proceeding the SDK descriptor.
283+
284+
:param str user_agent: The existing User Agent to modify.
285+
:param str moniker: The specific SDK moniker. The modification will immediately proceed azsdk-python-{moniker}.
286+
:param str encryption_version: The version of encryption being used.
287+
:param Dict[str, Any] request_options: The reuqest options to add the user agent override to.
288+
"""
289+
feature_flag = f"azstorage-clientsideencryption/{encryption_version}"
290+
if feature_flag not in user_agent:
291+
index = user_agent.find(f"azsdk-python-{moniker}")
292+
user_agent = f"{user_agent[:index]}{feature_flag} {user_agent[index:]}"
293+
294+
request_options['user_agent'] = user_agent
295+
request_options['user_agent_overwrite'] = True
296+
297+
273298
def get_adjusted_upload_size(length: int, encryption_version: str) -> int:
274299
"""
275300
Get the adjusted size of the blob upload which accounts for

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,8 @@ def __init__(
106106
primary_hostname = (parsed_url.netloc + parsed_url.path).rstrip('/')
107107
self._hosts = {LocationMode.PRIMARY: primary_hostname, LocationMode.SECONDARY: secondary_hostname}
108108

109-
self._config, self._pipeline = self._create_pipeline(self.credential, storage_sdk=service, **kwargs)
109+
self._sdk_moniker = f"storage-{service}/{VERSION}"
110+
self._config, self._pipeline = self._create_pipeline(self.credential, sdk_moniker=self._sdk_moniker, **kwargs)
110111

111112
def __enter__(self):
112113
self._client.__enter__()
@@ -318,6 +319,7 @@ def _batch_send(
318319
except HttpResponseError as error:
319320
process_storage_error(error)
320321

322+
321323
class TransportWrapper(HttpTransport):
322324
"""Wrapper class that ensures that an inner client created
323325
by a `get_client` method does not close the outer transport for the parent
@@ -410,8 +412,7 @@ def create_configuration(**kwargs):
410412
# type: (**Any) -> Configuration
411413
config = Configuration(**kwargs)
412414
config.headers_policy = StorageHeadersPolicy(**kwargs)
413-
config.user_agent_policy = UserAgentPolicy(
414-
sdk_moniker=f"storage-{kwargs.pop('storage_sdk')}/{VERSION}", **kwargs)
415+
config.user_agent_policy = UserAgentPolicy(sdk_moniker=kwargs.pop('sdk_moniker'), **kwargs)
415416
config.retry_policy = kwargs.get("retry_policy") or ExponentialRetry(**kwargs)
416417
config.logging_policy = StorageLoggingPolicy(**kwargs)
417418
config.proxy_policy = ProxyPolicy(**kwargs)

‎sdk/storage/azure-storage-blob/tests/test_blob_encryption_v2.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,9 @@ def test_encryption_kek_resolver(self, **kwargs):
238238
# Act
239239
self.bsc.key_encryption_key = None
240240
blob.upload_blob(content, overwrite=True)
241+
242+
# Set kek to None to test only resolver for download
243+
blob.key_encryption_key = None
241244
data = blob.download_blob().readall()
242245

243246
# Assert
@@ -1033,3 +1036,54 @@ def test_get_blob_large_blob(self, **kwargs):
10331036

10341037
# Assert
10351038
assert content == data
1039+
1040+
@BlobPreparer()
1041+
@recorded_by_proxy
1042+
@mock.patch('os.urandom', mock_urandom)
1043+
def test_encryption_user_agent(self, **kwargs):
1044+
storage_account_name = kwargs.pop("storage_account_name")
1045+
storage_account_key = kwargs.pop("storage_account_key")
1046+
1047+
self._setup(storage_account_name, storage_account_key)
1048+
kek = KeyWrapper('key1')
1049+
self.enable_encryption_v2(kek)
1050+
1051+
def assert_user_agent(request):
1052+
assert request.http_request.headers['User-Agent'].startswith('azstorage-clientsideencryption/2.0 ')
1053+
1054+
blob = self.bsc.get_blob_client(self.container_name, self._get_blob_reference())
1055+
content = b'Hello World Encrypted!'
1056+
1057+
# Act
1058+
blob.upload_blob(content, overwrite=True, raw_request_hook=assert_user_agent)
1059+
blob.download_blob(raw_request_hook=assert_user_agent).readall()
1060+
1061+
@BlobPreparer()
1062+
@recorded_by_proxy
1063+
@mock.patch('os.urandom', mock_urandom)
1064+
def test_encryption_user_agent_app_id(self, **kwargs):
1065+
storage_account_name = kwargs.pop("storage_account_name")
1066+
storage_account_key = kwargs.pop("storage_account_key")
1067+
1068+
self._setup(storage_account_name, storage_account_key)
1069+
1070+
app_id = 'TestAppId'
1071+
kek = KeyWrapper('key1')
1072+
bsc = BlobServiceClient(
1073+
self.bsc.url,
1074+
credential=storage_account_key,
1075+
require_encryption=True,
1076+
encryption_version='2.0',
1077+
key_encryption_key=kek,
1078+
user_agent=app_id)
1079+
1080+
def assert_user_agent(request):
1081+
start = f'{app_id} azstorage-clientsideencryption/2.0 '
1082+
assert request.http_request.headers['User-Agent'].startswith(start)
1083+
1084+
blob = bsc.get_blob_client(self.container_name, self._get_blob_reference())
1085+
content = b'Hello World Encrypted!'
1086+
1087+
# Act
1088+
blob.upload_blob(content, overwrite=True, raw_request_hook=assert_user_agent)
1089+
blob.download_blob(raw_request_hook=assert_user_agent).readall()

‎sdk/storage/azure-storage-blob/tests/test_blob_encryption_v2_async.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -240,6 +240,9 @@ async def test_encryption_kek_resolver(self, **kwargs):
240240
self.bsc.key_encryption_key = None
241241
with mock.patch('os.urandom', mock_urandom):
242242
await blob.upload_blob(content, overwrite=True)
243+
244+
# Set kek to None to test only resolver for download
245+
blob.key_encryption_key = None
243246
data = await (await blob.download_blob()).readall()
244247

245248
# Assert
@@ -1035,3 +1038,54 @@ async def test_get_blob_large_blob(self, **kwargs):
10351038

10361039
# Assert
10371040
assert content == data
1041+
1042+
@BlobPreparer()
1043+
@recorded_by_proxy_async
1044+
async def test_encryption_user_agent(self, **kwargs):
1045+
storage_account_name = kwargs.pop("storage_account_name")
1046+
storage_account_key = kwargs.pop("storage_account_key")
1047+
1048+
await self._setup(storage_account_name, storage_account_key)
1049+
kek = KeyWrapper('key1')
1050+
self.enable_encryption_v2(kek)
1051+
1052+
def assert_user_agent(request):
1053+
assert request.http_request.headers['User-Agent'].startswith('azstorage-clientsideencryption/2.0 ')
1054+
1055+
blob = self.bsc.get_blob_client(self.container_name, self._get_blob_reference())
1056+
content = b'Hello World Encrypted!'
1057+
1058+
# Act
1059+
with mock.patch('os.urandom', mock_urandom):
1060+
await blob.upload_blob(content, overwrite=True, raw_request_hook=assert_user_agent)
1061+
await (await blob.download_blob(raw_request_hook=assert_user_agent)).readall()
1062+
1063+
@BlobPreparer()
1064+
@recorded_by_proxy_async
1065+
async def test_encryption_user_agent_app_id(self, **kwargs):
1066+
storage_account_name = kwargs.pop("storage_account_name")
1067+
storage_account_key = kwargs.pop("storage_account_key")
1068+
1069+
await self._setup(storage_account_name, storage_account_key)
1070+
1071+
app_id = 'TestAppId'
1072+
kek = KeyWrapper('key1')
1073+
bsc = BlobServiceClient(
1074+
self.bsc.url,
1075+
credential=storage_account_key,
1076+
require_encryption=True,
1077+
encryption_version='2.0',
1078+
key_encryption_key=kek,
1079+
user_agent=app_id)
1080+
1081+
def assert_user_agent(request):
1082+
start = f'{app_id} azstorage-clientsideencryption/2.0 '
1083+
assert request.http_request.headers['User-Agent'].startswith(start)
1084+
1085+
blob = bsc.get_blob_client(self.container_name, self._get_blob_reference())
1086+
content = b'Hello World Encrypted!'
1087+
1088+
# Act
1089+
with mock.patch('os.urandom', mock_urandom):
1090+
await blob.upload_blob(content, overwrite=True, raw_request_hook=assert_user_agent)
1091+
await (await blob.download_blob(raw_request_hook=assert_user_agent)).readall()

‎sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_shared/base_client.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,8 @@ def __init__(
106106
primary_hostname = (parsed_url.netloc + parsed_url.path).rstrip('/')
107107
self._hosts = {LocationMode.PRIMARY: primary_hostname, LocationMode.SECONDARY: secondary_hostname}
108108

109-
self._config, self._pipeline = self._create_pipeline(self.credential, storage_sdk=service, **kwargs)
109+
self._sdk_moniker = f"storage-{service}/{VERSION}"
110+
self._config, self._pipeline = self._create_pipeline(self.credential, sdk_moniker=self._sdk_moniker, **kwargs)
110111

111112
def __enter__(self):
112113
self._client.__enter__()
@@ -318,6 +319,7 @@ def _batch_send(
318319
except HttpResponseError as error:
319320
process_storage_error(error)
320321

322+
321323
class TransportWrapper(HttpTransport):
322324
"""Wrapper class that ensures that an inner client created
323325
by a `get_client` method does not close the outer transport for the parent
@@ -410,8 +412,7 @@ def create_configuration(**kwargs):
410412
# type: (**Any) -> Configuration
411413
config = Configuration(**kwargs)
412414
config.headers_policy = StorageHeadersPolicy(**kwargs)
413-
config.user_agent_policy = UserAgentPolicy(
414-
sdk_moniker=f"storage-{kwargs.pop('storage_sdk')}/{VERSION}", **kwargs)
415+
config.user_agent_policy = UserAgentPolicy(sdk_moniker=kwargs.pop('sdk_moniker'), **kwargs)
415416
config.retry_policy = kwargs.get("retry_policy") or ExponentialRetry(**kwargs)
416417
config.logging_policy = StorageLoggingPolicy(**kwargs)
417418
config.proxy_policy = ProxyPolicy(**kwargs)

‎sdk/storage/azure-storage-file-share/azure/storage/fileshare/_shared/base_client.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,8 @@ def __init__(
106106
primary_hostname = (parsed_url.netloc + parsed_url.path).rstrip('/')
107107
self._hosts = {LocationMode.PRIMARY: primary_hostname, LocationMode.SECONDARY: secondary_hostname}
108108

109-
self._config, self._pipeline = self._create_pipeline(self.credential, storage_sdk=service, **kwargs)
109+
self._sdk_moniker = f"storage-{service}/{VERSION}"
110+
self._config, self._pipeline = self._create_pipeline(self.credential, sdk_moniker=self._sdk_moniker, **kwargs)
110111

111112
def __enter__(self):
112113
self._client.__enter__()
@@ -318,6 +319,7 @@ def _batch_send(
318319
except HttpResponseError as error:
319320
process_storage_error(error)
320321

322+
321323
class TransportWrapper(HttpTransport):
322324
"""Wrapper class that ensures that an inner client created
323325
by a `get_client` method does not close the outer transport for the parent
@@ -410,8 +412,7 @@ def create_configuration(**kwargs):
410412
# type: (**Any) -> Configuration
411413
config = Configuration(**kwargs)
412414
config.headers_policy = StorageHeadersPolicy(**kwargs)
413-
config.user_agent_policy = UserAgentPolicy(
414-
sdk_moniker=f"storage-{kwargs.pop('storage_sdk')}/{VERSION}", **kwargs)
415+
config.user_agent_policy = UserAgentPolicy(sdk_moniker=kwargs.pop('sdk_moniker'), **kwargs)
415416
config.retry_policy = kwargs.get("retry_policy") or ExponentialRetry(**kwargs)
416417
config.logging_policy = StorageLoggingPolicy(**kwargs)
417418
config.proxy_policy = ProxyPolicy(**kwargs)

‎sdk/storage/azure-storage-queue/azure/storage/queue/_encryption.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
# pylint: disable=too-many-lines
12
# -------------------------------------------------------------------------
23
# Copyright (c) Microsoft Corporation. All rights reserved.
34
# Licensed under the MIT License. See License.txt in the project root for
@@ -270,6 +271,30 @@ def is_encryption_v2(encryption_data: Optional[_EncryptionData]) -> bool:
270271
return encryption_data and encryption_data.encryption_agent.protocol == _ENCRYPTION_PROTOCOL_V2
271272

272273

274+
def modify_user_agent_for_encryption(
275+
user_agent: str,
276+
moniker: str,
277+
encryption_version: str,
278+
request_options: Dict[str, Any]
279+
) -> None:
280+
"""
281+
Modifies the request options to contain a user agent string updated with encryption information.
282+
Adds azstorage-clientsideencryption/<version> immediately proceeding the SDK descriptor.
283+
284+
:param str user_agent: The existing User Agent to modify.
285+
:param str moniker: The specific SDK moniker. The modification will immediately proceed azsdk-python-{moniker}.
286+
:param str encryption_version: The version of encryption being used.
287+
:param Dict[str, Any] request_options: The reuqest options to add the user agent override to.
288+
"""
289+
feature_flag = f"azstorage-clientsideencryption/{encryption_version}"
290+
if feature_flag not in user_agent:
291+
index = user_agent.find(f"azsdk-python-{moniker}")
292+
user_agent = f"{user_agent[:index]}{feature_flag} {user_agent[index:]}"
293+
294+
request_options['user_agent'] = user_agent
295+
request_options['user_agent_overwrite'] = True
296+
297+
273298
def get_adjusted_upload_size(length: int, encryption_version: str) -> int:
274299
"""
275300
Get the adjusted size of the blob upload which accounts for

0 commit comments

Comments
 (0)