diff --git a/src/util_mlkem.rs b/src/util_mlkem.rs new file mode 100644 index 0000000..9ebd0ed --- /dev/null +++ b/src/util_mlkem.rs @@ -0,0 +1,43 @@ +use ml_kem::kem::Encapsulate; +use ml_kem::{Encoded, EncodedSizeUser, KemCore, MlKem1024, MlKem768}; +use rust_util::{opt_result, simple_error, XResult}; + +pub fn ml_kem_768_encapsulate(public_key: &[u8]) -> XResult<(Vec, Vec)> { + let encapsulation_key_encoded: Encoded<::EncapsulationKey> = opt_result!( + public_key.try_into(), + "Parse ML-KEM 768 encapsulation key failed: {}" + ); + let encapsulation_key = + ::EncapsulationKey::from_bytes(&encapsulation_key_encoded); + let mut rng = rand::rngs::OsRng; + let (ciphertext, shared_key) = opt_result!( + encapsulation_key.encapsulate(&mut rng), + "Encapsulate shared key failed: {:?}" + ); + Ok((ciphertext.0.to_vec(), shared_key.0.to_vec())) +} + +pub fn ml_kem_1024_encapsulate(public_key: &[u8]) -> XResult<(Vec, Vec)> { + let encapsulation_key_encoded: Encoded<::EncapsulationKey> = opt_result!( + public_key.try_into(), + "Parse ML-KEM 1024 encapsulation key failed: {}" + ); + let encapsulation_key = + ::EncapsulationKey::from_bytes(&encapsulation_key_encoded); + let mut rng = rand::rngs::OsRng; + let (ciphertext, shared_key) = opt_result!( + encapsulation_key.encapsulate(&mut rng), + "Encapsulate shared key failed: {:?}" + ); + Ok((ciphertext.0.to_vec(), shared_key.0.to_vec())) +} + +pub fn try_ml_kem_encapsulate(public_key: &[u8]) -> XResult<(Vec, Vec)> { + if let Ok((ciphertext, shared_key)) = ml_kem_768_encapsulate(public_key) { + return Ok((ciphertext, shared_key)); + } + if let Ok((ciphertext, shared_key)) = ml_kem_1024_encapsulate(public_key) { + return Ok((ciphertext, shared_key)); + } + simple_error!("Not ML-KEM 768 or ML-KEM 1024.") +}