From 3d29fe6a6d8e4de525d9be8136da5f63f05fb083 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sat, 27 Sep 2025 00:49:37 +0800
Subject: [PATCH] feat: add generate_ml_kem_768, pending implement ML-KEM later

---
 Cargo.lock       | 52 ++++++++++++++++++++++++++++++++++++++++++++++++
 Cargo.toml       |  4 ++++
 src/main.rs      |  1 +
 src/mlkemutil.rs | 40 +++++++++++++++++++++++++++++++++++++
 4 files changed, 97 insertions(+)
 create mode 100644 src/mlkemutil.rs

diff --git a/Cargo.lock b/Cargo.lock
index 86bfca9..1a217a7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -614,6 +614,7 @@ dependencies = [
  "external-command-rs",
  "hex",
  "jwt",
+ "ml-kem",
  "openpgp-card 0.3.7",
  "openpgp-card-pcsc",
  "openpgp-card-sequoia",
@@ -1750,6 +1751,15 @@ version = "1.10.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
 
+[[package]]
+name = "hybrid-array"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2d35805454dc9f8662a98d6d61886ffe26bd465f5960e0e55345c70d5c0d2a9"
+dependencies = [
+ "typenum",
+]
+
 [[package]]
 name = "hyper"
 version = "1.6.0"
@@ -2071,6 +2081,25 @@ dependencies = [
  "sha2 0.10.9",
 ]
 
+[[package]]
+name = "keccak"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ecc2af9a1119c51f12a14607e783cb977bde58bc069ff0c3da1095e635d70654"
+dependencies = [
+ "cpufeatures",
+]
+
+[[package]]
+name = "kem"
+version = "0.3.0-pre.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2b8645470337db67b01a7f966decf7d0bafedbae74147d33e641c67a91df239f"
+dependencies = [
+ "rand_core 0.6.4",
+ "zeroize",
+]
+
 [[package]]
 name = "kernel32-sys"
 version = "0.2.2"
@@ -2356,6 +2385,19 @@ dependencies = [
  "ws2_32-sys",
 ]
 
+[[package]]
+name = "ml-kem"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "97befee0c869cb56f3118f49d0f9bb68c9e3f380dec23c1100aedc4ec3ba239a"
+dependencies = [
+ "hybrid-array",
+ "kem",
+ "rand_core 0.6.4",
+ "sha3",
+ "zeroize",
+]
+
 [[package]]
 name = "native-tls"
 version = "0.2.14"
@@ -3892,6 +3934,16 @@ dependencies = [
  "digest 0.10.7",
 ]
 
+[[package]]
+name = "sha3"
+version = "0.10.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "75872d278a8f37ef87fa0ddbda7802605cb18344497949862c0d4dcb291eba60"
+dependencies = [
+ "digest 0.10.7",
+ "keccak",
+]
+
 [[package]]
 name = "shlex"
 version = "1.3.0"
diff --git a/Cargo.toml b/Cargo.toml
index 501ae9a..45700e4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -67,5 +67,9 @@ ssh-key = { version = "0.6", features = ["ecdsa", "alloc"] }
 tokio = "1.45.1"
 ssh-encoding = { version = "0.2.0", features = ["alloc"] }
 zeroize = "1.8"
+ml-kem = { version = "0.2.1", features = ["zeroize"] }
 #lazy_static = "1.4.0"
 #ctap-hid-fido2 = "2.1.3"
+
+#[patch.crates-io]
+#ml-kem = { path = "externals/ml-kem" }
diff --git a/src/main.rs b/src/main.rs
index 24a0e7b..4aaeab8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -84,6 +84,7 @@ mod sshutil;
 mod util;
 mod yubikeyutil;
 mod cmd_yubikey;
+mod mlkemutil;
 
 pub struct DefaultCommandImpl;
 
diff --git a/src/mlkemutil.rs b/src/mlkemutil.rs
new file mode 100644
index 0000000..84110ec
--- /dev/null
+++ b/src/mlkemutil.rs
@@ -0,0 +1,40 @@
+use crate::util::base64_encode;
+use ml_kem::kem::{Decapsulate, Encapsulate};
+use ml_kem::{EncodedSizeUser, KemCore, MlKem768};
+use rust_util::XResult;
+use std::convert::TryInto;
+
+// #[test]
+pub fn generate_ml_kem_768() -> XResult<()> {
+    let mut rng = rand::thread_rng();
+    let (dk, ek) = <MlKem768 as KemCore>::generate(&mut rng);
+    println!("dk: {}", base64_encode(&dk.as_bytes().0.to_vec()));
+    println!("ek: {}", base64_encode(ek.as_bytes().0.to_vec()));
+
+    let ek_bytes = dk.as_bytes().0.to_vec();
+    let dk = <MlKem768 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
+        ek_bytes.as_slice().try_into(),
+        "Parse decapsulation key failed: {}"
+    ));
+
+    let (encoded_ciphertext, shared_key) = opt_result!(
+        ek.encapsulate(&mut rng),
+        "Encapsulation key encapsulate failed: {:?}"
+    );
+    println!(
+        "encoded_ciphertext: {}",
+        base64_encode(&encoded_ciphertext.0.to_vec())
+    );
+    println!("shared_key: {}", base64_encode(&shared_key.0.to_vec()));
+
+    let k_bytes = encoded_ciphertext.0.to_vec();
+    let shared_key_2 = opt_result!(
+        dk.decapsulate(opt_result!(
+            &k_bytes.as_slice().try_into(),
+            "Parse encoded ciphertext failed: {}"
+        )),
+        "Decapsulation key decapsulate failed: {:?}"
+    );
+    println!("shared_key2: {}", base64_encode(&shared_key_2.0.to_vec()));
+    Ok(())
+}