242 lines
8.4 KiB
Rust
242 lines
8.4 KiB
Rust
use crate::util::{base64_encode, to_pem};
|
|
use ml_kem::kem::{Decapsulate, Encapsulate};
|
|
use ml_kem::{EncodedSizeUser, KemCore, MlKem1024, MlKem512, MlKem768};
|
|
use rust_util::XResult;
|
|
use std::convert::TryInto;
|
|
use crate::pivutil::ToStr;
|
|
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub enum MlKemLen {
|
|
Len512,
|
|
Len768,
|
|
Len1024,
|
|
}
|
|
|
|
impl ToStr for MlKemLen {
|
|
fn to_str(&self) -> &str {
|
|
match self {
|
|
MlKemLen::Len512 => "mlkem512",
|
|
MlKemLen::Len768 => "mlkem768",
|
|
MlKemLen::Len1024 => "mlkem1024",
|
|
}
|
|
}
|
|
}
|
|
|
|
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(),
|
|
768 => generate_ml_kem_768(),
|
|
1024 => generate_ml_kem_1024(),
|
|
_ => return simple_error!("Invalid ML-KEM={}", len),
|
|
};
|
|
let secret_key_der_base64 = base64_encode(&dk_private);
|
|
let secret_key_pem = to_pem(&format!("ML-KEM-{} PRIVATE KEY", len), &dk_private);
|
|
let public_key_pem = to_pem(&format!("ML-KEM-{} PUBLIC KEY", len), &ek_public);
|
|
let public_key_der = ek_public;
|
|
let jwk_ec_key = "".to_string();
|
|
Ok((
|
|
secret_key_der_base64,
|
|
secret_key_pem,
|
|
public_key_pem,
|
|
public_key_der,
|
|
jwk_ec_key,
|
|
))
|
|
}
|
|
|
|
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);
|
|
(
|
|
dk_private.as_bytes().0.to_vec(),
|
|
ek_public.as_bytes().0.to_vec(),
|
|
)
|
|
}
|
|
|
|
pub fn generate_ml_kem_768() -> (Vec<u8>, Vec<u8>) {
|
|
let mut rng = rand::thread_rng();
|
|
let (dk_private, ek_public) = <MlKem768 as KemCore>::generate(&mut rng);
|
|
(
|
|
dk_private.as_bytes().0.to_vec(),
|
|
ek_public.as_bytes().0.to_vec(),
|
|
)
|
|
}
|
|
|
|
pub fn generate_ml_kem_1024() -> (Vec<u8>, Vec<u8>) {
|
|
let mut rng = rand::thread_rng();
|
|
let (dk_private, ek_public) = <MlKem1024 as KemCore>::generate(&mut rng);
|
|
(
|
|
dk_private.as_bytes().0.to_vec(),
|
|
ek_public.as_bytes().0.to_vec(),
|
|
)
|
|
}
|
|
|
|
pub fn try_parse_encapsulation_key_public_then_encapsulate(bytes: &[u8]) -> XResult<(MlKemLen, Vec<u8>, Vec<u8>)> {
|
|
if let Ok((ciphertext, shared_key)) = parse_encapsulation_key_512_public_then_encapsulate(bytes) {
|
|
return Ok((MlKemLen::Len512, ciphertext, shared_key));
|
|
}
|
|
if let Ok((ciphertext, shared_key)) = parse_encapsulation_key_768_public_then_encapsulate(bytes) {
|
|
return Ok((MlKemLen::Len768, ciphertext, shared_key));
|
|
}
|
|
if let Ok((ciphertext, shared_key)) = parse_encapsulation_key_1024_public_then_encapsulate(bytes) {
|
|
return Ok((MlKemLen::Len1024, ciphertext, shared_key));
|
|
}
|
|
simple_error!("Invalid encapsulation key, only allow MK-KEM-512, ML-KEM-768, ML-KEM-1024")
|
|
}
|
|
|
|
pub fn parse_encapsulation_key_512_public_then_encapsulate(
|
|
bytes: &[u8],
|
|
) -> XResult<(Vec<u8>, Vec<u8>)> {
|
|
let ek = <MlKem512 as KemCore>::EncapsulationKey::from_bytes(&opt_result!(
|
|
bytes.try_into(),
|
|
"Parse encapsulation key 512 failed: {}"
|
|
));
|
|
let (ciphertext, shared_key) = opt_result!(
|
|
ek.encapsulate(&mut rand::thread_rng()),
|
|
"Encapsulation key 512 encapsulate failed: {:?}"
|
|
);
|
|
Ok((ciphertext.0.to_vec(), shared_key.0.to_vec()))
|
|
}
|
|
|
|
pub fn parse_encapsulation_key_768_public_then_encapsulate(
|
|
bytes: &[u8],
|
|
) -> XResult<(Vec<u8>, Vec<u8>)> {
|
|
let ek = <MlKem768 as KemCore>::EncapsulationKey::from_bytes(&opt_result!(
|
|
bytes.try_into(),
|
|
"Parse encapsulation key 768 failed: {}"
|
|
));
|
|
let (ciphertext, shared_key) = opt_result!(
|
|
ek.encapsulate(&mut rand::thread_rng()),
|
|
"Encapsulation key 768 encapsulate failed: {:?}"
|
|
);
|
|
Ok((ciphertext.0.to_vec(), shared_key.0.to_vec()))
|
|
}
|
|
|
|
pub fn parse_encapsulation_key_1024_public_then_encapsulate(
|
|
bytes: &[u8],
|
|
) -> XResult<(Vec<u8>, Vec<u8>)> {
|
|
let ek = <MlKem1024 as KemCore>::EncapsulationKey::from_bytes(&opt_result!(
|
|
bytes.try_into(),
|
|
"Parse encapsulation key 1024 failed: {}"
|
|
));
|
|
let (ciphertext, shared_key) = opt_result!(
|
|
ek.encapsulate(&mut rand::thread_rng()),
|
|
"Encapsulation key 1024 encapsulate failed: {:?}"
|
|
);
|
|
Ok((ciphertext.0.to_vec(), shared_key.0.to_vec()))
|
|
}
|
|
|
|
pub fn parse_decapsulate_key_512_private_then_decapsulate(
|
|
key_bytes: &[u8],
|
|
ciphertext_bytes: &[u8],
|
|
) -> XResult<Vec<u8>> {
|
|
let dk = <MlKem512 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
|
|
key_bytes.try_into(),
|
|
"Parse decapsulation key 512 failed: {}"
|
|
));
|
|
let shared_key = opt_result!(
|
|
dk.decapsulate(opt_result!(
|
|
ciphertext_bytes.try_into(),
|
|
"Parse encoded ciphertext 512 failed: {}"
|
|
)),
|
|
"Decapsulation key 512 decapsulate failed: {:?}"
|
|
);
|
|
Ok(shared_key.0.to_vec())
|
|
}
|
|
|
|
pub fn parse_decapsulate_key_768_private_then_decapsulate(
|
|
key_bytes: &[u8],
|
|
ciphertext_bytes: &[u8],
|
|
) -> XResult<Vec<u8>> {
|
|
let dk = <MlKem768 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
|
|
key_bytes.try_into(),
|
|
"Parse decapsulation key 768 failed: {}"
|
|
));
|
|
let shared_key = opt_result!(
|
|
dk.decapsulate(opt_result!(
|
|
ciphertext_bytes.try_into(),
|
|
"Parse encoded ciphertext 768 failed: {}"
|
|
)),
|
|
"Decapsulation key 768 decapsulate failed: {:?}"
|
|
);
|
|
Ok(shared_key.0.to_vec())
|
|
}
|
|
|
|
pub fn parse_decapsulate_key_1024_private_then_decapsulate(
|
|
key_bytes: &[u8],
|
|
ciphertext_bytes: &[u8],
|
|
) -> XResult<Vec<u8>> {
|
|
let dk = <MlKem1024 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
|
|
key_bytes.try_into(),
|
|
"Parse decapsulation key 1024 failed: {}"
|
|
));
|
|
let shared_key = opt_result!(
|
|
dk.decapsulate(opt_result!(
|
|
ciphertext_bytes.try_into(),
|
|
"Parse encoded ciphertext 1024 failed: {}"
|
|
)),
|
|
"Decapsulation key 1024 decapsulate failed: {:?}"
|
|
);
|
|
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())
|
|
}
|