1#[cfg(not(feature = "std"))]
48use alloc::vec::Vec;
49use hashbrown::HashMap;
50
51use super::identity::{node_id_from_public_key, IdentityAttestation};
52use super::membership_token::{MembershipToken, MAX_CALLSIGN_LEN};
53use crate::NodeId;
54
55#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub enum RegistryResult {
58 Registered,
60
61 Verified,
63
64 InvalidSignature,
66
67 KeyMismatch {
69 node_id: NodeId,
71 },
72}
73
74impl RegistryResult {
75 pub fn is_trusted(&self) -> bool {
77 matches!(self, Self::Registered | Self::Verified)
78 }
79
80 pub fn is_violation(&self) -> bool {
82 matches!(self, Self::InvalidSignature | Self::KeyMismatch { .. })
83 }
84}
85
86#[derive(Debug, Clone)]
88pub struct IdentityRecord {
89 pub public_key: [u8; 32],
91
92 pub first_seen_ms: u64,
94
95 pub last_seen_ms: u64,
97
98 pub verification_count: u32,
100
101 pub callsign: Option<[u8; MAX_CALLSIGN_LEN]>,
104
105 pub token_expires_ms: Option<u64>,
107}
108
109impl IdentityRecord {
110 pub fn callsign_str(&self) -> Option<&str> {
112 self.callsign.as_ref().map(|cs| {
113 let len = cs.iter().position(|&b| b == 0).unwrap_or(MAX_CALLSIGN_LEN);
114 core::str::from_utf8(&cs[..len]).unwrap_or("")
115 })
116 }
117
118 pub fn is_token_expired(&self, now_ms: u64) -> bool {
120 match self.token_expires_ms {
121 Some(0) => false, Some(expires) => now_ms > expires,
123 None => false, }
125 }
126}
127
128#[derive(Debug, Clone)]
133pub struct IdentityRegistry {
134 known: HashMap<NodeId, IdentityRecord>,
136
137 max_identities: usize,
139}
140
141impl Default for IdentityRegistry {
142 fn default() -> Self {
143 Self::new()
144 }
145}
146
147impl IdentityRegistry {
148 pub const DEFAULT_MAX_IDENTITIES: usize = 256;
150
151 pub fn new() -> Self {
153 Self {
154 known: HashMap::new(),
155 max_identities: Self::DEFAULT_MAX_IDENTITIES,
156 }
157 }
158
159 pub fn with_capacity(max_identities: usize) -> Self {
161 Self {
162 known: HashMap::with_capacity(max_identities.min(64)),
163 max_identities,
164 }
165 }
166
167 pub fn verify_or_register(&mut self, attestation: &IdentityAttestation) -> RegistryResult {
174 self.verify_or_register_at(attestation, attestation.timestamp_ms)
175 }
176
177 pub fn verify_or_register_at(
179 &mut self,
180 attestation: &IdentityAttestation,
181 now_ms: u64,
182 ) -> RegistryResult {
183 if !attestation.verify() {
185 return RegistryResult::InvalidSignature;
186 }
187
188 let node_id = attestation.node_id;
189
190 if let Some(record) = self.known.get_mut(&node_id) {
192 if record.public_key == attestation.public_key {
194 record.last_seen_ms = now_ms;
196 record.verification_count = record.verification_count.saturating_add(1);
197 RegistryResult::Verified
198 } else {
199 RegistryResult::KeyMismatch { node_id }
201 }
202 } else {
203 if self.known.len() >= self.max_identities {
205 }
209
210 self.known.insert(
211 node_id,
212 IdentityRecord {
213 public_key: attestation.public_key,
214 first_seen_ms: now_ms,
215 last_seen_ms: now_ms,
216 verification_count: 1,
217 callsign: None,
218 token_expires_ms: None,
219 },
220 );
221 RegistryResult::Registered
222 }
223 }
224
225 pub fn is_known(&self, node_id: NodeId) -> bool {
227 self.known.contains_key(&node_id)
228 }
229
230 pub fn get_public_key(&self, node_id: NodeId) -> Option<&[u8; 32]> {
232 self.known.get(&node_id).map(|r| &r.public_key)
233 }
234
235 pub fn get_record(&self, node_id: NodeId) -> Option<&IdentityRecord> {
237 self.known.get(&node_id)
238 }
239
240 pub fn len(&self) -> usize {
242 self.known.len()
243 }
244
245 pub fn is_empty(&self) -> bool {
247 self.known.is_empty()
248 }
249
250 pub fn remove(&mut self, node_id: NodeId) -> Option<IdentityRecord> {
254 self.known.remove(&node_id)
255 }
256
257 pub fn clear(&mut self) {
261 self.known.clear();
262 }
263
264 pub fn known_nodes(&self) -> Vec<NodeId> {
266 self.known.keys().copied().collect()
267 }
268
269 pub fn pre_register(&mut self, node_id: NodeId, public_key: [u8; 32], now_ms: u64) {
274 self.known.insert(
275 node_id,
276 IdentityRecord {
277 public_key,
278 first_seen_ms: now_ms,
279 last_seen_ms: now_ms,
280 verification_count: 0,
281 callsign: None,
282 token_expires_ms: None,
283 },
284 );
285 }
286
287 pub fn register_member(
301 &mut self,
302 token: &MembershipToken,
303 authority_public_key: &[u8; 32],
304 now_ms: u64,
305 ) -> Result<NodeId, RegistryResult> {
306 if !token.verify(authority_public_key) {
308 return Err(RegistryResult::InvalidSignature);
309 }
310
311 if token.is_expired(now_ms) {
313 return Err(RegistryResult::InvalidSignature); }
315
316 let node_id = node_id_from_public_key(&token.public_key);
317
318 if let Some(existing) = self.known.get(&node_id) {
320 if existing.public_key != token.public_key {
321 return Err(RegistryResult::KeyMismatch { node_id });
322 }
323 }
324
325 self.known.insert(
327 node_id,
328 IdentityRecord {
329 public_key: token.public_key,
330 first_seen_ms: now_ms,
331 last_seen_ms: now_ms,
332 verification_count: 1,
333 callsign: Some(token.callsign),
334 token_expires_ms: Some(token.expires_at_ms),
335 },
336 );
337
338 Ok(node_id)
339 }
340
341 pub fn get_callsign(&self, node_id: NodeId) -> Option<&str> {
343 self.known.get(&node_id).and_then(|r| r.callsign_str())
344 }
345
346 pub fn find_by_callsign(&self, callsign: &str) -> Option<NodeId> {
348 for (node_id, record) in &self.known {
349 if let Some(cs) = record.callsign_str() {
350 if cs == callsign {
351 return Some(*node_id);
352 }
353 }
354 }
355 None
356 }
357
358 pub fn encode(&self) -> Vec<u8> {
373 let entry_size = 4 + 32 + 8 + 8 + 4 + 1 + MAX_CALLSIGN_LEN + 8; let mut buf = Vec::with_capacity(1 + 4 + self.known.len() * entry_size);
376
377 buf.push(2);
379
380 buf.extend_from_slice(&(self.known.len() as u32).to_le_bytes());
382
383 for (node_id, record) in &self.known {
384 buf.extend_from_slice(&node_id.as_u32().to_le_bytes());
385 buf.extend_from_slice(&record.public_key);
386 buf.extend_from_slice(&record.first_seen_ms.to_le_bytes());
387 buf.extend_from_slice(&record.last_seen_ms.to_le_bytes());
388 buf.extend_from_slice(&record.verification_count.to_le_bytes());
389
390 if let Some(callsign) = &record.callsign {
392 buf.push(1); buf.extend_from_slice(callsign);
394 buf.extend_from_slice(&record.token_expires_ms.unwrap_or(0).to_le_bytes());
395 } else {
396 buf.push(0); buf.extend_from_slice(&[0u8; MAX_CALLSIGN_LEN]);
398 buf.extend_from_slice(&0u64.to_le_bytes());
399 }
400 }
401
402 buf
403 }
404
405 pub fn decode(data: &[u8]) -> Option<Self> {
407 if data.is_empty() {
408 return None;
409 }
410
411 let version = data[0];
413
414 match version {
415 2 => Self::decode_v2(data),
416 _ => Self::decode_v1(data),
419 }
420 }
421
422 fn decode_v1(data: &[u8]) -> Option<Self> {
424 if data.len() < 4 {
425 return None;
426 }
427
428 let count = u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
429
430 if data.len() < 4 + count * 56 {
431 return None;
432 }
433
434 let mut registry = Self::new();
435 let mut offset = 4;
436
437 for _ in 0..count {
438 let node_id = NodeId::new(u32::from_le_bytes([
439 data[offset],
440 data[offset + 1],
441 data[offset + 2],
442 data[offset + 3],
443 ]));
444 offset += 4;
445
446 let mut public_key = [0u8; 32];
447 public_key.copy_from_slice(&data[offset..offset + 32]);
448 offset += 32;
449
450 let first_seen_ms = u64::from_le_bytes([
451 data[offset],
452 data[offset + 1],
453 data[offset + 2],
454 data[offset + 3],
455 data[offset + 4],
456 data[offset + 5],
457 data[offset + 6],
458 data[offset + 7],
459 ]);
460 offset += 8;
461
462 let last_seen_ms = u64::from_le_bytes([
463 data[offset],
464 data[offset + 1],
465 data[offset + 2],
466 data[offset + 3],
467 data[offset + 4],
468 data[offset + 5],
469 data[offset + 6],
470 data[offset + 7],
471 ]);
472 offset += 8;
473
474 let verification_count = u32::from_le_bytes([
475 data[offset],
476 data[offset + 1],
477 data[offset + 2],
478 data[offset + 3],
479 ]);
480 offset += 4;
481
482 registry.known.insert(
483 node_id,
484 IdentityRecord {
485 public_key,
486 first_seen_ms,
487 last_seen_ms,
488 verification_count,
489 callsign: None,
490 token_expires_ms: None,
491 },
492 );
493 }
494
495 Some(registry)
496 }
497
498 fn decode_v2(data: &[u8]) -> Option<Self> {
500 if data.len() < 5 {
501 return None;
502 }
503
504 let count = u32::from_le_bytes([data[1], data[2], data[3], data[4]]) as usize;
506 let entry_size = 77; if data.len() < 5 + count * entry_size {
509 return None;
510 }
511
512 let mut registry = Self::new();
513 let mut offset = 5;
514
515 for _ in 0..count {
516 let node_id = NodeId::new(u32::from_le_bytes([
517 data[offset],
518 data[offset + 1],
519 data[offset + 2],
520 data[offset + 3],
521 ]));
522 offset += 4;
523
524 let mut public_key = [0u8; 32];
525 public_key.copy_from_slice(&data[offset..offset + 32]);
526 offset += 32;
527
528 let first_seen_ms = u64::from_le_bytes([
529 data[offset],
530 data[offset + 1],
531 data[offset + 2],
532 data[offset + 3],
533 data[offset + 4],
534 data[offset + 5],
535 data[offset + 6],
536 data[offset + 7],
537 ]);
538 offset += 8;
539
540 let last_seen_ms = u64::from_le_bytes([
541 data[offset],
542 data[offset + 1],
543 data[offset + 2],
544 data[offset + 3],
545 data[offset + 4],
546 data[offset + 5],
547 data[offset + 6],
548 data[offset + 7],
549 ]);
550 offset += 8;
551
552 let verification_count = u32::from_le_bytes([
553 data[offset],
554 data[offset + 1],
555 data[offset + 2],
556 data[offset + 3],
557 ]);
558 offset += 4;
559
560 let has_callsign = data[offset] != 0;
561 offset += 1;
562
563 let (callsign, token_expires_ms) = if has_callsign {
564 let mut cs = [0u8; MAX_CALLSIGN_LEN];
565 cs.copy_from_slice(&data[offset..offset + MAX_CALLSIGN_LEN]);
566 offset += MAX_CALLSIGN_LEN;
567
568 let expires = u64::from_le_bytes([
569 data[offset],
570 data[offset + 1],
571 data[offset + 2],
572 data[offset + 3],
573 data[offset + 4],
574 data[offset + 5],
575 data[offset + 6],
576 data[offset + 7],
577 ]);
578 offset += 8;
579
580 (Some(cs), Some(expires))
581 } else {
582 offset += MAX_CALLSIGN_LEN + 8; (None, None)
584 };
585
586 registry.known.insert(
587 node_id,
588 IdentityRecord {
589 public_key,
590 first_seen_ms,
591 last_seen_ms,
592 verification_count,
593 callsign,
594 token_expires_ms,
595 },
596 );
597 }
598
599 Some(registry)
600 }
601}
602
603#[cfg(test)]
604mod tests {
605 use super::*;
606 use crate::security::DeviceIdentity;
607
608 #[test]
609 fn test_register_new_identity() {
610 let mut registry = IdentityRegistry::new();
611 let identity = DeviceIdentity::generate();
612 let attestation = identity.create_attestation(0);
613
614 let result = registry.verify_or_register(&attestation);
615 assert_eq!(result, RegistryResult::Registered);
616 assert!(result.is_trusted());
617 assert!(!result.is_violation());
618 assert_eq!(registry.len(), 1);
619 }
620
621 #[test]
622 fn test_verify_known_identity() {
623 let mut registry = IdentityRegistry::new();
624 let identity = DeviceIdentity::generate();
625 let attestation = identity.create_attestation(0);
626
627 registry.verify_or_register(&attestation);
629
630 let result = registry.verify_or_register(&attestation);
632 assert_eq!(result, RegistryResult::Verified);
633 assert!(result.is_trusted());
634 }
635
636 #[test]
637 fn test_key_mismatch_detection() {
638 let mut registry = IdentityRegistry::new();
639
640 let identity1 = DeviceIdentity::generate();
642 let attestation1 = identity1.create_attestation(0);
643 registry.verify_or_register(&attestation1);
644
645 let _identity2 = DeviceIdentity::generate();
651 let node_id = identity1.node_id();
652
653 registry.known.insert(
655 node_id,
656 IdentityRecord {
657 public_key: [0xAA; 32], first_seen_ms: 0,
659 last_seen_ms: 0,
660 verification_count: 1,
661 callsign: None,
662 token_expires_ms: None,
663 },
664 );
665
666 let result = registry.verify_or_register(&attestation1);
668 assert!(matches!(result, RegistryResult::KeyMismatch { .. }));
669 assert!(result.is_violation());
670 }
671
672 #[test]
673 fn test_invalid_signature_detection() {
674 let mut registry = IdentityRegistry::new();
675
676 let identity = DeviceIdentity::generate();
678 let mut attestation = identity.create_attestation(0);
679 attestation.signature[0] ^= 0xFF; let result = registry.verify_or_register(&attestation);
682 assert_eq!(result, RegistryResult::InvalidSignature);
683 assert!(result.is_violation());
684 }
685
686 #[test]
687 fn test_verification_count_increment() {
688 let mut registry = IdentityRegistry::new();
689 let identity = DeviceIdentity::generate();
690 let attestation = identity.create_attestation(0);
691 let node_id = identity.node_id();
692
693 registry.verify_or_register(&attestation);
695 registry.verify_or_register(&attestation);
696 registry.verify_or_register(&attestation);
697
698 let record = registry.get_record(node_id).unwrap();
699 assert_eq!(record.verification_count, 3);
700 }
701
702 #[test]
703 fn test_pre_register() {
704 let mut registry = IdentityRegistry::new();
705 let identity = DeviceIdentity::generate();
706 let node_id = identity.node_id();
707 let public_key = identity.public_key();
708
709 registry.pre_register(node_id, public_key, 1000);
711
712 assert!(registry.is_known(node_id));
713 assert_eq!(registry.get_public_key(node_id), Some(&public_key));
714
715 let attestation = identity.create_attestation(0);
717 let result = registry.verify_or_register(&attestation);
718 assert_eq!(result, RegistryResult::Verified);
719 }
720
721 #[test]
722 fn test_encode_decode_roundtrip() {
723 let mut registry = IdentityRegistry::new();
724
725 for _ in 0..5 {
727 let identity = DeviceIdentity::generate();
728 let attestation = identity.create_attestation(0);
729 registry.verify_or_register(&attestation);
730 }
731
732 let encoded = registry.encode();
733 let decoded = IdentityRegistry::decode(&encoded).unwrap();
734
735 assert_eq!(decoded.len(), registry.len());
736 for node_id in registry.known_nodes() {
737 assert!(decoded.is_known(node_id));
738 assert_eq!(
739 decoded.get_public_key(node_id),
740 registry.get_public_key(node_id)
741 );
742 }
743 }
744
745 #[test]
746 fn test_remove_identity() {
747 let mut registry = IdentityRegistry::new();
748 let identity = DeviceIdentity::generate();
749 let attestation = identity.create_attestation(0);
750 let node_id = identity.node_id();
751
752 registry.verify_or_register(&attestation);
753 assert!(registry.is_known(node_id));
754
755 registry.remove(node_id);
756 assert!(!registry.is_known(node_id));
757
758 let result = registry.verify_or_register(&attestation);
760 assert_eq!(result, RegistryResult::Registered);
761 }
762
763 #[test]
764 fn test_known_nodes() {
765 let mut registry = IdentityRegistry::new();
766 let mut expected_nodes = Vec::new();
767
768 for _ in 0..3 {
769 let identity = DeviceIdentity::generate();
770 let attestation = identity.create_attestation(0);
771 expected_nodes.push(identity.node_id());
772 registry.verify_or_register(&attestation);
773 }
774
775 let known = registry.known_nodes();
776 assert_eq!(known.len(), 3);
777 for node_id in expected_nodes {
778 assert!(known.contains(&node_id));
779 }
780 }
781
782 #[test]
783 fn test_register_member_with_token() {
784 use crate::security::{MembershipPolicy, MeshGenesis};
785
786 let mut registry = IdentityRegistry::new();
787 let authority = DeviceIdentity::generate();
788 let genesis = MeshGenesis::create("ALPHA", &authority, MembershipPolicy::Controlled);
789 let member = DeviceIdentity::generate();
790
791 let token = MembershipToken::issue(
792 &authority,
793 &genesis,
794 member.public_key(),
795 "BRAVO-07",
796 3600_000, );
798
799 let now = 1000u64;
800 let result = registry.register_member(&token, &authority.public_key(), now);
801 assert!(result.is_ok());
802
803 let node_id = result.unwrap();
804 assert!(registry.is_known(node_id));
805 assert_eq!(registry.get_callsign(node_id), Some("BRAVO-07"));
806 }
807
808 #[test]
809 fn test_find_by_callsign() {
810 use crate::security::{MembershipPolicy, MeshGenesis};
811
812 let mut registry = IdentityRegistry::new();
813 let authority = DeviceIdentity::generate();
814 let genesis = MeshGenesis::create("ALPHA", &authority, MembershipPolicy::Controlled);
815
816 let member1 = DeviceIdentity::generate();
818 let token1 =
819 MembershipToken::issue(&authority, &genesis, member1.public_key(), "ALPHA-01", 0);
820 let node1 = registry
821 .register_member(&token1, &authority.public_key(), 0)
822 .unwrap();
823
824 let member2 = DeviceIdentity::generate();
825 let token2 =
826 MembershipToken::issue(&authority, &genesis, member2.public_key(), "BRAVO-02", 0);
827 let _node2 = registry
828 .register_member(&token2, &authority.public_key(), 0)
829 .unwrap();
830
831 assert_eq!(registry.find_by_callsign("ALPHA-01"), Some(node1));
833 assert_eq!(registry.find_by_callsign("CHARLIE-03"), None);
834 }
835
836 #[test]
837 fn test_register_member_wrong_authority() {
838 use crate::security::{MembershipPolicy, MeshGenesis};
839
840 let mut registry = IdentityRegistry::new();
841 let authority = DeviceIdentity::generate();
842 let other = DeviceIdentity::generate();
843 let genesis = MeshGenesis::create("ALPHA", &authority, MembershipPolicy::Controlled);
844 let member = DeviceIdentity::generate();
845
846 let token =
847 MembershipToken::issue(&authority, &genesis, member.public_key(), "BRAVO-07", 0);
848
849 let result = registry.register_member(&token, &other.public_key(), 0);
851 assert!(matches!(result, Err(RegistryResult::InvalidSignature)));
852 }
853
854 #[test]
855 fn test_encode_decode_with_callsign() {
856 use crate::security::{MembershipPolicy, MeshGenesis};
857
858 let mut registry = IdentityRegistry::new();
859 let authority = DeviceIdentity::generate();
860 let genesis = MeshGenesis::create("ALPHA", &authority, MembershipPolicy::Controlled);
861
862 let member = DeviceIdentity::generate();
864 let token =
865 MembershipToken::issue(&authority, &genesis, member.public_key(), "ALPHA-01", 0);
866 let node_id = registry
867 .register_member(&token, &authority.public_key(), 0)
868 .unwrap();
869
870 let plain = DeviceIdentity::generate();
872 let attestation = plain.create_attestation(0);
873 registry.verify_or_register(&attestation);
874
875 let encoded = registry.encode();
877 let decoded = IdentityRegistry::decode(&encoded).unwrap();
878
879 assert_eq!(decoded.len(), 2);
880 assert_eq!(decoded.get_callsign(node_id), Some("ALPHA-01"));
881 assert_eq!(decoded.get_callsign(plain.node_id()), None);
882 }
883}