feat: external_public_key supports ML-KEM

This commit is contained in:
2025-09-27 12:02:26 +08:00
parent 9f544e3cb7
commit 10c38cda8a
3 changed files with 120 additions and 3 deletions

View File

@@ -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())
}