[Storage] Fix user_agent on method calls with encryption (#32183)

Azure · Sep 22, 2023 · eeacc10 · eeacc10
1 parent fa661e7
commit eeacc10
Show file tree

Hide file tree

Showing 8 changed files with 135 additions and 26 deletions.
diff --git a/sdk/storage/azure-storage-blob/assets.json b/sdk/storage/azure-storage-blob/assets.json
@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/storage/azure-storage-blob",
-  "Tag": "python/storage/azure-storage-blob_2fc7da368a"
+  "Tag": "python/storage/azure-storage-blob_d1c5556817"
 }
diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_encryption.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_encryption.py
@@ -286,13 +286,23 @@ def modify_user_agent_for_encryption(
     :param str encryption_version: The version of encryption being used.
     :param Dict[str, Any] request_options: The reuqest options to add the user agent override to.
     """
+    # If the user has specified user_agent_overwrite=True, don't make any modifications
+    if request_options.get('user_agent_overwrite'):
+        return
+
+    # If the feature flag is already present, don't add it again
     feature_flag = f"azstorage-clientsideencryption/{encryption_version}"
-    if feature_flag not in user_agent:
-        index = user_agent.find(f"azsdk-python-{moniker}")
-        user_agent = f"{user_agent[:index]}{feature_flag} {user_agent[index:]}"
+    if feature_flag in user_agent:
+        return
+
+    index = user_agent.find(f"azsdk-python-{moniker}")
+    user_agent = f"{user_agent[:index]}{feature_flag} {user_agent[index:]}"
+    # Since we are using user_agent_overwrite=True, we must prepend the user's user_agent if there is one
+    if request_options.get('user_agent'):
+        user_agent = f"{request_options.get('user_agent')} {user_agent}"
 
-        request_options['user_agent'] = user_agent
-        request_options['user_agent_overwrite'] = True
+    request_options['user_agent'] = user_agent
+    request_options['user_agent_overwrite'] = True
 
 
 def get_adjusted_upload_size(length: int, encryption_version: str) -> int:

diff --git a/sdk/storage/azure-storage-blob/tests/test_blob_encryption_v2.py b/sdk/storage/azure-storage-blob/tests/test_blob_encryption_v2.py
@@ -1066,9 +1066,23 @@ def test_encryption_user_agent_app_id(self, **kwargs):
         storage_account_key = kwargs.pop("storage_account_key")
 
         self._setup(storage_account_name, storage_account_key)
+        kek = KeyWrapper('key1')
+        self.enable_encryption_v2(kek)
 
         app_id = 'TestAppId'
-        kek = KeyWrapper('key1')
+        content = b'Hello World Encrypted!'
+
+        def assert_user_agent(request):
+            start = f'{app_id} azstorage-clientsideencryption/2.0 '
+            assert request.http_request.headers['User-Agent'].startswith(start)
+
+        # Test method level keyword
+        blob = self.bsc.get_blob_client(self.container_name, self._get_blob_reference())
+
+        blob.upload_blob(content, overwrite=True, raw_request_hook=assert_user_agent, user_agent=app_id)
+        blob.download_blob(raw_request_hook=assert_user_agent, user_agent=app_id).readall()
+
+        # Test client constructor level keyword
         bsc = BlobServiceClient(
             self.bsc.url,
             credential=storage_account_key,
@@ -1077,13 +1091,7 @@ def test_encryption_user_agent_app_id(self, **kwargs):
             key_encryption_key=kek,
             user_agent=app_id)
 
-        def assert_user_agent(request):
-            start = f'{app_id} azstorage-clientsideencryption/2.0 '
-            assert request.http_request.headers['User-Agent'].startswith(start)
-
         blob = bsc.get_blob_client(self.container_name, self._get_blob_reference())
-        content = b'Hello World Encrypted!'
 
-        # Act
         blob.upload_blob(content, overwrite=True, raw_request_hook=assert_user_agent)
         blob.download_blob(raw_request_hook=assert_user_agent).readall()
diff --git a/sdk/storage/azure-storage-blob/tests/test_blob_encryption_v2_async.py b/sdk/storage/azure-storage-blob/tests/test_blob_encryption_v2_async.py
@@ -1067,9 +1067,24 @@ async def test_encryption_user_agent_app_id(self, **kwargs):
         storage_account_key = kwargs.pop("storage_account_key")
 
         await self._setup(storage_account_name, storage_account_key)
+        kek = KeyWrapper('key1')
+        self.enable_encryption_v2(kek)
 
         app_id = 'TestAppId'
-        kek = KeyWrapper('key1')
+        content = b'Hello World Encrypted!'
+
+        def assert_user_agent(request):
+            start = f'{app_id} azstorage-clientsideencryption/2.0 '
+            assert request.http_request.headers['User-Agent'].startswith(start)
+
+        # Test method level keyword
+        blob = self.bsc.get_blob_client(self.container_name, self._get_blob_reference())
+
+        with mock.patch('os.urandom', mock_urandom):
+            await blob.upload_blob(content, overwrite=True, raw_request_hook=assert_user_agent, user_agent=app_id)
+        await (await blob.download_blob(raw_request_hook=assert_user_agent, user_agent=app_id)).readall()
+
+        # Test client constructor level keyword
         bsc = BlobServiceClient(
             self.bsc.url,
             credential=storage_account_key,
@@ -1078,14 +1093,8 @@ async def test_encryption_user_agent_app_id(self, **kwargs):
             key_encryption_key=kek,
             user_agent=app_id)
 
-        def assert_user_agent(request):
-            start = f'{app_id} azstorage-clientsideencryption/2.0 '
-            assert request.http_request.headers['User-Agent'].startswith(start)
-
         blob = bsc.get_blob_client(self.container_name, self._get_blob_reference())
-        content = b'Hello World Encrypted!'
 
-        # Act
         with mock.patch('os.urandom', mock_urandom):
             await blob.upload_blob(content, overwrite=True, raw_request_hook=assert_user_agent)
         await (await blob.download_blob(raw_request_hook=assert_user_agent)).readall()
diff --git a/sdk/storage/azure-storage-queue/assets.json b/sdk/storage/azure-storage-queue/assets.json
@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/storage/azure-storage-queue",
-  "Tag": "python/storage/azure-storage-queue_600330aa3b"
+  "Tag": "python/storage/azure-storage-queue_4096c52c38"
 }
diff --git a/sdk/storage/azure-storage-queue/azure/storage/queue/_encryption.py b/sdk/storage/azure-storage-queue/azure/storage/queue/_encryption.py
@@ -286,13 +286,23 @@ def modify_user_agent_for_encryption(
     :param str encryption_version: The version of encryption being used.
     :param Dict[str, Any] request_options: The reuqest options to add the user agent override to.
     """
+    # If the user has specified user_agent_overwrite=True, don't make any modifications
+    if request_options.get('user_agent_overwrite'):
+        return
+
+    # If the feature flag is already present, don't add it again
     feature_flag = f"azstorage-clientsideencryption/{encryption_version}"
-    if feature_flag not in user_agent:
-        index = user_agent.find(f"azsdk-python-{moniker}")
-        user_agent = f"{user_agent[:index]}{feature_flag} {user_agent[index:]}"
+    if feature_flag in user_agent:
+        return
+
+    index = user_agent.find(f"azsdk-python-{moniker}")
+    user_agent = f"{user_agent[:index]}{feature_flag} {user_agent[index:]}"
+    # Since we are using user_agent_overwrite=True, we must prepend the user's user_agent if there is one
+    if request_options.get('user_agent'):
+        user_agent = f"{request_options.get('user_agent')} {user_agent}"
 
-        request_options['user_agent'] = user_agent
-        request_options['user_agent_overwrite'] = True
+    request_options['user_agent'] = user_agent
+    request_options['user_agent_overwrite'] = True
 
 
 def get_adjusted_upload_size(length: int, encryption_version: str) -> int:

diff --git a/sdk/storage/azure-storage-queue/tests/test_queue_encryption.py b/sdk/storage/azure-storage-queue/tests/test_queue_encryption.py
@@ -772,6 +772,42 @@ def test_validate_encryption_v2(self, **kwargs):
         # Assert
         assert content == decrypted_data
 
+    @QueuePreparer()
+    @recorded_by_proxy
+    def test_encryption_user_agent(self, **kwargs):
+        storage_account_name = kwargs.pop("storage_account_name")
+        storage_account_key = kwargs.pop("storage_account_key")
+
+        app_id = 'TestAppId'
+        content = 'Hello World Encrypted!'
+        kek = KeyWrapper('key1')
+
+        def assert_user_agent(request):
+            start = f'{app_id} azstorage-clientsideencryption/2.0 '
+            assert request.http_request.headers['User-Agent'].startswith(start)
+
+        # Test method level keyword
+        qsc = QueueServiceClient(
+            self.account_url(storage_account_name, "queue"),
+            storage_account_key,
+            require_encryption=True,
+            encryption_version='2.0',
+            key_encryption_key=kek)
+        queue = self._create_queue(qsc)
+        queue.send_message(content, raw_request_hook=assert_user_agent, user_agent=app_id)
+
+        # Test client constructor level keyword
+        qsc = QueueServiceClient(
+            self.account_url(storage_account_name, "queue"),
+            storage_account_key,
+            require_encryption=True,
+            encryption_version='2.0',
+            key_encryption_key=kek,
+            user_agent=app_id)
+
+        queue = self._get_queue_reference(qsc)
+        queue.send_message(content, raw_request_hook=assert_user_agent)
+
 
 # ------------------------------------------------------------------------------
 if __name__ == '__main__':

diff --git a/sdk/storage/azure-storage-queue/tests/test_queue_encryption_async.py b/sdk/storage/azure-storage-queue/tests/test_queue_encryption_async.py
@@ -781,6 +781,42 @@ async def test_validate_encryption_v2(self, **kwargs):
         # Assert
         assert content == decrypted_data
 
+    @QueuePreparer()
+    @recorded_by_proxy_async
+    async def test_encryption_user_agent(self, **kwargs):
+        storage_account_name = kwargs.pop("storage_account_name")
+        storage_account_key = kwargs.pop("storage_account_key")
+
+        app_id = 'TestAppId'
+        content = 'Hello World Encrypted!'
+        kek = KeyWrapper('key1')
+
+        def assert_user_agent(request):
+            start = f'{app_id} azstorage-clientsideencryption/2.0 '
+            assert request.http_request.headers['User-Agent'].startswith(start)
+
+        # Test method level keyword
+        qsc = QueueServiceClient(
+            self.account_url(storage_account_name, "queue"),
+            storage_account_key,
+            require_encryption=True,
+            encryption_version='2.0',
+            key_encryption_key=kek)
+        queue = await self._create_queue(qsc)
+        await queue.send_message(content, raw_request_hook=assert_user_agent, user_agent=app_id)
+
+        # Test client constructor level keyword
+        qsc = QueueServiceClient(
+            self.account_url(storage_account_name, "queue"),
+            storage_account_key,
+            require_encryption=True,
+            encryption_version='2.0',
+            key_encryption_key=kek,
+            user_agent=app_id)
+
+        queue = self._get_queue_reference(qsc)
+        await queue.send_message(content, raw_request_hook=assert_user_agent)
+
 
 # ------------------------------------------------------------------------------
 if __name__ == '__main__':