1use hmac::{Hmac, Mac};
35use rand_core::{OsRng, RngCore};
36use sha2::Sha256;
37
38use super::error::SecurityError;
39
40pub const FORMATION_CHALLENGE_SIZE: usize = 32;
42
43pub const FORMATION_RESPONSE_SIZE: usize = 32;
45
46const HKDF_INFO_FORMATION: &[u8] = b"peat-protocol-v1-formation";
48
49type HmacSha256 = Hmac<Sha256>;
50
51#[derive(Clone)]
56pub struct FormationKey {
57 formation_id: String,
59 hmac_key: [u8; 32],
61}
62
63impl FormationKey {
64 pub fn new(formation_id: &str, shared_secret: &[u8; 32]) -> Self {
70 use hkdf::Hkdf;
71
72 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 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 let mut arr = [0u8; 32];
102 arr.copy_from_slice(&secret_bytes);
103 arr
104 } else {
105 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 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 pub fn formation_id(&self) -> &str {
127 &self.formation_id
128 }
129
130 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 pub fn respond_to_challenge(&self, nonce: &[u8]) -> [u8; FORMATION_RESPONSE_SIZE] {
148 self.compute_response(nonce)
149 }
150
151 pub fn verify_response(&self, nonce: &[u8], response: &[u8; FORMATION_RESPONSE_SIZE]) -> bool {
153 let expected = self.compute_response(nonce);
154
155 use subtle::ConstantTimeEq;
157 expected.ct_eq(response).into()
158 }
159
160 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#[derive(Debug, Clone)]
188pub struct FormationChallenge {
189 pub formation_id: String,
191 pub nonce: [u8; FORMATION_CHALLENGE_SIZE],
193}
194
195impl FormationChallenge {
196 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 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#[derive(Debug, Clone)]
242pub struct FormationChallengeResponse {
243 pub response: [u8; FORMATION_RESPONSE_SIZE],
245}
246
247impl FormationChallengeResponse {
248 pub fn to_bytes(&self) -> Vec<u8> {
250 self.response.to_vec()
251 }
252
253 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
270pub enum FormationAuthResult {
271 Accepted,
273 Rejected,
275}
276
277impl FormationAuthResult {
278 pub fn to_byte(self) -> u8 {
280 match self {
281 Self::Accepted => 0x01,
282 Self::Rejected => 0x00,
283 }
284 }
285
286 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 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 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 let key1 = FormationKey::from_base64("test", &short_secret).unwrap();
341 let key2 = FormationKey::from_base64("test", &long_secret).unwrap();
342
343 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 let (nonce, _expected) = key.create_challenge();
359
360 let response = key.respond_to_challenge(&nonce);
362
363 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]; let key1 = FormationKey::new("alpha-company", &secret1);
373 let key2 = FormationKey::new("alpha-company", &secret2);
374
375 let (nonce, _expected) = key1.create_challenge();
377
378 let response = key2.respond_to_challenge(&nonce);
380
381 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); let (nonce, _expected) = key1.create_challenge();
394
395 let response = key2.respond_to_challenge(&nonce);
397
398 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 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 assert_ne!(secret1, secret2);
468
469 use base64::{engine::general_purpose::STANDARD, Engine};
471 let decoded = STANDARD.decode(&secret1).unwrap();
472 assert_eq!(decoded.len(), 32);
473 }
474
475 #[test]
478 fn test_wire_protocol_accept_with_matching_key() {
479 let secret = [0x42u8; 32];
480 let acceptor_key = FormationKey::new("test-formation", &secret);
481 let connector_key = FormationKey::new("test-formation", &secret);
482
483 let (nonce, _) = acceptor_key.create_challenge();
485 let challenge = FormationChallenge {
486 formation_id: acceptor_key.formation_id().to_string(),
487 nonce,
488 };
489 let challenge_bytes = challenge.to_bytes();
490
491 let decoded_challenge = FormationChallenge::from_bytes(&challenge_bytes).unwrap();
493 assert_eq!(decoded_challenge.formation_id, "test-formation");
494 let response = connector_key.respond_to_challenge(&decoded_challenge.nonce);
495 let resp = FormationChallengeResponse { response };
496 let resp_bytes = resp.to_bytes();
497
498 let decoded_resp = FormationChallengeResponse::from_bytes(&resp_bytes).unwrap();
500 assert!(
501 acceptor_key.verify_response(&nonce, &decoded_resp.response),
502 "Matching keys should produce accepted auth"
503 );
504 }
505
506 #[test]
508 fn test_wire_protocol_reject_with_wrong_key() {
509 let acceptor_key = FormationKey::new("test-formation", &[0x42u8; 32]);
510 let connector_key = FormationKey::new("test-formation", &[0xFF; 32]); let (nonce, _) = acceptor_key.create_challenge();
514 let challenge = FormationChallenge {
515 formation_id: acceptor_key.formation_id().to_string(),
516 nonce,
517 };
518 let challenge_bytes = challenge.to_bytes();
519
520 let decoded_challenge = FormationChallenge::from_bytes(&challenge_bytes).unwrap();
522 let response = connector_key.respond_to_challenge(&decoded_challenge.nonce);
523 let resp = FormationChallengeResponse { response };
524 let resp_bytes = resp.to_bytes();
525
526 let decoded_resp = FormationChallengeResponse::from_bytes(&resp_bytes).unwrap();
528 assert!(
529 !acceptor_key.verify_response(&nonce, &decoded_resp.response),
530 "Wrong key should produce rejected auth"
531 );
532 }
533
534 #[test]
537 fn test_wire_protocol_formation_id_mismatch() {
538 let acceptor_key = FormationKey::new("alpha", &[0x42u8; 32]);
539 let connector_key = FormationKey::new("bravo", &[0x42u8; 32]);
540
541 let (nonce, _) = acceptor_key.create_challenge();
542 let challenge = FormationChallenge {
543 formation_id: acceptor_key.formation_id().to_string(),
544 nonce,
545 };
546 let challenge_bytes = challenge.to_bytes();
547
548 let decoded_challenge = FormationChallenge::from_bytes(&challenge_bytes).unwrap();
549
550 assert_ne!(
552 decoded_challenge.formation_id,
553 connector_key.formation_id(),
554 "Connector should detect formation ID mismatch before responding"
555 );
556 }
557}