feat: external_public_key supports ML-KEM

src/cmd_external_public_key.rs

@@ -11,6 +11,7 @@ use std::collections::BTreeMap;
 use rsa::RsaPrivateKey;
 use spki::EncodePublicKey;
 use x509_parser::parse_x509_certificate;
+use crate::mlkemutil::try_parse_decapsulate_key_private_get_encapsulate;
 use crate::pivutil::ToStr;
 
 pub struct CommandImpl;
@@ -94,6 +95,11 @@ fn fetch_public_key(parameter: &str, serial_opt: &Option<&str>) -> XResult<Vec<u
                 let private_key_der = base64_decode(&private_key)?;
                 let rsa_private_key = RsaPrivateKey::from_pkcs8_der(&private_key_der)?;
                 Ok(rsa_private_key.to_public_key().to_public_key_der()?.to_vec())
+            } else if key.algorithm.is_mlkem() {
+                let private_key = cmd_hmac_decrypt::try_decrypt(&mut None, &key.hmac_enc_private_key)?;
+                let private_key_der = base64_decode(&private_key)?;
+                let (_, ek_public) = try_parse_decapsulate_key_private_get_encapsulate(&private_key_der)?;
+                Ok(ek_public)
             } else {
                 simple_error!("Invalid algorithm: {}", key.algorithm.to_str())
             }

src/keyutil.rs

@@ -37,6 +37,10 @@ impl KeyUri {
             KeyAlgorithmId::EccP256 => AlgorithmType::Es256,
             KeyAlgorithmId::EccP384 => AlgorithmType::Es384,
             KeyAlgorithmId::EccP521 => AlgorithmType::Es512,
+            // ML-KEM not supports JWS
+            KeyAlgorithmId::MlKem512
+            | KeyAlgorithmId::MlKem768
+            | KeyAlgorithmId::MlKem1024 => AlgorithmType::None,
         }
     }
 }
@@ -97,6 +101,9 @@ pub enum KeyAlgorithmId {
     EccP256,
     EccP384,
     EccP521,
+    MlKem512,
+    MlKem768,
+    MlKem1024,
 }
 
 impl KeyAlgorithmId {
@@ -118,6 +125,9 @@ impl KeyAlgorithmId {
             KeyAlgorithmId::EccP256 => Some(AlgorithmId::EccP256),
             KeyAlgorithmId::EccP384 => Some(AlgorithmId::EccP384),
             KeyAlgorithmId::EccP521 => None,
+            KeyAlgorithmId::MlKem512 => None,
+            KeyAlgorithmId::MlKem768 => None,
+            KeyAlgorithmId::MlKem1024 => None,
         }
     }
 
@@ -127,17 +137,41 @@ impl KeyAlgorithmId {
             | KeyAlgorithmId::Rsa2048
             | KeyAlgorithmId::Rsa3072
             | KeyAlgorithmId::Rsa4096 => true,
-            KeyAlgorithmId::EccP256 | KeyAlgorithmId::EccP384 | KeyAlgorithmId::EccP521 => false,
+            KeyAlgorithmId::EccP256
+            | KeyAlgorithmId::EccP384
+            | KeyAlgorithmId::EccP521
+            | KeyAlgorithmId::MlKem512
+            | KeyAlgorithmId::MlKem768
+            | KeyAlgorithmId::MlKem1024 => false,
         }
     }
 
     pub fn is_ecc(&self) -> bool {
         match self {
+            KeyAlgorithmId::EccP256 | KeyAlgorithmId::EccP384 | KeyAlgorithmId::EccP521 => true,
             KeyAlgorithmId::Rsa1024
             | KeyAlgorithmId::Rsa2048
             | KeyAlgorithmId::Rsa3072
+            | KeyAlgorithmId::Rsa4096
+            | KeyAlgorithmId::MlKem512
+            | KeyAlgorithmId::MlKem768
+            | KeyAlgorithmId::MlKem1024 => false,
+        }
+    }
+
+
+    pub fn is_mlkem(&self) -> bool {
+        match self {
+            | KeyAlgorithmId::MlKem512
+            | KeyAlgorithmId::MlKem768
+            | KeyAlgorithmId::MlKem1024 => true,
+            KeyAlgorithmId::EccP256
+            | KeyAlgorithmId::EccP384
+            | KeyAlgorithmId::EccP521
+            | KeyAlgorithmId::Rsa1024
+            | KeyAlgorithmId::Rsa2048
+            | KeyAlgorithmId::Rsa3072
             | KeyAlgorithmId::Rsa4096 => false,
-            KeyAlgorithmId::EccP256 | KeyAlgorithmId::EccP384 | KeyAlgorithmId::EccP521 => true,
         }
     }
 
@@ -150,6 +184,9 @@ impl KeyAlgorithmId {
             KeyAlgorithmId::EccP256 => "ES256,",
             KeyAlgorithmId::EccP384 => "ES384",
             KeyAlgorithmId::EccP521 => "ES512",
+            KeyAlgorithmId::MlKem512
+            | KeyAlgorithmId::MlKem768
+            | KeyAlgorithmId::MlKem1024 => "__UNKNOWN__",
         }
     }
 }
@@ -167,6 +204,9 @@ impl FromStr for KeyAlgorithmId {
             "p256" => Some(KeyAlgorithmId::EccP256),
             "p384" => Some(KeyAlgorithmId::EccP384),
             "p521" => Some(KeyAlgorithmId::EccP521),
+            "mlkem512" => Some(KeyAlgorithmId::MlKem512),
+            "mlkem768" => Some(KeyAlgorithmId::MlKem768),
+            "mlkem1024" => Some(KeyAlgorithmId::MlKem1024),
             _ => None,
         }
     }
@@ -182,6 +222,9 @@ impl ToStr for KeyAlgorithmId {
             KeyAlgorithmId::EccP256 => "p256",
             KeyAlgorithmId::EccP384 => "p384",
             KeyAlgorithmId::EccP521 => "p521",
+            KeyAlgorithmId::MlKem512 => "mlkem512",
+            KeyAlgorithmId::MlKem768 => "mlkem768",
+            KeyAlgorithmId::MlKem1024 => "mlkem1024",
         }
     }
 }

src/mlkemutil.rs

@@ -6,6 +6,13 @@ use ml_kem::{EncodedSizeUser, KemCore, MlKem1024, MlKem512, MlKem768};
 use rust_util::XResult;
 use std::convert::TryInto;
 
+#[derive(Debug, Clone, Copy)]
+pub enum MlKemLen {
+    Len512,
+    Len768,
+    Len1024,
+}
+
 pub fn generate_mlkem_keypair(len: usize) -> XResult<(String, String, String, Vec<u8>, String)> {
     let (dk_private, ek_public) = match len {
         512 => generate_ml_kem_512(),
@@ -27,6 +34,43 @@ pub fn generate_mlkem_keypair(len: usize) -> XResult<(String, String, String, Ve
     ))
 }
 
+pub fn try_parse_decapsulate_key_private_then_decapsulate(
+    key_bytes: &[u8],
+    ciphertext_bytes: &[u8],
+) -> XResult<(MlKemLen, Vec<u8>)> {
+    if let Ok(shared_secret) =
+        parse_decapsulate_key_512_private_then_decapsulate(key_bytes, ciphertext_bytes)
+    {
+        return Ok((MlKemLen::Len512, shared_secret.to_vec()));
+    }
+    if let Ok(shared_secret) =
+        parse_decapsulate_key_768_private_then_decapsulate(key_bytes, ciphertext_bytes)
+    {
+        return Ok((MlKemLen::Len768, shared_secret.to_vec()));
+    }
+    if let Ok(shared_secret) =
+        parse_decapsulate_key_1024_private_then_decapsulate(key_bytes, ciphertext_bytes)
+    {
+        return Ok((MlKemLen::Len1024, shared_secret.to_vec()));
+    }
+    simple_error!("Invalid decapsulation key, only allow MK-KEM-512, ML-KEM-768, ML-KEM-1024")
+}
+
+pub fn try_parse_decapsulate_key_private_get_encapsulate(
+    key_bytes: &[u8],
+) -> XResult<(MlKemLen, Vec<u8>)> {
+    if let Ok(encapsulate_key) = parse_decapsulate_key_512_private_get_encapsulate(key_bytes) {
+        return Ok((MlKemLen::Len512, encapsulate_key));
+    }
+    if let Ok(encapsulate_key) = parse_decapsulate_key_768_private_get_encapsulate(key_bytes) {
+        return Ok((MlKemLen::Len768, encapsulate_key));
+    }
+    if let Ok(encapsulate_key) = parse_decapsulate_key_1024_private_get_encapsulate(key_bytes) {
+        return Ok((MlKemLen::Len1024, encapsulate_key));
+    }
+    simple_error!("Invalid decapsulation key, only allow MK-KEM-512, ML-KEM-768, ML-KEM-1024")
+}
+
 pub fn generate_ml_kem_512() -> (Vec<u8>, Vec<u8>) {
     let mut rng = rand::thread_rng();
     let (dk_private, ek_public) = <MlKem512 as KemCore>::generate(&mut rng);
@@ -136,7 +180,7 @@ pub fn parse_decapsulate_key_1024_private_then_decapsulate(
     key_bytes: &[u8],
     ciphertext_bytes: &[u8],
 ) -> XResult<Vec<u8>> {
-    let dk = <MlKem768 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
+    let dk = <MlKem1024 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
         key_bytes.try_into(),
         "Parse decapsulation key 1024 failed: {}"
     ));
@@ -149,3 +193,27 @@ pub fn parse_decapsulate_key_1024_private_then_decapsulate(
     );
     Ok(shared_key.0.to_vec())
 }
+
+pub fn parse_decapsulate_key_512_private_get_encapsulate(key_bytes: &[u8]) -> XResult<Vec<u8>> {
+    let dk = <MlKem512 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
+        key_bytes.try_into(),
+        "Parse decapsulation key 512 failed: {}"
+    ));
+    Ok(dk.encapsulation_key().as_bytes().0.to_vec())
+}
+
+pub fn parse_decapsulate_key_768_private_get_encapsulate(key_bytes: &[u8]) -> XResult<Vec<u8>> {
+    let dk = <MlKem768 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
+        key_bytes.try_into(),
+        "Parse decapsulation key 768 failed: {}"
+    ));
+    Ok(dk.encapsulation_key().as_bytes().0.to_vec())
+}
+
+pub fn parse_decapsulate_key_1024_private_get_encapsulate(key_bytes: &[u8]) -> XResult<Vec<u8>> {
+    let dk = <MlKem1024 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
+        key_bytes.try_into(),
+        "Parse decapsulation key 1024 failed: {}"
+    ));
+    Ok(dk.encapsulation_key().as_bytes().0.to_vec())
+}