Skip to main content

shadow_crypto/
keyexchange.rs

1use x25519_dalek::{EphemeralSecret, PublicKey, SharedSecret};
2use anyhow::Result;
3use serde::{Deserialize, Serialize};
4use crate::encryption::EncryptionKey;
5
6/// X25519 key exchange for forward secrecy
7pub struct KeyExchange {
8    secret: EphemeralSecret,
9    public: PublicKey,
10}
11
12impl KeyExchange {
13    /// Generate new ephemeral keypair
14    pub fn new() -> Self {
15        let secret = EphemeralSecret::random_from_rng(rand::rngs::OsRng);
16        let public = PublicKey::from(&secret);
17        
18        Self { secret, public }
19    }
20
21    /// Get public key for exchange
22    pub fn public_key(&self) -> &PublicKey {
23        &self.public
24    }
25
26    /// Get public key as bytes
27    pub fn public_key_bytes(&self) -> [u8; 32] {
28        self.public.to_bytes()
29    }
30
31    /// Perform key exchange and derive encryption key
32    pub fn exchange(self, peer_public: &PublicKey) -> Result<EncryptionKey> {
33        let shared_secret = self.secret.diffie_hellman(peer_public);
34        
35        // Derive encryption key from shared secret using HKDF
36        let key_bytes = derive_key_from_shared_secret(&shared_secret)?;
37        
38        Ok(EncryptionKey::from_bytes(key_bytes))
39    }
40
41    /// Exchange with peer public key bytes
42    pub fn exchange_bytes(self, peer_public_bytes: &[u8; 32]) -> Result<EncryptionKey> {
43        let peer_public = PublicKey::from(*peer_public_bytes);
44        self.exchange(&peer_public)
45    }
46}
47
48impl Default for KeyExchange {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54/// Derive encryption key from X25519 shared secret
55fn derive_key_from_shared_secret(shared_secret: &SharedSecret) -> Result<[u8; 32]> {
56    use hkdf::Hkdf;
57    use sha2::Sha256;
58
59    let hk = Hkdf::<Sha256>::new(None, shared_secret.as_bytes());
60    let mut okm = [0u8; 32];
61    hk.expand(b"shadow-network-dh-key", &mut okm)
62        .map_err(|e| anyhow::anyhow!("HKDF expansion failed: {:?}", e))?;
63    
64    Ok(okm)
65}
66
67/// Serializable public key for network transmission
68#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct PublicKeyBytes(pub [u8; 32]);
70
71impl From<&PublicKey> for PublicKeyBytes {
72    fn from(pk: &PublicKey) -> Self {
73        Self(pk.to_bytes())
74    }
75}
76
77impl From<PublicKeyBytes> for PublicKey {
78    fn from(pkb: PublicKeyBytes) -> Self {
79        PublicKey::from(pkb.0)
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn test_key_exchange() {
89        // Alice and Bob generate ephemeral keypairs
90        let alice = KeyExchange::new();
91        let bob = KeyExchange::new();
92
93        // Exchange public keys
94        let alice_public = *alice.public_key();
95        let bob_public = *bob.public_key();
96
97        // Both derive the same shared secret
98        let alice_key = alice.exchange(&bob_public).unwrap();
99        let bob_key = bob.exchange(&alice_public).unwrap();
100
101        assert_eq!(alice_key.as_bytes(), bob_key.as_bytes());
102    }
103
104    #[test]
105    fn test_key_exchange_bytes() {
106        let alice = KeyExchange::new();
107        let bob = KeyExchange::new();
108
109        let alice_pub_bytes = alice.public_key_bytes();
110        let bob_pub_bytes = bob.public_key_bytes();
111
112        let alice_key = alice.exchange_bytes(&bob_pub_bytes).unwrap();
113        let bob_key = bob.exchange_bytes(&alice_pub_bytes).unwrap();
114
115        assert_eq!(alice_key.as_bytes(), bob_key.as_bytes());
116    }
117}