use hkdf::Hkdf;
use rand::RngCore;
use sha2::{Digest, Sha256};
use x25519_dalek::{PublicKey, StaticSecret};
use crate::error::{ProtocolError, Result};
pub const TX_ID_LEN: usize = 16;
const SAS_INFO_V1: &[u8] = b"huddle-sas-v1";
const SAS_INFO_PQBIND: &[u8] = b"huddle-sas-pqbind-v1";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SasCode {
pub emoji_indices: [u8; 7],
pub decimal: String,
}
impl SasCode {
pub fn emoji_string(&self) -> String {
self.emoji_indices
.iter()
.map(|i| SAS_EMOJI[*i as usize].0)
.collect::<Vec<_>>()
.join(" ")
}
pub fn emoji_labels(&self) -> String {
self.emoji_indices
.iter()
.map(|i| SAS_EMOJI[*i as usize].1)
.collect::<Vec<_>>()
.join(" / ")
}
}
pub fn new_session() -> ([u8; TX_ID_LEN], StaticSecret, PublicKey) {
let mut tx_id = [0u8; TX_ID_LEN];
rand::thread_rng().fill_bytes(&mut tx_id);
let secret = StaticSecret::random_from_rng(rand::thread_rng());
let public = PublicKey::from(&secret);
(tx_id, secret, public)
}
pub fn derive_sas_code(
our_secret: &StaticSecret,
their_public: &PublicKey,
tx_id: &[u8; TX_ID_LEN],
our_mlkem_ek: Option<&[u8]>,
their_mlkem_ek: Option<&[u8]>,
) -> Result<SasCode> {
let shared = our_secret.diffie_hellman(their_public);
if !shared.was_contributory() {
return Err(ProtocolError::Session(
"SAS rejected: peer X25519 ephemeral is non-contributory (small-order point)".into(),
));
}
let hk = Hkdf::<Sha256>::new(Some(tx_id), shared.as_bytes());
let info = sas_info(our_mlkem_ek, their_mlkem_ek);
let mut okm = [0u8; 11];
hk.expand(&info, &mut okm)
.expect("11 bytes is well within HKDF output limit");
let b = &okm[..6];
let mut raw_emoji = [0u8; 7];
raw_emoji[0] = b[0] >> 2;
raw_emoji[1] = ((b[0] & 0x03) << 4) | (b[1] >> 4);
raw_emoji[2] = ((b[1] & 0x0f) << 2) | (b[2] >> 6);
raw_emoji[3] = b[2] & 0x3f;
raw_emoji[4] = b[3] >> 2;
raw_emoji[5] = ((b[3] & 0x03) << 4) | (b[4] >> 4);
raw_emoji[6] = ((b[4] & 0x0f) << 2) | (b[5] >> 6);
let emoji_indices = derive_emoji_indices_rejection(&hk, raw_emoji);
let d = &okm[6..11];
let chunk0 = ((u32::from(d[0]) << 5) | (u32::from(d[1]) >> 3)) & 0x1fff;
let chunk1 =
((u32::from(d[1] & 0x07) << 10) | (u32::from(d[2]) << 2) | (u32::from(d[3]) >> 6)) & 0x1fff;
let chunk2 = ((u32::from(d[3] & 0x3f) << 7) | (u32::from(d[4]) >> 1)) & 0x1fff;
let decimal = format!("{}-{}-{}", chunk0 + 1000, chunk1 + 1000, chunk2 + 1000);
Ok(SasCode {
emoji_indices,
decimal,
})
}
fn sas_info(our_mlkem_ek: Option<&[u8]>, their_mlkem_ek: Option<&[u8]>) -> Vec<u8> {
let their_ek = match their_mlkem_ek {
None => return SAS_INFO_V1.to_vec(),
Some(ek) => ek,
};
let mut eks: Vec<&[u8]> = Vec::with_capacity(2);
if let Some(ours) = our_mlkem_ek {
eks.push(ours);
}
eks.push(their_ek);
eks.sort_unstable();
let mut hasher = Sha256::new();
for ek in &eks {
hasher.update(ek);
}
let digest = hasher.finalize();
let mut info = Vec::with_capacity(SAS_INFO_PQBIND.len() + digest.len());
info.extend_from_slice(SAS_INFO_PQBIND);
info.extend_from_slice(&digest);
info
}
pub const SAS_EMOJI: [(&str, &str); 49] = [
("🐶", "dog"),
("🐱", "cat"),
("🦁", "lion"),
("🐎", "horse"),
("🦄", "unicorn"),
("🐷", "pig"),
("🐘", "elephant"),
("🐰", "rabbit"),
("🐼", "panda"),
("🐓", "rooster"),
("🐧", "penguin"),
("🐢", "turtle"),
("🐟", "fish"),
("🐙", "octopus"),
("🦋", "butterfly"),
("🌷", "flower"),
("🌳", "tree"),
("🌵", "cactus"),
("🍄", "mushroom"),
("🌏", "globe"),
("🌙", "moon"),
("☁️", "cloud"),
("🔥", "fire"),
("🍌", "banana"),
("🍎", "apple"),
("🍓", "strawberry"),
("🌽", "corn"),
("🍕", "pizza"),
("🎂", "cake"),
("❤️", "heart"),
("🙂", "smiley"),
("🤖", "robot"),
("🎩", "hat"),
("👓", "glasses"),
("🔧", "spanner"),
("🎅", "santa"),
("👍", "thumbs up"),
("☂️", "umbrella"),
("⌛", "hourglass"),
("⏰", "clock"),
("🎁", "gift"),
("💡", "light bulb"),
("📕", "book"),
("✏️", "pencil"),
("📎", "paperclip"),
("✂️", "scissors"),
("🔒", "lock"),
("🔑", "key"),
("🔨", "hammer"),
];
fn derive_emoji_indices_rejection(hk: &Hkdf<Sha256>, initial: [u8; 7]) -> [u8; 7] {
let mut out = [0u8; 7];
let mut accepted = 0usize;
for &v in &initial {
if v < 49 {
out[accepted] = v;
accepted += 1;
if accepted == 7 {
return out;
}
}
}
let mut counter: u32 = 0;
while accepted < 7 {
let info = {
let mut buf = [0u8; 24];
buf[..16].copy_from_slice(b"huddle-sas-v1-rs");
buf[16..20].copy_from_slice(&counter.to_be_bytes());
buf
};
let mut block = [0u8; 32];
if hk.expand(&info, &mut block).is_err() {
for v in &mut initial.iter().copied() {
if accepted < 7 {
out[accepted] = v % 49;
accepted += 1;
}
}
break;
}
for &byte in block.iter() {
let candidate = byte & 0x3f;
if candidate < 49 {
out[accepted] = candidate;
accepted += 1;
if accepted == 7 {
return out;
}
}
}
counter += 1;
}
out
}
pub fn parse_pubkey(b64: &str) -> Result<PublicKey> {
use base64::engine::general_purpose::STANDARD as B64;
use base64::Engine;
let bytes = B64
.decode(b64)
.map_err(|e| ProtocolError::Session(format!("bad x25519 pubkey b64: {e}")))?;
if bytes.len() != 32 {
return Err(ProtocolError::Session(format!(
"x25519 pubkey is {} bytes, expected 32",
bytes.len()
)));
}
let mut arr = [0u8; 32];
arr.copy_from_slice(&bytes);
Ok(PublicKey::from(arr))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn both_sides_derive_same_code() {
let (tx_id, alice_secret, alice_pub) = new_session();
let (_, bob_secret, bob_pub) = new_session();
let alice_code = derive_sas_code(&alice_secret, &bob_pub, &tx_id, None, None).unwrap();
let bob_code = derive_sas_code(&bob_secret, &alice_pub, &tx_id, None, None).unwrap();
assert_eq!(alice_code, bob_code);
let parts: Vec<&str> = alice_code.decimal.split('-').collect();
assert_eq!(parts.len(), 3);
for p in parts {
assert_eq!(p.len(), 4);
let n: u32 = p.parse().unwrap();
assert!((1000..=9191).contains(&n));
}
for i in alice_code.emoji_indices {
assert!((i as usize) < SAS_EMOJI.len());
}
}
#[test]
fn different_tx_id_yields_different_code() {
let (tx_id_a, alice_secret, _) = new_session();
let (_, bob_secret, bob_pub) = new_session();
let alice_code = derive_sas_code(&alice_secret, &bob_pub, &tx_id_a, None, None).unwrap();
let mut tx_id_b = tx_id_a;
tx_id_b[0] ^= 0xff;
let alice_code_b = derive_sas_code(&alice_secret, &bob_pub, &tx_id_b, None, None).unwrap();
let _ = bob_secret;
assert_ne!(alice_code, alice_code_b);
}
#[test]
fn mitm_substitute_yields_different_code() {
let (tx_id, alice_secret, alice_pub) = new_session();
let (_, bob_secret, bob_pub) = new_session();
let (_, _mallory_secret, mallory_pub) = new_session();
let alice_thinks_bob =
derive_sas_code(&alice_secret, &mallory_pub, &tx_id, None, None).unwrap();
let bob_thinks_alice =
derive_sas_code(&bob_secret, &mallory_pub, &tx_id, None, None).unwrap();
assert_ne!(alice_thinks_bob, bob_thinks_alice);
let alice_real = derive_sas_code(&alice_secret, &bob_pub, &tx_id, None, None).unwrap();
let bob_real = derive_sas_code(&bob_secret, &alice_pub, &tx_id, None, None).unwrap();
assert_eq!(alice_real, bob_real);
}
#[test]
fn rejects_small_order_ephemeral() {
let (tx_id, our_secret, _) = new_session();
let zero_pub = PublicKey::from([0u8; 32]);
assert!(derive_sas_code(&our_secret, &zero_pub, &tx_id, None, None).is_err());
}
fn fake_ek(fill: u8) -> Vec<u8> {
vec![fill; crate::crypto::pqc::MLKEM_EK_LEN]
}
#[test]
fn pq_binding_changes_the_code() {
let (tx_id, alice_secret, _) = new_session();
let (_, _bob_secret, bob_pub) = new_session();
let our = fake_ek(0xA4);
let ek = fake_ek(0xA5);
let classical = derive_sas_code(&alice_secret, &bob_pub, &tx_id, None, None).unwrap();
let bound =
derive_sas_code(&alice_secret, &bob_pub, &tx_id, Some(&our), Some(&ek)).unwrap();
assert_ne!(
classical, bound,
"binding the ML-KEM ek must change the derived SAS code"
);
}
#[test]
fn both_sides_distinct_eks_agree() {
let (tx_id, alice_secret, alice_pub) = new_session();
let (_, bob_secret, bob_pub) = new_session();
let ek_a = fake_ek(0x11);
let ek_b = fake_ek(0x22);
let alice =
derive_sas_code(&alice_secret, &bob_pub, &tx_id, Some(&ek_a), Some(&ek_b)).unwrap();
let bob =
derive_sas_code(&bob_secret, &alice_pub, &tx_id, Some(&ek_b), Some(&ek_a)).unwrap();
assert_eq!(
alice, bob,
"sorted dual-ek binding must agree across peers in opposite roles"
);
}
#[test]
fn one_side_bound_other_not_diverges() {
let (tx_id, alice_secret, alice_pub) = new_session();
let (_, bob_secret, bob_pub) = new_session();
let ek_a = fake_ek(0x77);
let ek_b = fake_ek(0x88);
let alice_bound =
derive_sas_code(&alice_secret, &bob_pub, &tx_id, Some(&ek_a), Some(&ek_b)).unwrap();
let bob_stripped =
derive_sas_code(&bob_secret, &alice_pub, &tx_id, Some(&ek_b), None).unwrap();
assert_ne!(
alice_bound, bob_stripped,
"a stripped (classical) side must not match a bound side"
);
}
#[test]
fn different_ek_yields_different_code() {
let (tx_id, secret, _) = new_session();
let (_, _b, peer_pub) = new_session();
let our = fake_ek(0x00);
let a =
derive_sas_code(&secret, &peer_pub, &tx_id, Some(&our), Some(&fake_ek(0x01))).unwrap();
let b =
derive_sas_code(&secret, &peer_pub, &tx_id, Some(&our), Some(&fake_ek(0x02))).unwrap();
assert_ne!(a, b, "different bound eks must yield different codes");
}
#[test]
fn pqbind_is_order_independent() {
let ek_a = fake_ek(0x33);
let ek_b = fake_ek(0x44);
assert_eq!(
sas_info(Some(&ek_a), Some(&ek_b)),
sas_info(Some(&ek_b), Some(&ek_a)),
"dual-ek binding must be order-independent"
);
}
#[test]
fn classical_none_path_is_unchanged_golden() {
assert_eq!(sas_info(None, None), b"huddle-sas-v1".to_vec());
assert_eq!(
sas_info(Some(&fake_ek(0x01)), None),
b"huddle-sas-v1".to_vec(),
"PQ-capable self + classical partner must stay on the 1.x transcript"
);
let ek_a = fake_ek(0x5C);
let ek_b = fake_ek(0x6D);
let info = sas_info(Some(&ek_a), Some(&ek_b));
assert_eq!(info.len(), SAS_INFO_PQBIND.len() + 32);
assert!(info.starts_with(SAS_INFO_PQBIND));
let mut sorted = [ek_a.as_slice(), ek_b.as_slice()];
sorted.sort_unstable();
let mut h = Sha256::new();
for e in sorted {
h.update(e);
}
assert_eq!(&info[SAS_INFO_PQBIND.len()..], &h.finalize()[..]);
}
#[test]
fn pubkey_round_trip() {
let (_, _, pub_) = new_session();
use base64::engine::general_purpose::STANDARD as B64;
use base64::Engine;
let encoded = B64.encode(pub_.as_bytes());
let decoded = parse_pubkey(&encoded).unwrap();
assert_eq!(decoded.as_bytes(), pub_.as_bytes());
}
}