Add ML-KEM support for KRA key recovery#5362
Conversation
This commit implements ML-KEM decapsulation support for recovering
archived private keys from KRA storage, completing the PQC key
archival/recovery workflow.
Changes:
- StorageKeyUnit.java:
* Updated unwrap() for PrivateKey recovery to detect ML-KEM storage
cert and use CryptoUtil.decapsulateMLKEM() to recover the shared
secret from the KEM ciphertext
* Updated unwrap() for SymmetricKey recovery with same ML-KEM logic
* Both methods maintain backward compatibility with RSA/EC storage
certs by branching based on storage public key algorithm
* Added comments documenting the SEQUENCE structure semantics for
RSA/EC (wrapped session key) vs ML-KEM (KEM ciphertext)
- RecoveryService.java:
* Added note that allowEncDecrypt_recovery is not supported for
ML-KEM storage at this time (matching archival code pattern)
Recovery workflow:
1. Read privateKeyData SEQUENCE from KeyRecord
2. Extract ciphertext (first OCTET STRING) and wrapped key (second)
3. If storage cert is ML-KEM: decapsulate to get shared secret
4. If storage cert is RSA/EC: unwrap session key (existing code path)
5. Unwrap user's private key with shared secret/session key
6. Create PKCS#12 with recovered key (works for PQC keys)
The PKCS#12 creation code requires no changes as it already works
with PQC keys (verified by hsmCompatVerifyServ tool).
IDM-5473
Assisted-By: Claude
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 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 |
There was a problem hiding this comment.
Code Review
This pull request introduces support for ML-KEM (Post-Quantum Cryptography) storage certificates within the Key Recovery Authority (KRA). Key changes include updating StorageKeyUnit to handle ML-KEM decapsulation for both symmetric and private key recovery paths. Review feedback highlights a critical bug where the session key length is passed in bits instead of bytes, suggests making the algorithm string matching more robust to support variations like 'MLKEM', and recommends refactoring the duplicated decapsulation logic into a shared helper method.
| sk = CryptoUtil.decapsulateMLKEM( | ||
| storagePrivKey, | ||
| session, | ||
| params.getSkLength()); |
There was a problem hiding this comment.
The params.getSkLength() method returns the key size in bits (e.g., 128 or 256), but CryptoUtil.decapsulateMLKEM() expects the size in bytes. Passing the bit count will result in an incorrect derived key size, causing the subsequent unwrap operation to fail. The value should be divided by 8.
| sk = CryptoUtil.decapsulateMLKEM( | |
| storagePrivKey, | |
| session, | |
| params.getSkLength()); | |
| sk = CryptoUtil.decapsulateMLKEM( | |
| storagePrivKey, | |
| session, | |
| params.getSkLength() / 8); |
| sk = CryptoUtil.decapsulateMLKEM( | ||
| storagePrivKey, | ||
| session, | ||
| params.getSkLength()); |
There was a problem hiding this comment.
The params.getSkLength() method returns the key size in bits (e.g., 128 or 256), but CryptoUtil.decapsulateMLKEM() expects the size in bytes. Passing the bit count will result in an incorrect derived key size, causing the subsequent unwrap operation to fail. The value should be divided by 8.
| sk = CryptoUtil.decapsulateMLKEM( | |
| storagePrivKey, | |
| session, | |
| params.getSkLength()); | |
| sk = CryptoUtil.decapsulateMLKEM( | |
| storagePrivKey, | |
| session, | |
| params.getSkLength() / 8); |
| SymmetricKey sk; | ||
|
|
||
| // Check if storage cert is ML-KEM (PQC) | ||
| if (storageAlg != null && storageAlg.toUpperCase().startsWith("ML-KEM")) { |
There was a problem hiding this comment.
The algorithm check for ML-KEM should be more robust to match the logic used in CryptoUtil.unwrap(). It should also account for the algorithm name "MLKEM" (without the hyphen) to ensure compatibility with different provider naming conventions.
| if (storageAlg != null && storageAlg.toUpperCase().startsWith("ML-KEM")) { | |
| if (storageAlg != null && (storageAlg.toUpperCase().startsWith("ML-KEM") || storageAlg.equalsIgnoreCase("MLKEM"))) { |
| SymmetricKey sk; | ||
|
|
||
| // Check if storage cert is ML-KEM (PQC) | ||
| if (storageAlg != null && storageAlg.toUpperCase().startsWith("ML-KEM")) { |
There was a problem hiding this comment.
The algorithm check for ML-KEM should be more robust to match the logic used in CryptoUtil.unwrap(). It should also account for the algorithm name "MLKEM" (without the hyphen) to ensure compatibility with different provider naming conventions.
| if (storageAlg != null && storageAlg.toUpperCase().startsWith("ML-KEM")) { | |
| if (storageAlg != null && (storageAlg.toUpperCase().startsWith("ML-KEM") || storageAlg.equalsIgnoreCase("MLKEM"))) { |
| SymmetricKey sk; | ||
|
|
||
| // Check if storage cert is ML-KEM (PQC) | ||
| if (storageAlg != null && storageAlg.toUpperCase().startsWith("ML-KEM")) { | ||
| logger.debug("StorageKeyUnit:unwrap() Using ML-KEM storage cert for symmetric key"); | ||
|
|
||
| // (1) ML-KEM decapsulation to recover shared secret | ||
| PrivateKey storagePrivKey = getPrivateKey(); | ||
| sk = CryptoUtil.decapsulateMLKEM( | ||
| storagePrivKey, | ||
| session, | ||
| params.getSkLength()); | ||
|
|
||
| logger.debug("StorageKeyUnit:unwrap() ML-KEM decapsulation complete - recovered shared secret"); | ||
|
|
||
| } else { | ||
| // (1) RSA/EC: unwrap the session key | ||
| logger.debug("StorageKeyUnit:unwrap() Using RSA/EC storage cert for symmetric key"); | ||
| sk = unwrap_session_key(token, session, SymmetricKey.Usage.UNWRAP, params); | ||
| logger.debug("StorageKeyUnit:unwrap() session key unwrapped"); | ||
| } |
There was a problem hiding this comment.
The logic for selecting the session key recovery method (ML-KEM decapsulation vs. RSA/EC unwrapping) is duplicated in both unwrap() methods. Consider refactoring this into a private helper method to improve maintainability and ensure consistent behavior across both symmetric and private key recovery paths.
Change encapsulateMLKEM() and decapsulateMLKEM() to accept key size in bits instead of bytes, consistent with PKI convention and other CryptoUtil methods like generateKey(). The methods now convert bits to bytes internally when calling the Java KEM API, making them compatible with params.getSkLength() which returns key size in bits. IDM-6295
|



This commit implements ML-KEM decapsulation support for recovering archived private keys from KRA storage, completing the PQC key archival/recovery workflow.
Depends on #5363 (CryptoUtil ML-KEM bits fix)
Changes:
StorageKeyUnit.java:
RecoveryService.java:
Recovery workflow:
The PKCS#12 creation code requires no changes as it already works with PQC keys (verified by hsmCompatVerifyServ tool).
IDM-5473
Assisted-By: Claude