1use 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
23pub const TAG_SIZE: usize = 8;
25
26pub const X25519_PUBLIC_KEY_SIZE: usize = 32;
28
29pub const X25519_PRIVATE_KEY_SIZE: usize = 32;
31
32pub const CHACHA20_KEY_SIZE: usize = 32;
34
35pub const POLY1305_TAG_SIZE: usize = 16;
37
38pub const NONCE_SIZE: usize = 12;
40
41pub const DEFAULT_WINDOW_MS: u64 = 10_000;
43
44const 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#[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#[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 pub fn generate() -> Self {
72 let mut private_key_bytes = [0u8; 32];
73 OsRng.fill_bytes(&mut private_key_bytes);
74
75 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 pub fn from_private_key(mut key_bytes: [u8; 32]) -> Self {
91 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 pub fn public_key_bytes(&self) -> [u8; X25519_PUBLIC_KEY_SIZE] {
105 self.public_key_bytes
106 }
107
108 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 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
122pub 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 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 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
155pub 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
168pub 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
181pub 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 if tag[0] >= 1 && tag[0] <= 4 {
204 tag[0] = tag[0].wrapping_add(5); }
206 tag
207}
208
209pub fn compute_time_window(timestamp_ms: u64, window_ms: u64) -> u64 {
211 timestamp_ms / window_ms
212}
213
214pub 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
223pub fn random_bytes(len: usize) -> Vec<u8> {
225 let mut buf = vec![0u8; len];
226 OsRng.fill_bytes(&mut buf);
227 buf
228}
229
230pub fn blake3_hash(data: &[u8]) -> [u8; 32] {
232 blake3::hash(data).into()
233}
234
235pub 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
244pub 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); assert_eq!(tag1, tag3); }
295}