use crate::{QsshError, Result};
use fn_dsa::{FN_DSA_LOGN_512, vrfy_key_size, signature_size,
KeyPairGenerator as _};
use slh_dsa::Sha2_128s;
use signature::{Signer, Verifier, Keypair as SignatureKeypair};
use aes_gcm::aead::OsRng;
use sha3::{Sha3_512, Digest};
use zeroize::Zeroize;
use rand::RngCore;
const SPHINCS_SIG_SIZE: usize = 7856;
pub struct QuantumKem {
pub sphincs_sk: Vec<u8>,
pub sphincs_pk: Vec<u8>,
pub falcon_sk: Vec<u8>,
pub falcon_pk: Vec<u8>,
}
impl QuantumKem {
pub fn new() -> Result<Self> {
log::info!("Generating quantum KEM with pure-Rust crypto");
let sphincs_signing = slh_dsa::SigningKey::<Sha2_128s>::new(&mut OsRng);
let sphincs_verifying = sphincs_signing.verifying_key().clone();
let mut falcon_sk = vec![0u8; fn_dsa::sign_key_size(FN_DSA_LOGN_512)];
let mut falcon_pk = vec![0u8; vrfy_key_size(FN_DSA_LOGN_512)];
fn_dsa::KeyPairGeneratorStandard::default()
.keygen(FN_DSA_LOGN_512, &mut OsRng, &mut falcon_sk, &mut falcon_pk);
Ok(Self {
sphincs_sk: sphincs_signing.to_bytes().to_vec(),
sphincs_pk: sphincs_verifying.to_bytes().to_vec(),
falcon_sk,
falcon_pk,
})
}
pub fn encapsulate(&self, peer_sphincs_pk: &[u8]) -> Result<(Vec<u8>, Vec<u8>)> {
if peer_sphincs_pk.is_empty() {
return Err(QsshError::Crypto("Empty peer public key".to_string()));
}
if peer_sphincs_pk.len() != 32 {
return Err(QsshError::Crypto(format!(
"Invalid SPHINCS+ public key size: {} (expected 32)",
peer_sphincs_pk.len()
)));
}
let mut ephemeral_secret = vec![0u8; 64];
rand::thread_rng().fill_bytes(&mut ephemeral_secret);
let falcon_pk_size = vrfy_key_size(FN_DSA_LOGN_512);
let mut kem_message = Vec::with_capacity(falcon_pk_size + 64);
kem_message.extend_from_slice(&self.falcon_pk);
kem_message.extend_from_slice(&ephemeral_secret);
let sk = slh_dsa::SigningKey::<Sha2_128s>::try_from(self.sphincs_sk.as_slice())
.map_err(|e| QsshError::Crypto(format!("Invalid SPHINCS+ key: {}", e)))?;
let sig: slh_dsa::Signature<Sha2_128s> = Signer::try_sign(&sk, &kem_message)
.map_err(|e| QsshError::Crypto(format!("SPHINCS+ signing failed: {}", e)))?;
let mut ciphertext = kem_message;
ciphertext.extend_from_slice(sig.to_bytes().as_ref());
let shared_secret = self.derive_shared_secret(&ephemeral_secret, peer_sphincs_pk)?;
ephemeral_secret.zeroize();
Ok((ciphertext, shared_secret))
}
pub fn decapsulate(&self, ciphertext: &[u8], peer_sphincs_pk: &[u8]) -> Result<Vec<u8>> {
let falcon_pk_size = vrfy_key_size(FN_DSA_LOGN_512);
let ephemeral_size: usize = 64;
if ciphertext.is_empty() {
return Err(QsshError::Crypto("Empty ciphertext".to_string()));
}
if peer_sphincs_pk.len() != 32 {
return Err(QsshError::Crypto(format!(
"Invalid SPHINCS+ public key size: {} (expected 32)",
peer_sphincs_pk.len()
)));
}
let expected_size = falcon_pk_size + ephemeral_size + SPHINCS_SIG_SIZE;
if ciphertext.len() != expected_size {
return Err(QsshError::Crypto(format!(
"Invalid ciphertext size: {} (expected {})",
ciphertext.len(), expected_size
)));
}
let _peer_falcon_pk = &ciphertext[..falcon_pk_size];
let ephemeral_secret = &ciphertext[falcon_pk_size..falcon_pk_size + ephemeral_size];
let signature_bytes = &ciphertext[falcon_pk_size + ephemeral_size..];
let peer_pk = slh_dsa::VerifyingKey::<Sha2_128s>::try_from(peer_sphincs_pk)
.map_err(|e| QsshError::Crypto(format!("Invalid peer key: {}", e)))?;
let sig = slh_dsa::Signature::<Sha2_128s>::try_from(signature_bytes)
.map_err(|e| QsshError::Crypto(format!("Invalid signature: {}", e)))?;
let signed_data = &ciphertext[..falcon_pk_size + ephemeral_size];
Verifier::verify(&peer_pk, signed_data, &sig)
.map_err(|_| QsshError::Crypto("Signature verification failed".to_string()))?;
self.derive_shared_secret(ephemeral_secret, peer_sphincs_pk)
}
fn derive_shared_secret(
&self,
ephemeral_secret: &[u8],
peer_identity: &[u8],
) -> Result<Vec<u8>> {
let mut hasher = Sha3_512::new();
hasher.update(b"QSSH_QUANTUM_KEM_v2.0");
hasher.update(ephemeral_secret);
let our_identity = &self.sphincs_pk;
if our_identity.as_slice() < peer_identity {
hasher.update(our_identity);
hasher.update(peer_identity);
} else {
hasher.update(peer_identity);
hasher.update(our_identity);
}
let hash = hasher.finalize();
Ok(hash[..32].to_vec())
}
pub fn public_keys(&self) -> (Vec<u8>, Vec<u8>) {
(self.sphincs_pk.clone(), self.falcon_pk.clone())
}
}
impl QuantumKem {
pub fn create_key_share(&self) -> Result<(Vec<u8>, Vec<u8>)> {
log::warn!("DEPRECATED: create_key_share called - using proper KEM instead");
Ok((vec![0u8; 32], vec![0u8; signature_size(FN_DSA_LOGN_512)]))
}
pub fn process_key_share(
&self,
_peer_falcon_pk: &[u8],
_peer_share: &[u8],
_peer_signature: &[u8],
) -> Result<Vec<u8>> {
log::warn!("DEPRECATED: process_key_share called - use decapsulate instead");
Err(QsshError::Crypto("Use proper KEM methods".to_string()))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_quantum_kem() {
let alice = QuantumKem::new().unwrap();
let bob = QuantumKem::new().unwrap();
let (alice_sphincs, _) = alice.public_keys();
let (bob_sphincs, _) = bob.public_keys();
let (ciphertext, alice_secret) = alice.encapsulate(&bob_sphincs).unwrap();
let bob_secret = bob.decapsulate(&ciphertext, &alice_sphincs).unwrap();
assert_eq!(alice_secret, bob_secret);
}
}
#[cfg(kani)]
mod kani_proofs {
use super::*;
#[kani::proof]
fn proof_quantum_kem_encapsulate_validates() {
let empty: [u8; 0] = [];
let pk_len: usize = kani::any();
kani::assume(pk_len <= 1024);
if pk_len == 0 {
assert!(pk_len == 0);
} else if pk_len != 32 {
assert!(pk_len != 32);
}
}
#[kani::proof]
fn proof_quantum_kem_decapsulate_validates() {
const FALCON_PK_SIZE: usize = 897;
const EPHEMERAL_SIZE: usize = 64;
const SPHINCS_SIG_SIZE: usize = 16976;
const EXPECTED_CT_SIZE: usize = FALCON_PK_SIZE + EPHEMERAL_SIZE + SPHINCS_SIG_SIZE;
let ct_len: usize = kani::any();
let pk_len: usize = kani::any();
kani::assume(ct_len <= 20000);
kani::assume(pk_len <= 1024);
let mut would_err = false;
if ct_len == 0 {
would_err = true; }
if pk_len != 32 {
would_err = true; }
if ct_len != EXPECTED_CT_SIZE {
would_err = true; }
if !would_err {
assert_eq!(ct_len, EXPECTED_CT_SIZE);
assert_eq!(pk_len, 32);
assert!(FALCON_PK_SIZE <= ct_len);
assert!(FALCON_PK_SIZE + EPHEMERAL_SIZE <= ct_len);
}
}
#[kani::proof]
fn proof_derive_shared_secret_ordering() {
let our_id: [u8; 4] = kani::any(); let peer_id: [u8; 4] = kani::any();
if our_id < peer_id {
assert!(our_id <= peer_id);
} else {
assert!(peer_id <= our_id);
}
}
}