Skip to main content

aivpn_common/
crypto.rs

1//! Cryptographic primitives for AIVPN
2//!
3//! Implements:
4//! - X25519 key exchange
5//! - ChaCha20-Poly1305 AEAD encryption
6//! - BLAKE3 hashing and HMAC
7//! - Resonance Tag generation
8
9use blake3::Hasher;
10use chacha20poly1305::{
11    aead::{Aead, KeyInit, OsRng},
12    ChaCha20Poly1305, Key as ChachaKey, Nonce,
13};
14use hmac::Hmac;
15use rand::RngCore;
16use sha2::Sha256;
17use subtle::ConstantTimeEq;
18use x25519_dalek;
19use zeroize::{Zeroize, ZeroizeOnDrop};
20
21use crate::error::{Error, Result};
22
23/// Size of resonance tag in bytes
24pub const TAG_SIZE: usize = 8;
25
26/// Size of X25519 public key in bytes
27pub const X25519_PUBLIC_KEY_SIZE: usize = 32;
28
29/// Size of X25519 private key in bytes
30pub const X25519_PRIVATE_KEY_SIZE: usize = 32;
31
32/// Size of ChaCha20-Poly1305 key in bytes
33pub const CHACHA20_KEY_SIZE: usize = 32;
34
35/// Size of Poly1305 tag in bytes
36pub const POLY1305_TAG_SIZE: usize = 16;
37
38/// Size of nonce in bytes
39pub const NONCE_SIZE: usize = 12;
40
41/// Default time window for tag rotation in milliseconds (optimized: increased from 5s to 10s)
42pub const DEFAULT_WINDOW_MS: u64 = 10_000;
43
44/// HKDF context strings
45const HKDF_SESSION_KEY_CONTEXT: &str = "aivpn-session-key-v1";
46const HKDF_TAG_SECRET_CONTEXT: &str = "aivpn-tag-secret-v1";
47const HKDF_PRNG_SEED_CONTEXT: &str = "aivpn-prng-seed-v1";
48
49/// Session keys derived from key exchange
50#[derive(Clone, Zeroize, ZeroizeOnDrop)]
51pub struct SessionKeys {
52    pub session_key: [u8; CHACHA20_KEY_SIZE],
53    pub tag_secret: [u8; 32],
54    pub prng_seed: [u8; 32],
55}
56
57/// X25519 keypair for key exchange
58#[derive(Debug, Clone)]
59pub struct KeyPair {
60    private_key_bytes: [u8; X25519_PRIVATE_KEY_SIZE],
61    public_key_bytes: [u8; X25519_PUBLIC_KEY_SIZE],
62}
63
64impl Drop for KeyPair {
65    fn drop(&mut self) {
66        self.private_key_bytes.zeroize();
67    }
68}
69impl KeyPair {
70    /// Generate a new ephemeral keypair
71    pub fn generate() -> Self {
72        let mut private_key_bytes = [0u8; 32];
73        OsRng.fill_bytes(&mut private_key_bytes);
74
75        // X25519 clamping (RFC 7748)
76        private_key_bytes[0] &= 248;
77        private_key_bytes[31] &= 127;
78        private_key_bytes[31] |= 64;
79
80        let public_key_bytes =
81            x25519_dalek::x25519(private_key_bytes, x25519_dalek::X25519_BASEPOINT_BYTES);
82
83        Self {
84            private_key_bytes,
85            public_key_bytes,
86        }
87    }
88
89    /// Create keypair from existing private key bytes (loaded from file)
90    pub fn from_private_key(mut key_bytes: [u8; 32]) -> Self {
91        // X25519 clamping (RFC 7748)
92        key_bytes[0] &= 248;
93        key_bytes[31] &= 127;
94        key_bytes[31] |= 64;
95        let public_key_bytes =
96            x25519_dalek::x25519(key_bytes, x25519_dalek::X25519_BASEPOINT_BYTES);
97        Self {
98            private_key_bytes: key_bytes,
99            public_key_bytes,
100        }
101    }
102
103    /// Get the public key as bytes
104    pub fn public_key_bytes(&self) -> [u8; X25519_PUBLIC_KEY_SIZE] {
105        self.public_key_bytes
106    }
107
108    /// Compute shared secret with remote public key
109    /// Returns error if the result is all-zero (small subgroup attack)
110    pub fn compute_shared(&self, remote_public: &[u8; X25519_PUBLIC_KEY_SIZE]) -> Result<[u8; 32]> {
111        let shared = x25519_dalek::x25519(self.private_key_bytes, *remote_public);
112        // Reject all-zero shared secret (small subgroup / identity point attack)
113        if shared.ct_eq(&[0u8; 32]).into() {
114            return Err(Error::Crypto(
115                "DH result is all-zero (possible small subgroup attack)".into(),
116            ));
117        }
118        Ok(shared)
119    }
120}
121
122/// Derive session keys from DH result using HKDF-BLAKE3
123pub fn derive_session_keys(
124    dh_result: &[u8; 32],
125    preshared_key: Option<&[u8; 32]>,
126    eph_pub: &[u8; X25519_PUBLIC_KEY_SIZE],
127) -> SessionKeys {
128    // IKM = dh_result || preshared_key (or just dh_result if no PSK)
129    let ikm: Vec<u8> = if let Some(psk) = preshared_key {
130        let mut buf = [0u8; 64];
131        buf[..32].copy_from_slice(dh_result);
132        buf[32..].copy_from_slice(psk);
133        buf.to_vec()
134    } else {
135        dh_result.to_vec()
136    };
137
138    // Derive keys using BLAKE3 derive_key with different contexts
139    // Context strings are combined with key material for domain separation
140    let session_key_input: Vec<u8> = [ikm.clone(), eph_pub.to_vec()].concat();
141    let tag_secret_input: Vec<u8> = [ikm.clone(), eph_pub.to_vec()].concat();
142    let prng_seed_input: Vec<u8> = [ikm, eph_pub.to_vec()].concat();
143
144    let session_key_hash = blake3::derive_key(HKDF_SESSION_KEY_CONTEXT, &session_key_input);
145    let tag_secret_hash = blake3::derive_key(HKDF_TAG_SECRET_CONTEXT, &tag_secret_input);
146    let prng_seed_hash = blake3::derive_key(HKDF_PRNG_SEED_CONTEXT, &prng_seed_input);
147
148    SessionKeys {
149        session_key: session_key_hash[..CHACHA20_KEY_SIZE].try_into().unwrap(),
150        tag_secret: tag_secret_hash[..32].try_into().unwrap(),
151        prng_seed: prng_seed_hash[..32].try_into().unwrap(),
152    }
153}
154
155/// Encrypt payload using ChaCha20-Poly1305
156pub fn encrypt_payload(
157    key: &[u8; CHACHA20_KEY_SIZE],
158    nonce: &[u8; NONCE_SIZE],
159    plaintext: &[u8],
160) -> Result<Vec<u8>> {
161    let cipher = ChaCha20Poly1305::new(ChachaKey::from_slice(key));
162    let nonce = Nonce::from_slice(nonce);
163
164    let ciphertext = cipher.encrypt(nonce, plaintext)?;
165    Ok(ciphertext)
166}
167
168/// Decrypt payload using ChaCha20-Poly1305
169pub fn decrypt_payload(
170    key: &[u8; CHACHA20_KEY_SIZE],
171    nonce: &[u8; NONCE_SIZE],
172    ciphertext: &[u8],
173) -> Result<Vec<u8>> {
174    let cipher = ChaCha20Poly1305::new(ChachaKey::from_slice(key));
175    let nonce = Nonce::from_slice(nonce);
176
177    let plaintext = cipher.decrypt(nonce, ciphertext)?;
178    Ok(plaintext)
179}
180
181/// Generate Resonance Tag using HMAC-BLAKE3
182///
183/// Tag = HMAC-BLAKE3(tag_secret, counter_bytes || time_window_bytes)
184/// truncated to first 8 bytes.
185/// The first byte is guaranteed NOT to be 1–4 (WireGuard message types),
186/// preventing heuristic WireGuard detection by Wireshark / DPI (Issue #30).
187pub fn generate_resonance_tag(
188    tag_secret: &[u8; 32],
189    counter: u64,
190    time_window: u64,
191) -> [u8; TAG_SIZE] {
192    let mut hasher = Hasher::new_keyed(tag_secret);
193    hasher.update(&counter.to_le_bytes());
194    hasher.update(&time_window.to_le_bytes());
195
196    let hash = hasher.finalize();
197    let mut tag = [0u8; TAG_SIZE];
198    tag.copy_from_slice(&hash.as_bytes()[..TAG_SIZE]);
199    // Avoid WireGuard message type signatures: 0x01 (Initiation), 0x02 (Response),
200    // 0x03 (Cookie), 0x04 (Transport).  DPI/Wireshark checks byte[0] ∈ {1..4}
201    // followed by three zero bytes.  Shifting byte[0] out of that range eliminates
202    // the heuristic match without reducing tag entropy (the secret is still 256-bit).
203    if tag[0] >= 1 && tag[0] <= 4 {
204        tag[0] = tag[0].wrapping_add(5); // 1→6, 2→7, 3→8, 4→9
205    }
206    tag
207}
208
209/// Compute time window from timestamp
210pub fn compute_time_window(timestamp_ms: u64, window_ms: u64) -> u64 {
211    timestamp_ms / window_ms
212}
213
214/// Get current timestamp in milliseconds
215pub fn current_timestamp_ms() -> u64 {
216    use std::time::{SystemTime, UNIX_EPOCH};
217    SystemTime::now()
218        .duration_since(UNIX_EPOCH)
219        .unwrap()
220        .as_millis() as u64
221}
222
223/// Generate random bytes
224pub fn random_bytes(len: usize) -> Vec<u8> {
225    let mut buf = vec![0u8; len];
226    OsRng.fill_bytes(&mut buf);
227    buf
228}
229
230/// Compute BLAKE3 hash
231pub fn blake3_hash(data: &[u8]) -> [u8; 32] {
232    blake3::hash(data).into()
233}
234
235/// Obfuscate/deobfuscate ephemeral public key using server's static public key.
236/// XOR with BLAKE3-derived mask makes eph_pub indistinguishable from random. (HIGH-9)
237pub fn obfuscate_eph_pub(eph_pub: &mut [u8; 32], server_static_pub: &[u8; 32]) {
238    let mask = blake3::derive_key("aivpn-eph-obfuscation-v1", server_static_pub);
239    for i in 0..32 {
240        eph_pub[i] ^= mask[i];
241    }
242}
243
244/// Compute HMAC-SHA256
245pub fn hmac_sha256(key: &[u8], data: &[u8]) -> [u8; 32] {
246    use hmac::Mac;
247    type HmacSha256 = Hmac<Sha256>;
248    let mut mac = <HmacSha256 as Mac>::new_from_slice(key).expect("HMAC can take key of any size");
249    mac.update(data);
250    let result = mac.finalize();
251    result.into_bytes().into()
252}
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257
258    #[test]
259    fn test_key_exchange() {
260        let client_keys = KeyPair::generate();
261        let server_keys = KeyPair::generate();
262
263        let client_shared = client_keys
264            .compute_shared(&server_keys.public_key_bytes())
265            .unwrap();
266        let server_shared = server_keys
267            .compute_shared(&client_keys.public_key_bytes())
268            .unwrap();
269
270        assert_eq!(client_shared, server_shared);
271    }
272
273    #[test]
274    fn test_encrypt_decrypt() {
275        let key = [1u8; CHACHA20_KEY_SIZE];
276        let nonce = [2u8; NONCE_SIZE];
277        let plaintext = b"Hello, AIVPN!";
278
279        let ciphertext = encrypt_payload(&key, &nonce, plaintext).unwrap();
280        let decrypted = decrypt_payload(&key, &nonce, &ciphertext).unwrap();
281
282        assert_eq!(plaintext.to_vec(), decrypted);
283    }
284
285    #[test]
286    fn test_resonance_tag() {
287        let tag_secret = [3u8; 32];
288        let tag1 = generate_resonance_tag(&tag_secret, 1, 100);
289        let tag2 = generate_resonance_tag(&tag_secret, 2, 100);
290        let tag3 = generate_resonance_tag(&tag_secret, 1, 100);
291
292        assert_ne!(tag1, tag2); // Different counter
293        assert_eq!(tag1, tag3); // Same counter and window
294    }
295}