[Storage] Add argument to perf tests to use client-side encryption (#24978)

Author: jalauzon-msft

sdk/storage/azure-storage-blob/tests/perfstress_tests/README.md

@@ -54,6 +54,7 @@ The options are available for all Blob perf tests:
 - `--max-put-size` Maximum size of data uploading in single HTTP PUT. Default is 64\*1024\*1024.
 - `--max-block-size` Maximum size of data in a block within a blob. Defaults to 4\*1024\*1024.
 - `--buffer-threshold` Minimum block size to prevent full block buffering. Defaults to 4\*1024\*1024+1.
+- `--client-encryption` The version of client-side encryption to use. Leave out for no encryption.
 
 #### List Blobs command line options
 This option is only available to the List Blobs test (T1 and T2).

sdk/storage/azure-storage-blob/tests/perfstress_tests/_test_base.py

@@ -3,15 +3,14 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-import os
 import uuid
 
 from azure_devtools.perfstress_tests import PerfStressTest
-
-from azure.core.exceptions import ResourceNotFoundError
 from azure.storage.blob import BlobServiceClient as SyncBlobServiceClient
 from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
 
+from .key_wrapper import KeyWrapper
+
 
 class _ServiceTest(PerfStressTest):
  service_client = None
@@ -25,13 +24,18 @@ def __init__(self, arguments):
  self._client_kwargs['max_single_put_size'] = self.args.max_put_size
  self._client_kwargs['max_block_size'] = self.args.max_block_size
  self._client_kwargs['min_large_block_upload_threshold'] = self.args.buffer_threshold
+ if self.args.client_encryption:
+ self.key_encryption_key = KeyWrapper()
+ self._client_kwargs['require_encryption'] = True
+ self._client_kwargs['key_encryption_key'] = self.key_encryption_key
+ self._client_kwargs['encryption_version'] = self.args.client_encryption
  # self._client_kwargs['api_version'] = '2019-02-02' # Used only for comparison with T1 legacy tests
 
  if not _ServiceTest.service_client or self.args.no_client_share:
  _ServiceTest.service_client = SyncBlobServiceClient.from_connection_string(conn_str=connection_string, **self._client_kwargs)
  _ServiceTest.async_service_client = AsyncBlobServiceClient.from_connection_string(conn_str=connection_string, **self._client_kwargs)
  self.service_client = _ServiceTest.service_client
- self.async_service_client =_ServiceTest.async_service_client
+ self.async_service_client = _ServiceTest.async_service_client
 
  async def close(self):
  await self.async_service_client.close()
@@ -43,6 +47,7 @@ def add_arguments(parser):
  parser.add_argument('--max-put-size', nargs='?', type=int, help='Maximum size of data uploading in single HTTP PUT. Defaults to 64*1024*1024', default=64*1024*1024)
  parser.add_argument('--max-block-size', nargs='?', type=int, help='Maximum size of data in a block within a blob. Defaults to 4*1024*1024', default=4*1024*1024)
  parser.add_argument('--buffer-threshold', nargs='?', type=int, help='Minimum block size to prevent full block buffering. Defaults to 4*1024*1024+1', default=4*1024*1024+1)
+ parser.add_argument('--client-encryption', nargs='?', type=str, help='The version of client-side encryption to use. Leave out for no encryption.', default=None)
  parser.add_argument('--max-concurrency', nargs='?', type=int, help='Maximum number of concurrent threads used for data transfer. Defaults to 1', default=1)
  parser.add_argument('-s', '--size', nargs='?', type=int, help='Size of data to transfer. Default is 10240.', default=10240)
  parser.add_argument('--no-client-share', action='store_true', help='Create one ServiceClient per test instance. Default is to share a single ServiceClient.', default=False)

sdk/storage/azure-storage-blob/tests/perfstress_tests/key_wrapper.py

@@ -0,0 +1,34 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import os
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.keywrap import aes_key_wrap, aes_key_unwrap
+
+
+class KeyWrapper:
+ def __init__(self, kid='local:key1'):
+ self.kek = os.urandom(32)
+ self.backend = default_backend()
+ self.kid = kid
+
+ def wrap_key(self, key, algorithm='A256KW'):
+ if algorithm == 'A256KW':
+ return aes_key_wrap(self.kek, key, self.backend)
+
+ raise ValueError("Unknown key wrap algorithm.")
+
+ def unwrap_key(self, key, algorithm):
+ if algorithm == 'A256KW':
+ return aes_key_unwrap(self.kek, key, self.backend)
+
+ raise ValueError("Unknown key wrap algorithm.")
+
+ def get_key_wrap_algorithm(self):
+ return 'A256KW'
+
+ def get_kid(self):
+ return self.kid