feat(flo-cloud): add Azure Key Vault KMS provider#255
Conversation
Implements AzureKMS backed by azure-keyvault-keys, completing the Azure provider set alongside the recently added AzureBlobStorage. - flo_cloud/azure/key_vault.py: new AzureKMS class implementing the FloKMS interface (sign, verify, encrypt, decrypt, get_public_key_pem) using CryptographyClient with PS256 / RSA-OAEP-256 algorithms; auth mirrors AzureBlobStorage (ClientSecretCredential or DefaultAzureCredential) - flo_cloud/kms.py: routes CloudProvider.AZURE to AzureKMS in FloKmsService - flo_cloud/azure/__init__.py: exports AzureKMS - pyproject.toml: adds azure-keyvault-keys>=4.9.0 dependency - rag_ingestion/main.py: enables KMS for Azure when AZURE_KEY_VAULT_URL and AZURE_KEY_VAULT_KEY_NAME is set
📝 WalkthroughWalkthroughAdds Azure Key Vault KMS support: new Changes
Sequence Diagram(s)sequenceDiagram
participant Worker as Worker
participant FloKmsService as FloKmsService
participant AzureKMS as AzureKMS
participant AzureKeyVault as Azure Key Vault
Worker->>FloKmsService: request encrypt/decrypt/sign
FloKmsService->>AzureKMS: instantiate/forward operation
AzureKMS->>AzureKeyVault: KeyClient / CryptographyClient API call
AzureKeyVault-->>AzureKMS: ciphertext / plaintext / signature
AzureKMS-->>FloKmsService: return result
FloKmsService-->>Worker: return result
Estimated code review effort🎯 4 (Complex) | ⏱️ ~40 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@wavefront/server/background_jobs/rag_ingestion/rag_ingestion/main.py`:
- Around line 24-28: The condition enabling Azure KMS currently checks
AZURE_KEY_VAULT_URL and AZURE_KEY_VAULT_KEY_NAME with "is not None", which lets
empty strings through and causes FloKmsService/AzureKMS init errors; change the
checks in the CLOUD_PROVIDER conditional to treat empty values as missing (e.g.,
use truthiness or os.getenv(...).strip() to ensure non-empty) so
FloKmsService(...) is only constructed when both AZURE_KEY_VAULT_URL and
AZURE_KEY_VAULT_KEY_NAME are non-empty strings.
In `@wavefront/server/packages/flo_cloud/flo_cloud/azure/__init__.py`:
- Line 5: The __all__ export list in __init__.py is unsorted and triggers Ruff
RUF022; sort the list of exported symbols alphabetically so it reads
['AzureBlobStorage', 'AzureKMS', 'StorageQueue'] (update the __all__ variable
containing AzureBlobStorage, StorageQueue, AzureKMS accordingly) to satisfy the
lint rule.
In `@wavefront/server/packages/flo_cloud/flo_cloud/azure/key_vault.py`:
- Around line 94-102: The sign and verify methods currently pass raw message
bytes to CryptographyClient.sign/verify, but the Python SDK requires a
precomputed digest; update flo_cloud.azure.key_vault.KeyVaultSigner.sign and
.verify to compute the SHA-256 digest (because SignatureAlgorithm.ps256 uses
SHA-256) from the input message (e.g., message_digest =
hashlib.sha256(message).digest()) and pass that digest to
self.crypto_client.sign(...) and self.crypto_client.verify(...), using the same
algorithm selection logic and ensure hashlib is imported; keep variable names
like algorithm and result to minimize changes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6d38bbe4-2db8-4815-8e60-7ae318a7006b
⛔ Files ignored due to path filters (1)
wavefront/server/uv.lockis excluded by!**/*.lock
📒 Files selected for processing (5)
wavefront/server/background_jobs/rag_ingestion/rag_ingestion/main.pywavefront/server/packages/flo_cloud/flo_cloud/azure/__init__.pywavefront/server/packages/flo_cloud/flo_cloud/azure/key_vault.pywavefront/server/packages/flo_cloud/flo_cloud/kms.pywavefront/server/packages/flo_cloud/pyproject.toml
| def sign(self, message: bytes, **kwargs) -> bytes: | ||
| algorithm = kwargs.get('signing_algorithm', SignatureAlgorithm.ps256) | ||
| result = self.crypto_client.sign(algorithm, message) | ||
| return result.signature | ||
|
|
||
| def verify(self, message: bytes, signature: bytes, **kwargs) -> bool: | ||
| algorithm = kwargs.get('signing_algorithm', SignatureAlgorithm.ps256) | ||
| result = self.crypto_client.verify(algorithm, message, signature) | ||
| return result.is_valid |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
In azure-keyvault-keys v4.9.x, what inputs are expected by CryptographyClient.sign/verify vs sign_data/verify_data?
💡 Result:
In azure-keyvault-keys v4.9.x for Python, the CryptographyClient class provides sign/verify methods that expect a pre-computed digest (hash) as input, along with the algorithm and signature. There are no separate sign_data/verify_data methods in the Python SDK, unlike other language SDKs (.NET, Java, JavaScript) which distinguish between signing/verifying digests vs raw data (where _data variants handle hashing automatically). sign(algorithm: SignatureAlgorithm, digest: bytes, kwargs) -> SignResult - algorithm: e.g., SignatureAlgorithm.rs256 - digest: Pre-hashed bytes (e.g., SHA-256 hash of message) verify(algorithm: SignatureAlgorithm, digest: bytes, signature: bytes, kwargs) -> VerifyResult - algorithm: Verification algorithm - digest: Pre-hashed bytes matching the signature - signature: Signature bytes to verify Users must compute the digest client-side (e.g., using hashlib.sha256(message).digest) before calling sign/verify. This matches the official Python API docs for v4.9.x and source code analysis, where only digest-based signing is exposed directly on CryptographyClient. RSA key wrappers (KeyVaultRSAPrivateKey/KeyVaultRSAPublicKey, added in later versions) support raw data signing via cryptography library interfaces, but are not direct client methods.
Citations:
- 1: https://learn.microsoft.com/en-us/python/api/azure-keyvault-keys/azure.keyvault.keys.crypto.cryptographyclient?view=azure-python
- 2: https://azuresdkdocs.blob.core.windows.net/$web/python/azure-keyvault-keys/latest/_modules/azure/keyvault/keys/_client.html
- 3: https://azuresdkdocs.blob.core.windows.net/$web/python/azure-keyvault-keys/latest/_modules/azure/keyvault/keys/crypto/_models.html
- 4: [Key Vault] Implement abstract RSA key interfaces from
cryptographyAzure/azure-sdk-for-python#31996 - 5: https://azuresdkdocs.blob.core.windows.net/$web/python/azure-keyvault-keys/latest/index.html
- 6: https://learn.microsoft.com/en-us/python/api/overview/azure/keyvault-keys-readme?view=azure-python
- 7: https://learn.microsoft.com/en-us/azure/key-vault/keys/javascript-developer-guide-sign-verify-key
- 8: https://learn.microsoft.com/en-us/dotnet/api/azure.security.keyvault.keys.cryptography.cryptographyclient.signdata?view=azure-dotnet
- 9: https://learn.microsoft.com/en-us/java/api/com.azure.security.keyvault.keys.cryptography.cryptographyclient?view=azure-java-stable
- 10: https://docs.azure.cn/en-us/key-vault/keys/javascript-developer-guide-sign-verify-key
Compute digest before signing; Python SDK expects pre-hashed bytes, not raw messages.
Lines 96 and 101 pass raw message bytes to sign() and verify(), but the Azure Key Vault Python SDK (v4.9.x) CryptographyClient methods expect a pre-computed digest (hash). This will cause a runtime failure or incorrect behavior.
The Python SDK differs from other SDKs (.NET, Java, JavaScript) which provide sign_data()/verify_data() methods for raw data. In Python, you must compute the digest client-side before calling these methods:
Proposed fix
+import hashlib
+
def sign(self, message: bytes, **kwargs) -> bytes:
algorithm = kwargs.get('signing_algorithm', SignatureAlgorithm.ps256)
+ digest = hashlib.sha256(message).digest()
- result = self.crypto_client.sign(algorithm, message)
+ result = self.crypto_client.sign(algorithm, digest)
return result.signature
def verify(self, message: bytes, signature: bytes, **kwargs) -> bool:
algorithm = kwargs.get('signing_algorithm', SignatureAlgorithm.ps256)
+ digest = hashlib.sha256(message).digest()
- result = self.crypto_client.verify(algorithm, message, signature)
+ result = self.crypto_client.verify(algorithm, digest, signature)
return result.is_valid🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@wavefront/server/packages/flo_cloud/flo_cloud/azure/key_vault.py` around
lines 94 - 102, The sign and verify methods currently pass raw message bytes to
CryptographyClient.sign/verify, but the Python SDK requires a precomputed
digest; update flo_cloud.azure.key_vault.KeyVaultSigner.sign and .verify to
compute the SHA-256 digest (because SignatureAlgorithm.ps256 uses SHA-256) from
the input message (e.g., message_digest = hashlib.sha256(message).digest()) and
pass that digest to self.crypto_client.sign(...) and
self.crypto_client.verify(...), using the same algorithm selection logic and
ensure hashlib is imported; keep variable names like algorithm and result to
minimize changes.
There was a problem hiding this comment.
This is a false positive. The codebase contract is to pass a pre-hashed SHA-256 digest into sign()/verify() — not raw message bytes. AWS KMS uses MessageType='DIGEST' and GCP KMS uses kms_v1.Digest(sha256=message), both explicitly signalling digest input. The Azure SDK's CryptographyClient.sign(algorithm, digest) also expects a digest for PS256, which is what we pass. Adding hashlib.sha256(message).digest() here would double-hash the input and break compatibility with all callers (e.g. token_service.py). No change needed.
There was a problem hiding this comment.
@rootflo-hardik Thanks for clarifying. Agreed — this is a false positive. The FloKMS contract in this codebase is to pass a pre-computed SHA-256 digest into sign/verify, and Azure’s CryptographyClient.sign/verify for PS256 indeed expects a digest. Hashing again inside AzureKMS would double-hash and break existing callers. No change needed on this PR.
If helpful later, I can add a brief docstring or rename the parameter to “digest” in AzureKMS.sign/verify to make the expectation explicit, but I’ll leave it as-is per your guidance.
(•̀ᴗ•́)و thoughtful bunny
✏️ Learnings added
Learnt from: rootflo-hardik
Repo: rootflo/wavefront PR: 255
File: wavefront/server/packages/flo_cloud/flo_cloud/azure/key_vault.py:94-102
Timestamp: 2026-03-25T05:38:36.270Z
Learning: In wavefront/server/packages/flo_cloud, the FloKMS interface expects callers to pass a pre-computed SHA-256 digest to sign(message_digest: bytes) and verify(message_digest: bytes, signature: bytes). This applies across providers:
- AWS KMS uses MessageType='DIGEST'.
- GCP KMS uses kms_v1.Digest(sha256=...).
- Azure Key Vault’s CryptographyClient.sign/verify with PS256 also require a digest (not raw data).
Do not hash again inside provider adapters; callers (e.g., token_service.py) already supply the SHA-256 digest.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@wavefront/server/background_jobs/rag_ingestion/rag_ingestion/main.py`:
- Around line 19-20: The boolean expression combining CLOUD_PROVIDER and GCP env
checks is ambiguous; wrap the entire GCP branch in parentheses so precedence is
explicit—for example replace "or CLOUD_PROVIDER == 'gcp' and
(os.getenv('GCP_KMS_KEY_RING') and os.getenv('GCP_KMS_CRYPTO_KEY'))" with "or
(CLOUD_PROVIDER == 'gcp' and (os.getenv('GCP_KMS_KEY_RING') and
os.getenv('GCP_KMS_CRYPTO_KEY')))" to silence RUF021 and improve readability,
updating the expression where CLOUD_PROVIDER and the GCP env vars are evaluated
in main.py.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 5822dd9b-ef59-4301-84c2-47b7bae7923b
📒 Files selected for processing (2)
wavefront/server/background_jobs/rag_ingestion/rag_ingestion/main.pywavefront/server/packages/flo_cloud/flo_cloud/azure/__init__.py
✅ Files skipped from review due to trivial changes (1)
- wavefront/server/packages/flo_cloud/flo_cloud/azure/init.py
- Remove conditional enable_cloud_kms selector from ApplicationContainer and AuthContainer; KMS service is now always initialized via cloud_provider config - Add cloud_config section to floconsole config.ini to expose CLOUD_PROVIDER env var - Suppress verbose Azure SDK logs by setting the 'azure' logger to WARNING level
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@wavefront/server/apps/floconsole/floconsole/di/application_container.py`:
- Around line 69-71: The container now always provides kms_service which makes
TokenService's defensive checks and optional type incorrect; update
floconsole/services/token_service.py by changing the kms_service annotation from
"FloKMS | None" to "FloKMS", remove the two unreachable guards that raise
ValueError (the checks around "if self.kms_service is None" at the spots
referenced), and delete any code paths that rely on kms_service being None so
the class consistently assumes a present FloKMS; keep the providers.Singleton in
application_container.py as-is if the intended behavior is to require a
configured cloud_provider.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 7d978e9e-20d5-43cb-aed7-b766ae6a550d
📒 Files selected for processing (4)
wavefront/server/apps/floconsole/floconsole/config.iniwavefront/server/apps/floconsole/floconsole/di/application_container.pywavefront/server/modules/auth_module/auth_module/auth_container.pywavefront/server/packages/flo_cloud/flo_cloud/azure/__init__.py
✅ Files skipped from review due to trivial changes (1)
- wavefront/server/apps/floconsole/floconsole/config.ini
🚧 Files skipped from review as they are similar to previous changes (1)
- wavefront/server/packages/flo_cloud/flo_cloud/azure/init.py
| kms_service = providers.Singleton( | ||
| FloKmsService, cloud_provider=config.cloud_config.cloud_provider | ||
| ) |
There was a problem hiding this comment.
Same breaking change as auth_container.py; also creates dead code in TokenService.
This change mirrors auth_container.py and has the same issue: container initialization will fail if cloud_config.cloud_provider is not properly configured.
Additionally, the floconsole TokenService has explicit defensive guards:
- Lines 100-101:
if self.kms_service is None: raise ValueError(...) - Lines 135-136:
if self.kms_service is None: raise ValueError(...)
These guards become unreachable dead code since kms_service will never be None. The type hint kms_service: FloKMS | None in floconsole/services/token_service.py:30 is also now misleading.
Consider cleaning up the dead code and updating the type hint to FloKMS if this is the intended behavior going forward.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@wavefront/server/apps/floconsole/floconsole/di/application_container.py`
around lines 69 - 71, The container now always provides kms_service which makes
TokenService's defensive checks and optional type incorrect; update
floconsole/services/token_service.py by changing the kms_service annotation from
"FloKMS | None" to "FloKMS", remove the two unreachable guards that raise
ValueError (the checks around "if self.kms_service is None" at the spots
referenced), and delete any code paths that rely on kms_service being None so
the class consistently assumes a present FloKMS; keep the providers.Singleton in
application_container.py as-is if the intended behavior is to require a
configured cloud_provider.
* feat(flo-cloud): add Azure Key Vault KMS provider Implements AzureKMS backed by azure-keyvault-keys, completing the Azure provider set alongside the recently added AzureBlobStorage. - flo_cloud/azure/key_vault.py: new AzureKMS class implementing the FloKMS interface (sign, verify, encrypt, decrypt, get_public_key_pem) using CryptographyClient with PS256 / RSA-OAEP-256 algorithms; auth mirrors AzureBlobStorage (ClientSecretCredential or DefaultAzureCredential) - flo_cloud/kms.py: routes CloudProvider.AZURE to AzureKMS in FloKmsService - flo_cloud/azure/__init__.py: exports AzureKMS - pyproject.toml: adds azure-keyvault-keys>=4.9.0 dependency - rag_ingestion/main.py: enables KMS for Azure when AZURE_KEY_VAULT_URL and AZURE_KEY_VAULT_KEY_NAME is set * resolved review comments * resolved review comment 2 * simplify KMS service wiring and suppress Azure logs - Remove conditional enable_cloud_kms selector from ApplicationContainer and AuthContainer; KMS service is now always initialized via cloud_provider config - Add cloud_config section to floconsole config.ini to expose CLOUD_PROVIDER env var - Suppress verbose Azure SDK logs by setting the 'azure' logger to WARNING level
2 participants
Implements AzureKMS backed by azure-keyvault-keys, completing the Azure provider set alongside the recently added AzureBlobStorage.
Summary by CodeRabbit
New Features
Changes
Chores