Using aio KeyClient fails when aio BlobClient attempts to upload/download files with client-side encryption

[katagaki]

• Package Name:
  azure-storage-blob
  azure-keyvault-keys
• Package Version:
  azure-storage-blob: 12.26.0
  azure-keyvault-keys: 4.11.0
• Operating System:
  Debian Bookworm
• Python Version:
  3.11.13

Describe the bug
When using a KeyWrapper class to turn on client-side encryption using the aio variant of both the BlobClient and KeyClient, BlobClient upload/downloads fail due to key wrap functions being coroutines (marked with async).

To Reproduce
Steps to reproduce the behavior:

1. Prepare the below KeyWrapper class:

from azure.keyvault.keys import KeyVaultKey
from azure.keyvault.keys.crypto import KeyWrapAlgorithm
from azure.keyvault.keys.crypto.aio import CryptographyClient

...

class KeyWrapper:
    algorithm: KeyWrapAlgorithm
    key_vault_key: KeyVaultKey
    client: CryptographyClient

    def __init__(self, kek: KeyVaultKey, credential):
        self.algorithm = KeyWrapAlgorithm.rsa_oaep_256
        self.key_vault_key = kek
        self.client = CryptographyClient(kek, credential)

    async def wrap_key(self, key):
        wrapped = await self.client.wrap_key(key=key, algorithm=self.algorithm)
        return wrapped.encrypted_key

    async def unwrap_key(self, key, _):
        unwrapped = await self.client.unwrap_key(encrypted_key=key, algorithm=self.algorithm)
        return unwrapped.key

    def get_key_wrap_algorithm(self):
        return self.algorithm

    def get_kid(self):
        return self.key_vault_key.id

2. Turn on encryption on the BlobClient:

from azure.keyvault.keys import KeyVaultKey
from azure.keyvault.keys.aio import KeyClient
from azure.storage.blob.aio import BlobClient

...

async def enable_encryption(blob_client: BlobClient) -> BlobClient:
    credential = ... # authentication

    key_client: KeyClient = KeyClient(vault_url=key_vault_address, credential=credential)
    key_vault_key: KeyVaultKey = await key_client.get_key(key_name)
    key_encryption_key: KeyWrapper = KeyWrapper(key_vault_key, credential)

    blob_client.require_encryption = True
    blob_client.key_encryption_key = key_encryption_key
    blob_client.encryption_version = "2.0"

    return blob_client

3. Attempt to upload a block blob with the BlobClient.

Expected behavior
Azure Storage Blob Client should adapt and await the coroutine.

Actual behavior

  File \"/app/upload.py\", line 25, in upload_blob
    await blob_client.upload_blob(file, overwrite=True)
  File \"/app/.venv/lib/python3.11/site-packages/azure/core/tracing/decorator_async.py\", line 119, in wrapper_use_tracer
    return await func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File \"/app/.venv/lib/python3.11/site-packages/azure/storage/blob/aio/_blob_client_async.py\", line 624, in upload_blob
    return cast(Dict[str, Any], await upload_block_blob(**options))
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File \"/app/.venv/lib/python3.11/site-packages/azure/storage/blob/aio/_upload_helpers.py\", line 82, in upload_block_blob
    encryption_data, data = encrypt_blob(data, encryption_options['key'], encryption_options['version'])
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File \"/app/.venv/lib/python3.11/site-packages/azure/storage/blob/_encryption.py\", line 802, in encrypt_blob
    encryption_data = _generate_encryption_data_dict(
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File \"/app/.venv/lib/python3.11/site-packages/azure/storage/blob/_encryption.py\", line 513, in _generate_encryption_data_dict
    wrapped_content_key[\"EncryptedKey\"] = encode_base64(wrapped_cek)
                                          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File \"/app/.venv/lib/python3.11/site-packages/azure/storage/blob/_shared/__init__.py\", line 28, in encode_base64
    encoded = base64.b64encode(data)
              ^^^^^^^^^^^^^^^^^^^^^^
  File \"/usr/local/lib/python3.11/base64.py\", line 58, in b64encode
    encoded = binascii.b2a_base64(s, newline=False)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: a bytes-like object is required, not 'coroutine'

Additional context
Using asyncio.run was not possible in this context, as the application was running in a FastAPI event loop.

This was temporarily mitigated by implementing each function using the CryptographyClient like the below code snippet.

executor: ThreadPoolExecutor = ThreadPoolExecutor()

...

def __init__(self, kek: KeyVaultKey, credential):
    ...
    self.credential = credential
    ...

def wrap_key(self, key):
    async def awrap_key():
        async with CryptographyClient(self.key_vault_key, self.credential) as client:
            return await client.wrap_key(key=key, algorithm=self.algorithm)

    wrap_key_task = executor.submit(asyncio.run, awrap_key())
    wrapped = wrap_key_task.result()
    return wrapped.encrypted_key

[github-actions]

Thanks for the feedback! We are routing this to the appropriate team for follow-up. cc @jalauzon-msft @vincenttran-msft.

[jalauzon-msft]

Hi @katagaki, thanks for reaching out. You are correct that we do not support async KeyWrapper methods in our async clients. We can take this as a feature request to correct this, but it may be a while until we are able to get to it. In the meantime, it seems like you have a workaround, or you could consider using the synchronous CryptographyClient and keeping your KeyWrapper methods synchronous.

You could also consider whether a type of server-side encryption would work for your needs. Azure Storage offers different ways to manage your keys and still use server-side encryption, and we recommend server-side encryption for most scenarios. See here for more information if you would like.

[katagaki]

@jalauzon-msft

Thanks for clarifying!

We currently have requirements that don't allow us to rely on server side encryption or use the sync implementation, so for now we'll stick to the workaround we've implemented until the feature has been implemented in the future.

Leaving the snippet up as well for anyone that might run into a similar challenge.