From c446a52462c8ca45515220f2a21739c12dfe31a9 Mon Sep 17 00:00:00 2001 From: Hatter Jiang Date: Fri, 26 Sep 2025 23:25:53 +0800 Subject: [PATCH] feat: encrypt ML-KEM-768&1024 --- src/cmd_encrypt.rs | 29 ++++++++++++++++++++--------- src/consts.rs | 4 ++++ src/util_mlkem.rs | 22 ++++++++++++++-------- 3 files changed, 38 insertions(+), 17 deletions(-) diff --git a/src/cmd_encrypt.rs b/src/cmd_encrypt.rs index 70b3b73..e97ce24 100644 --- a/src/cmd_encrypt.rs +++ b/src/cmd_encrypt.rs @@ -12,24 +12,20 @@ use rust_util::{debugging, failure, iff, information, opt_result, simple_error, use crate::compress::GzStreamEncoder; use crate::config::{TinyEncryptConfig, TinyEncryptConfigEnvelop}; -use crate::consts::{ - ENC_AES256_GCM_KYBER1204, ENC_AES256_GCM_P256, ENC_AES256_GCM_P384, ENC_AES256_GCM_X25519, - ENC_CHACHA20_POLY1305_KYBER1204, ENC_CHACHA20_POLY1305_P256, ENC_CHACHA20_POLY1305_P384, - ENC_CHACHA20_POLY1305_X25519, SALT_COMMENT, TINY_ENC_FILE_EXT, TINY_ENC_PEM_FILE_EXT, - TINY_ENC_PEM_NAME, -}; +use crate::consts::{ENC_AES256_GCM_KYBER1204, ENC_AES256_GCM_MLKEM1024, ENC_AES256_GCM_MLKEM768, ENC_AES256_GCM_P256, ENC_AES256_GCM_P384, ENC_AES256_GCM_X25519, ENC_CHACHA20_POLY1305_KYBER1204, ENC_CHACHA20_POLY1305_MLKEM1024, ENC_CHACHA20_POLY1305_MLKEM768, ENC_CHACHA20_POLY1305_P256, ENC_CHACHA20_POLY1305_P384, ENC_CHACHA20_POLY1305_X25519, SALT_COMMENT, TINY_ENC_FILE_EXT, TINY_ENC_PEM_FILE_EXT, TINY_ENC_PEM_NAME}; use crate::crypto_cryptor::{Cryptor, KeyNonce}; use crate::spec::{ EncEncryptedMeta, EncMetadata, TinyEncryptEnvelop, TinyEncryptEnvelopType, TinyEncryptMeta, }; -use crate::util::{is_tiny_enc_file, to_pem}; +use crate::util::{decode_base64, is_tiny_enc_file, to_pem}; use crate::util_ecdh::{ecdh_kyber1024, ecdh_p256, ecdh_p384, ecdh_x25519}; use crate::util_progress::Progress; -use crate::util_rsa; +use crate::{util_mlkem, util_rsa}; use crate::wrap_key::{WrapKey, WrapKeyHeader}; use crate::{crypto_cryptor, crypto_simple, util, util_enc_file, util_env, util_gpg}; use crate::temporary_key::parse_temporary_keys; +use crate::util_mlkem::MlKemAlgo; #[derive(Debug, Args)] pub struct CmdEncrypt { @@ -333,10 +329,12 @@ pub fn encrypt_envelops(cryptor: Cryptor, key: &[u8], envelops: &[&TinyEncryptCo TinyEncryptEnvelopType::PivP384 | TinyEncryptEnvelopType::ExtP384 => { encrypted_envelops.push(encrypt_envelop_ecdh_p384(cryptor, key, envelop)?); } - // TODO ML-KEM 768 & 1024 TinyEncryptEnvelopType::StaticKyber1024 => { encrypted_envelops.push(encrypt_envelop_ecdh_kyber1204(cryptor, key, envelop)?); } + TinyEncryptEnvelopType::ExtMlKem768 | TinyEncryptEnvelopType::ExtMlKem1024 => { + encrypted_envelops.push(encrypt_envelop_ecdh_ml_kem(cryptor, key, envelop)?); + } _ => return simple_error!("Not supported type: {:?}", envelop.r#type), } } @@ -383,6 +381,19 @@ fn encrypt_envelop_ecdh_kyber1204(cryptor: Cryptor, key: &[u8], envelop: &TinyEn encrypt_envelop_shared_secret(cryptor, key, &shared_secret, &ephemeral_spki, enc_type, envelop) } +fn encrypt_envelop_ecdh_ml_kem(cryptor: Cryptor, key: &[u8], envelop: &TinyEncryptConfigEnvelop) -> XResult { + let public_key_base64 = &envelop.public_part; + let public_key = opt_result!(decode_base64(public_key_base64), "Decode ML-KEM public key from base64 failed: {}"); + let (shared_secret, ciphertext, ml_kem_algo) = util_mlkem::try_ml_kem_encapsulate(&public_key)?; + let enc_type = match (cryptor, ml_kem_algo) { + (Cryptor::Aes256Gcm, MlKemAlgo::MlKem768) => ENC_AES256_GCM_MLKEM768, + (Cryptor::Aes256Gcm, MlKemAlgo::MlKem1024) => ENC_AES256_GCM_MLKEM1024, + (Cryptor::ChaCha20Poly1305, MlKemAlgo::MlKem768) => ENC_CHACHA20_POLY1305_MLKEM768, + (Cryptor::ChaCha20Poly1305, MlKemAlgo::MlKem1024) => ENC_CHACHA20_POLY1305_MLKEM1024, + }; + encrypt_envelop_shared_secret(cryptor, key, &shared_secret, &ciphertext, enc_type, envelop) +} + fn encrypt_envelop_shared_secret(cryptor: Cryptor, key: &[u8], shared_secret: &[u8], diff --git a/src/consts.rs b/src/consts.rs index 524014d..9328180 100644 --- a/src/consts.rs +++ b/src/consts.rs @@ -3,10 +3,14 @@ pub const ENC_AES256_GCM_P256: &str = "aes256-gcm-p256"; pub const ENC_AES256_GCM_P384: &str = "aes256-gcm-p384"; pub const ENC_AES256_GCM_X25519: &str = "aes256-gcm-x25519"; pub const ENC_AES256_GCM_KYBER1204: &str = "aes256-gcm-kyber1204"; +pub const ENC_AES256_GCM_MLKEM768: &str = "aes256-gcm-mlkem768"; +pub const ENC_AES256_GCM_MLKEM1024: &str = "aes256-gcm-mlkem1024"; pub const ENC_CHACHA20_POLY1305_P256: &str = "chacha20-poly1305-p256"; pub const ENC_CHACHA20_POLY1305_P384: &str = "chacha20-poly1305-p384"; pub const ENC_CHACHA20_POLY1305_X25519: &str = "chacha20-poly1305-x25519"; pub const ENC_CHACHA20_POLY1305_KYBER1204: &str = "chacha20-poly1305-kyber1204"; +pub const ENC_CHACHA20_POLY1305_MLKEM768: &str = "chacha20-poly1305-mlkem768"; +pub const ENC_CHACHA20_POLY1305_MLKEM1024: &str = "chacha20-poly1305-mlkem1024"; // Extend and config file pub const TINY_ENC_FILE_EXT: &str = ".tinyenc"; diff --git a/src/util_mlkem.rs b/src/util_mlkem.rs index 9ebd0ed..a71ec47 100644 --- a/src/util_mlkem.rs +++ b/src/util_mlkem.rs @@ -2,6 +2,12 @@ use ml_kem::kem::Encapsulate; use ml_kem::{Encoded, EncodedSizeUser, KemCore, MlKem1024, MlKem768}; use rust_util::{opt_result, simple_error, XResult}; +#[derive(Clone, Copy, Debug)] +pub enum MlKemAlgo { + MlKem768, + MlKem1024, +} + pub fn ml_kem_768_encapsulate(public_key: &[u8]) -> XResult<(Vec, Vec)> { let encapsulation_key_encoded: Encoded<::EncapsulationKey> = opt_result!( public_key.try_into(), @@ -14,7 +20,7 @@ pub fn ml_kem_768_encapsulate(public_key: &[u8]) -> XResult<(Vec, Vec)> encapsulation_key.encapsulate(&mut rng), "Encapsulate shared key failed: {:?}" ); - Ok((ciphertext.0.to_vec(), shared_key.0.to_vec())) + Ok((shared_key.0.to_vec(), ciphertext.0.to_vec())) } pub fn ml_kem_1024_encapsulate(public_key: &[u8]) -> XResult<(Vec, Vec)> { @@ -29,15 +35,15 @@ pub fn ml_kem_1024_encapsulate(public_key: &[u8]) -> XResult<(Vec, Vec)> encapsulation_key.encapsulate(&mut rng), "Encapsulate shared key failed: {:?}" ); - Ok((ciphertext.0.to_vec(), shared_key.0.to_vec())) + Ok((shared_key.0.to_vec(), ciphertext.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)); +pub fn try_ml_kem_encapsulate(public_key: &[u8]) -> XResult<(Vec, Vec, MlKemAlgo)> { + if let Ok((shared_key, ciphertext)) = ml_kem_768_encapsulate(public_key) { + return Ok((shared_key, ciphertext, MlKemAlgo::MlKem768)); } - if let Ok((ciphertext, shared_key)) = ml_kem_1024_encapsulate(public_key) { - return Ok((ciphertext, shared_key)); + if let Ok((shared_key, ciphertext)) = ml_kem_1024_encapsulate(public_key) { + return Ok((shared_key, ciphertext, MlKemAlgo::MlKem1024)); } - simple_error!("Not ML-KEM 768 or ML-KEM 1024.") + simple_error!("Only supports ML-KEM 768 or ML-KEM 1024.") }