Skip to main content

peat_mesh/security/
formation_key.rs

1//! # Formation Key - Shared Secret Authentication for Mesh Formations
2//!
3//! Provides pre-shared key (PSK) authentication for formation membership.
4//!
5//! ## Overview
6//!
7//! A formation key is a shared secret that all nodes in a formation possess.
8//! When two nodes connect, they perform a challenge-response handshake to
9//! verify mutual possession of the key before allowing sync operations.
10//!
11//! ## Security Model
12//!
13//! - **HMAC-SHA256** challenge-response (not replay-able)
14//! - **Key derivation** from shared secret using HKDF
15//! - **Formation isolation** - different formations can't sync
16//!
17//! ## Wire Protocol
18//!
19//! After QUIC connection establishment:
20//!
21//! ```text
22//! Initiator                              Responder
23//!     |                                      |
24//!     |  ---- QUIC Connect (TLS) ---->       |
25//!     |                                      |
26//!     |  <---- Challenge (32-byte nonce) --- |
27//!     |                                      |
28//!     |  ---- Response (32-byte HMAC) ---->  |
29//!     |                                      |
30//!     |  <---- OK / REJECT ---------------   |
31//!     |                                      |
32//! ```
33
34use hmac::{Hmac, Mac};
35use rand_core::{OsRng, RngCore};
36use sha2::Sha256;
37
38use super::error::SecurityError;
39
40/// Size of challenge nonce in bytes
41pub const FORMATION_CHALLENGE_SIZE: usize = 32;
42
43/// Size of challenge response (HMAC-SHA256 output)
44pub const FORMATION_RESPONSE_SIZE: usize = 32;
45
46/// HKDF info string for formation key derivation
47const HKDF_INFO_FORMATION: &[u8] = b"peat-protocol-v1-formation";
48
49type HmacSha256 = Hmac<Sha256>;
50
51/// Formation key for shared secret authentication
52///
53/// All nodes in a formation share this key. Used for challenge-response
54/// authentication during connection establishment.
55#[derive(Clone)]
56pub struct FormationKey {
57    /// Formation identifier (e.g., "alpha-company")
58    formation_id: String,
59    /// Derived HMAC key (32 bytes)
60    hmac_key: [u8; 32],
61}
62
63impl FormationKey {
64    /// Create a new formation key from a shared secret
65    ///
66    /// The key is derived using HKDF-SHA256 with the formation ID as context,
67    /// ensuring different formations have different derived keys even with
68    /// the same base secret.
69    pub fn new(formation_id: &str, shared_secret: &[u8; 32]) -> Self {
70        use hkdf::Hkdf;
71
72        // Derive HMAC key using HKDF
73        // Salt = formation_id, IKM = shared_secret, info = HKDF_INFO_FORMATION
74        let hk = Hkdf::<Sha256>::new(Some(formation_id.as_bytes()), shared_secret);
75        let mut hmac_key = [0u8; 32];
76        hk.expand(HKDF_INFO_FORMATION, &mut hmac_key)
77            .expect("HKDF expand should never fail with 32-byte output");
78
79        Self {
80            formation_id: formation_id.to_string(),
81            hmac_key,
82        }
83    }
84
85    /// Create a formation key from a base64-encoded shared secret
86    ///
87    /// # Key Derivation
88    ///
89    /// - If the decoded secret is exactly 32 bytes, it's used directly
90    /// - Otherwise, SHA-256 is used to derive a 32-byte key from the input
91    pub fn from_base64(formation_id: &str, base64_secret: &str) -> Result<Self, SecurityError> {
92        use base64::{engine::general_purpose::STANDARD, Engine};
93        use sha2::{Digest, Sha256};
94
95        let secret_bytes = STANDARD.decode(base64_secret.trim()).map_err(|e| {
96            SecurityError::AuthenticationFailed(format!("Invalid base64 shared secret: {}", e))
97        })?;
98
99        let secret: [u8; 32] = if secret_bytes.len() == 32 {
100            // Exact 32 bytes - use directly (backwards compatible)
101            let mut arr = [0u8; 32];
102            arr.copy_from_slice(&secret_bytes);
103            arr
104        } else {
105            // Derive 32-byte key using SHA-256 (supports EC keys, etc.)
106            let mut hasher = Sha256::new();
107            hasher.update(&secret_bytes);
108            hasher.finalize().into()
109        };
110
111        Ok(Self::new(formation_id, &secret))
112    }
113
114    /// Generate a random 32-byte shared secret (for initial setup)
115    ///
116    /// Returns base64-encoded 32-byte random secret.
117    pub fn generate_secret() -> String {
118        use base64::{engine::general_purpose::STANDARD, Engine};
119
120        let mut secret = [0u8; 32];
121        OsRng.fill_bytes(&mut secret);
122        STANDARD.encode(secret)
123    }
124
125    /// Get the formation ID
126    pub fn formation_id(&self) -> &str {
127        &self.formation_id
128    }
129
130    /// Create a challenge for the peer to respond to
131    ///
132    /// Returns tuple of (nonce, expected_response).
133    pub fn create_challenge(
134        &self,
135    ) -> (
136        [u8; FORMATION_CHALLENGE_SIZE],
137        [u8; FORMATION_RESPONSE_SIZE],
138    ) {
139        let mut nonce = [0u8; FORMATION_CHALLENGE_SIZE];
140        OsRng.fill_bytes(&mut nonce);
141
142        let expected = self.compute_response(&nonce);
143        (nonce, expected)
144    }
145
146    /// Respond to a challenge from a peer
147    pub fn respond_to_challenge(&self, nonce: &[u8]) -> [u8; FORMATION_RESPONSE_SIZE] {
148        self.compute_response(nonce)
149    }
150
151    /// Verify a peer's response to our challenge
152    pub fn verify_response(&self, nonce: &[u8], response: &[u8; FORMATION_RESPONSE_SIZE]) -> bool {
153        let expected = self.compute_response(nonce);
154
155        // Constant-time comparison to prevent timing attacks
156        use subtle::ConstantTimeEq;
157        expected.ct_eq(response).into()
158    }
159
160    /// Compute HMAC response for a nonce
161    ///
162    /// Response = HMAC-SHA256(key, nonce || formation_id)
163    fn compute_response(&self, nonce: &[u8]) -> [u8; FORMATION_RESPONSE_SIZE] {
164        let mut mac =
165            HmacSha256::new_from_slice(&self.hmac_key).expect("HMAC key should be valid length");
166
167        mac.update(nonce);
168        mac.update(self.formation_id.as_bytes());
169
170        let result = mac.finalize();
171        let mut response = [0u8; FORMATION_RESPONSE_SIZE];
172        response.copy_from_slice(&result.into_bytes());
173        response
174    }
175}
176
177impl std::fmt::Debug for FormationKey {
178    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
179        f.debug_struct("FormationKey")
180            .field("formation_id", &self.formation_id)
181            .field("hmac_key", &"[REDACTED]")
182            .finish()
183    }
184}
185
186/// Challenge message sent from responder to initiator
187#[derive(Debug, Clone)]
188pub struct FormationChallenge {
189    /// Formation ID (so initiator knows which key to use)
190    pub formation_id: String,
191    /// Random nonce
192    pub nonce: [u8; FORMATION_CHALLENGE_SIZE],
193}
194
195impl FormationChallenge {
196    /// Serialize to bytes for wire transmission
197    ///
198    /// Format: formation_id_len (2 bytes) || formation_id || nonce (32 bytes)
199    pub fn to_bytes(&self) -> Vec<u8> {
200        let id_bytes = self.formation_id.as_bytes();
201        let mut bytes = Vec::with_capacity(2 + id_bytes.len() + FORMATION_CHALLENGE_SIZE);
202
203        bytes.extend_from_slice(&(id_bytes.len() as u16).to_le_bytes());
204        bytes.extend_from_slice(id_bytes);
205        bytes.extend_from_slice(&self.nonce);
206
207        bytes
208    }
209
210    /// Deserialize from bytes
211    pub fn from_bytes(bytes: &[u8]) -> Result<Self, SecurityError> {
212        if bytes.len() < 2 {
213            return Err(SecurityError::AuthenticationFailed(
214                "Challenge too short".to_string(),
215            ));
216        }
217
218        let id_len = u16::from_le_bytes([bytes[0], bytes[1]]) as usize;
219
220        if bytes.len() < 2 + id_len + FORMATION_CHALLENGE_SIZE {
221            return Err(SecurityError::AuthenticationFailed(
222                "Challenge truncated".to_string(),
223            ));
224        }
225
226        let formation_id = String::from_utf8(bytes[2..2 + id_len].to_vec()).map_err(|e| {
227            SecurityError::AuthenticationFailed(format!("Invalid formation ID: {}", e))
228        })?;
229
230        let mut nonce = [0u8; FORMATION_CHALLENGE_SIZE];
231        nonce.copy_from_slice(&bytes[2 + id_len..2 + id_len + FORMATION_CHALLENGE_SIZE]);
232
233        Ok(Self {
234            formation_id,
235            nonce,
236        })
237    }
238}
239
240/// Response message sent from initiator to responder
241#[derive(Debug, Clone)]
242pub struct FormationChallengeResponse {
243    /// HMAC response
244    pub response: [u8; FORMATION_RESPONSE_SIZE],
245}
246
247impl FormationChallengeResponse {
248    /// Serialize to bytes
249    pub fn to_bytes(&self) -> Vec<u8> {
250        self.response.to_vec()
251    }
252
253    /// Deserialize from bytes
254    pub fn from_bytes(bytes: &[u8]) -> Result<Self, SecurityError> {
255        if bytes.len() < FORMATION_RESPONSE_SIZE {
256            return Err(SecurityError::AuthenticationFailed(
257                "Response too short".to_string(),
258            ));
259        }
260
261        let mut response = [0u8; FORMATION_RESPONSE_SIZE];
262        response.copy_from_slice(&bytes[..FORMATION_RESPONSE_SIZE]);
263
264        Ok(Self { response })
265    }
266}
267
268/// Result of formation authentication handshake
269#[derive(Debug, Clone, Copy, PartialEq, Eq)]
270pub enum FormationAuthResult {
271    /// Authentication successful
272    Accepted,
273    /// Authentication failed (wrong key or formation)
274    Rejected,
275}
276
277impl FormationAuthResult {
278    /// Serialize to single byte
279    pub fn to_byte(self) -> u8 {
280        match self {
281            Self::Accepted => 0x01,
282            Self::Rejected => 0x00,
283        }
284    }
285
286    /// Deserialize from byte
287    pub fn from_byte(byte: u8) -> Self {
288        if byte == 0x01 {
289            Self::Accepted
290        } else {
291            Self::Rejected
292        }
293    }
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299
300    #[test]
301    fn test_formation_key_creation() {
302        let secret = [0x42u8; 32];
303        let key = FormationKey::new("alpha-company", &secret);
304
305        assert_eq!(key.formation_id(), "alpha-company");
306    }
307
308    #[test]
309    fn test_formation_key_from_base64() {
310        let secret = FormationKey::generate_secret();
311        let key = FormationKey::from_base64("test-formation", &secret).unwrap();
312
313        assert_eq!(key.formation_id(), "test-formation");
314    }
315
316    #[test]
317    fn test_formation_key_from_base64_invalid() {
318        let result = FormationKey::from_base64("test", "not-valid-base64!!!");
319        assert!(result.is_err());
320    }
321
322    #[test]
323    fn test_formation_key_from_base64_derives_key_for_non_32_bytes() {
324        use base64::{engine::general_purpose::STANDARD, Engine};
325
326        // Short secret (16 bytes) - should be derived via SHA-256
327        let short_secret = STANDARD.encode([0u8; 16]);
328        let result = FormationKey::from_base64("test", &short_secret);
329        assert!(result.is_ok(), "Short key should be derived via SHA-256");
330
331        // Long secret (138 bytes like EC key) - should be derived via SHA-256
332        let long_secret = STANDARD.encode([0xABu8; 138]);
333        let result = FormationKey::from_base64("test", &long_secret);
334        assert!(
335            result.is_ok(),
336            "Long key (EC format) should be derived via SHA-256"
337        );
338
339        // Verify that different inputs produce different keys
340        let key1 = FormationKey::from_base64("test", &short_secret).unwrap();
341        let key2 = FormationKey::from_base64("test", &long_secret).unwrap();
342
343        // Create challenges to verify keys are different
344        let (nonce, _) = key1.create_challenge();
345        let response1 = key1.respond_to_challenge(&nonce);
346        assert!(
347            !key2.verify_response(&nonce, &response1),
348            "Different input keys should produce different derived keys"
349        );
350    }
351
352    #[test]
353    fn test_challenge_response_success() {
354        let secret = [0x42u8; 32];
355        let key = FormationKey::new("alpha-company", &secret);
356
357        // Responder creates challenge
358        let (nonce, _expected) = key.create_challenge();
359
360        // Initiator responds (with same key)
361        let response = key.respond_to_challenge(&nonce);
362
363        // Responder verifies
364        assert!(key.verify_response(&nonce, &response));
365    }
366
367    #[test]
368    fn test_challenge_response_wrong_key() {
369        let secret1 = [0x42u8; 32];
370        let secret2 = [0x43u8; 32]; // Different secret
371
372        let key1 = FormationKey::new("alpha-company", &secret1);
373        let key2 = FormationKey::new("alpha-company", &secret2);
374
375        // Responder (key1) creates challenge
376        let (nonce, _expected) = key1.create_challenge();
377
378        // Initiator (key2) responds with wrong key
379        let response = key2.respond_to_challenge(&nonce);
380
381        // Responder verifies - should fail
382        assert!(!key1.verify_response(&nonce, &response));
383    }
384
385    #[test]
386    fn test_challenge_response_wrong_formation() {
387        let secret = [0x42u8; 32];
388
389        let key1 = FormationKey::new("alpha-company", &secret);
390        let key2 = FormationKey::new("bravo-company", &secret); // Different formation
391
392        // Responder (key1) creates challenge
393        let (nonce, _expected) = key1.create_challenge();
394
395        // Initiator (key2) responds with wrong formation
396        let response = key2.respond_to_challenge(&nonce);
397
398        // Responder verifies - should fail (different formation IDs in HMAC)
399        assert!(!key1.verify_response(&nonce, &response));
400    }
401
402    #[test]
403    fn test_different_nonces_produce_different_responses() {
404        let secret = [0x42u8; 32];
405        let key = FormationKey::new("alpha-company", &secret);
406
407        let (nonce1, _) = key.create_challenge();
408        let (nonce2, _) = key.create_challenge();
409
410        let response1 = key.respond_to_challenge(&nonce1);
411        let response2 = key.respond_to_challenge(&nonce2);
412
413        // Different nonces should produce different responses
414        assert_ne!(response1, response2);
415    }
416
417    #[test]
418    fn test_challenge_serialization() {
419        let mut nonce = [0u8; FORMATION_CHALLENGE_SIZE];
420        nonce[0] = 0x42;
421
422        let challenge = FormationChallenge {
423            formation_id: "test-formation".to_string(),
424            nonce,
425        };
426
427        let bytes = challenge.to_bytes();
428        let restored = FormationChallenge::from_bytes(&bytes).unwrap();
429
430        assert_eq!(challenge.formation_id, restored.formation_id);
431        assert_eq!(challenge.nonce, restored.nonce);
432    }
433
434    #[test]
435    fn test_response_serialization() {
436        let mut response_bytes = [0u8; FORMATION_RESPONSE_SIZE];
437        response_bytes[0] = 0x42;
438
439        let response = FormationChallengeResponse {
440            response: response_bytes,
441        };
442
443        let bytes = response.to_bytes();
444        let restored = FormationChallengeResponse::from_bytes(&bytes).unwrap();
445
446        assert_eq!(response.response, restored.response);
447    }
448
449    #[test]
450    fn test_auth_result_serialization() {
451        assert_eq!(
452            FormationAuthResult::from_byte(FormationAuthResult::Accepted.to_byte()),
453            FormationAuthResult::Accepted
454        );
455        assert_eq!(
456            FormationAuthResult::from_byte(FormationAuthResult::Rejected.to_byte()),
457            FormationAuthResult::Rejected
458        );
459    }
460
461    #[test]
462    fn test_generate_secret() {
463        let secret1 = FormationKey::generate_secret();
464        let secret2 = FormationKey::generate_secret();
465
466        // Should be different (random)
467        assert_ne!(secret1, secret2);
468
469        // Should be valid base64 that decodes to 32 bytes
470        use base64::{engine::general_purpose::STANDARD, Engine};
471        let decoded = STANDARD.decode(&secret1).unwrap();
472        assert_eq!(decoded.len(), 32);
473    }
474}