feat: pending ML-KEM encryption and decryption

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<u8>, Vec<u8>)> {
+    let encapsulation_key_encoded: Encoded<<MlKem768 as KemCore>::EncapsulationKey> = opt_result!(
+        public_key.try_into(),
+        "Parse ML-KEM 768 encapsulation key failed: {}"
+    );
+    let encapsulation_key =
+        <MlKem768 as KemCore>::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<u8>, Vec<u8>)> {
+    let encapsulation_key_encoded: Encoded<<MlKem1024 as KemCore>::EncapsulationKey> = opt_result!(
+        public_key.try_into(),
+        "Parse ML-KEM 1024 encapsulation key failed: {}"
+    );
+    let encapsulation_key =
+        <MlKem1024 as KemCore>::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<u8>, Vec<u8>)> {
+    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.")
+}