Azure / Azure/azure-sdk-for-python

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

Open
#42,216 3 comments 0 reactions 2 assignees Claimed by @vincenttran-msft View on GitHub
Client customer-reported feature-request needs-team-attention Service Attention Storage
Dominant language
Python
Stars
5.6k
Forks
3.4k
Avg merge
1d 21h
Merged PRs (30d)
193

Description

- **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:

```py
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:

```py
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.

```py
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
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.