saorsa_core/messaging/
key_exchange.rs

1use crate::identity::FourWordAddress;
2use crate::messaging::DhtClient;
3use anyhow::Result;
4use chrono::{DateTime, Duration, Utc};
5use hkdf::Hkdf;
6use serde::{Deserialize, Serialize};
7use sha2::Sha256;
8use std::collections::HashMap;
9use std::sync::Arc;
10use tokio::sync::RwLock;
11
12use crate::quantum_crypto::ant_quic_integration::{
13    MlKemCiphertext, MlKemPublicKey, MlKemSecretKey,
14};
15
16const DHT_KEM_PREFIX: &str = "pqc:kem:";
17
18#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
19pub enum KeyExchangeType {
20    Initiation,
21    Response,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct KeyExchangeMessage {
26    pub sender: FourWordAddress,
27    pub recipient: FourWordAddress,
28    pub message_type: KeyExchangeType,
29    pub payload: Vec<u8>,
30    pub timestamp: DateTime<Utc>,
31}
32
33#[derive(Debug, Clone)]
34pub struct EstablishedKey {
35    _peer: FourWordAddress,
36    encryption_key: Vec<u8>,
37    _established_at: DateTime<Utc>,
38    expires_at: DateTime<Utc>,
39    _messages_sent: u64,
40    _messages_received: u64,
41}
42
43#[derive(Debug, Clone, Serialize, Deserialize)]
44struct KemInitiationPayload {
45    ciphertext: Vec<u8>,
46}
47
48/// PQC-only Key Exchange using ML-KEM with DHT-published public keys
49pub struct KeyExchange {
50    identity: FourWordAddress,
51    dht: DhtClient,
52    kem_public: MlKemPublicKey,
53    kem_secret: MlKemSecretKey,
54    established_keys: Arc<RwLock<HashMap<FourWordAddress, EstablishedKey>>>,
55}
56
57impl KeyExchange {
58    pub async fn new(identity: FourWordAddress, dht: DhtClient) -> Result<Self> {
59        let (kem_public, kem_secret) =
60            crate::quantum_crypto::ant_quic_integration::generate_ml_kem_keypair()?;
61        // Publish KEM pubkey
62        let dht_key = format!("{}{}", DHT_KEM_PREFIX, identity);
63        dht.put(dht_key, kem_public.as_bytes().to_vec()).await?;
64        Ok(Self {
65            identity,
66            dht,
67            kem_public,
68            kem_secret,
69            established_keys: Arc::new(RwLock::new(HashMap::new())),
70        })
71    }
72
73    /// Get the public KEM key for this node
74    pub fn kem_public_key(&self) -> &MlKemPublicKey {
75        &self.kem_public
76    }
77    async fn fetch_peer_kem_public(&self, peer: &FourWordAddress) -> Result<MlKemPublicKey> {
78        let key = format!("{}{}", DHT_KEM_PREFIX, peer);
79        let bytes = self
80            .dht
81            .get(key)
82            .await?
83            .ok_or_else(|| anyhow::anyhow!("No KEM public key for {}", peer))?;
84        MlKemPublicKey::from_bytes(&bytes)
85            .map_err(|e| anyhow::anyhow!("Invalid KEM public key for {}: {:?}", peer, e))
86    }
87
88    fn derive_session_key(shared_secret: &[u8]) -> Result<Vec<u8>> {
89        let hkdf = Hkdf::<Sha256>::new(None, shared_secret);
90        let mut okm = [0u8; 32];
91        hkdf.expand(b"saorsa-messaging-session", &mut okm)
92            .map_err(|e| anyhow::anyhow!("HKDF expand failed: {:?}", e))?;
93        Ok(okm.to_vec())
94    }
95
96    pub async fn initiate_exchange(&self, peer: FourWordAddress) -> Result<KeyExchangeMessage> {
97        let peer_pub = self.fetch_peer_kem_public(&peer).await?;
98        let (ciphertext, shared) =
99            crate::quantum_crypto::ant_quic_integration::ml_kem_encapsulate(&peer_pub)?;
100        let session_key = Self::derive_session_key(shared.as_bytes())?;
101        self.store_session(peer.clone(), session_key).await;
102        let payload = KemInitiationPayload {
103            ciphertext: ciphertext.as_bytes().to_vec(),
104        };
105        let payload_bytes = bincode::serialize(&payload)?;
106        Ok(KeyExchangeMessage {
107            sender: self.identity.clone(),
108            recipient: peer,
109            message_type: KeyExchangeType::Initiation,
110            payload: payload_bytes,
111            timestamp: Utc::now(),
112        })
113    }
114
115    pub async fn respond_to_exchange(&self, msg: KeyExchangeMessage) -> Result<KeyExchangeMessage> {
116        if msg.message_type != KeyExchangeType::Initiation {
117            return Err(anyhow::anyhow!("Unexpected message type"));
118        }
119        let payload: KemInitiationPayload = bincode::deserialize(&msg.payload)?;
120        let kem_ct = MlKemCiphertext::from_bytes(&payload.ciphertext)
121            .map_err(|e| anyhow::anyhow!("Invalid KEM ciphertext: {:?}", e))?;
122        let shared = crate::quantum_crypto::ant_quic_integration::ml_kem_decapsulate(
123            &self.kem_secret,
124            &kem_ct,
125        )?;
126        let session_key = Self::derive_session_key(shared.as_bytes())?;
127        self.store_session(msg.sender.clone(), session_key).await;
128        Ok(KeyExchangeMessage {
129            sender: self.identity.clone(),
130            recipient: msg.sender,
131            message_type: KeyExchangeType::Response,
132            payload: Vec::new(),
133            timestamp: Utc::now(),
134        })
135    }
136
137    pub async fn complete_exchange(&self, _message: KeyExchangeMessage) -> Result<()> {
138        Ok(())
139    }
140
141    async fn store_session(&self, peer: FourWordAddress, key: Vec<u8>) {
142        let mut map = self.established_keys.write().await;
143        map.insert(
144            peer.clone(),
145            EstablishedKey {
146                _peer: peer,
147                encryption_key: key,
148                _established_at: Utc::now(),
149                expires_at: Utc::now() + Duration::hours(24),
150                _messages_sent: 0,
151                _messages_received: 0,
152            },
153        );
154    }
155
156    pub async fn get_session_key(&self, peer: &FourWordAddress) -> Result<Vec<u8>> {
157        let keys = self.established_keys.read().await;
158        if let Some(established) = keys.get(peer)
159            && established.expires_at > Utc::now()
160        {
161            return Ok(established.encryption_key.clone());
162        }
163        Err(anyhow::anyhow!("No established PQC session with {}", peer))
164    }
165}
saorsa_core/messaging/key_exchange.rs

saorsa_core/messaging/
key_exchange.rs
