shadow_crypto/
keyexchange.rs1use x25519_dalek::{EphemeralSecret, PublicKey, SharedSecret};
2use anyhow::Result;
3use serde::{Deserialize, Serialize};
4use crate::encryption::EncryptionKey;
5
6pub struct KeyExchange {
8 secret: EphemeralSecret,
9 public: PublicKey,
10}
11
12impl KeyExchange {
13 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 pub fn public_key(&self) -> &PublicKey {
23 &self.public
24 }
25
26 pub fn public_key_bytes(&self) -> [u8; 32] {
28 self.public.to_bytes()
29 }
30
31 pub fn exchange(self, peer_public: &PublicKey) -> Result<EncryptionKey> {
33 let shared_secret = self.secret.diffie_hellman(peer_public);
34
35 let key_bytes = derive_key_from_shared_secret(&shared_secret)?;
37
38 Ok(EncryptionKey::from_bytes(key_bytes))
39 }
40
41 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
54fn 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#[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 let alice = KeyExchange::new();
91 let bob = KeyExchange::new();
92
93 let alice_public = *alice.public_key();
95 let bob_public = *bob.public_key();
96
97 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}