saorsa_core/
security.rs

1// Copyright 2024 Saorsa Labs Limited
2//
3// This software is dual-licensed under:
4// - GNU Affero General Public License v3.0 or later (AGPL-3.0-or-later)
5// - Commercial License
6//
7// For AGPL-3.0 license, see LICENSE-AGPL-3.0
8// For commercial licensing, contact: saorsalabs@gmail.com
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under these licenses is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
14//! Security module
15//!
16//! This module provides cryptographic functionality and Sybil protection for the P2P network.
17//! It implements IPv6-based node ID generation and IP diversity enforcement to prevent
18//! large-scale Sybil attacks while maintaining network openness.
19
20use crate::PeerId;
21use crate::quantum_crypto::ant_quic_integration::{
22    MlDsaPublicKey, MlDsaSecretKey, MlDsaSignature, ml_dsa_sign, ml_dsa_verify,
23};
24use anyhow::{Result, anyhow};
25use serde::{Deserialize, Serialize};
26use sha2::{Digest, Sha256};
27use std::collections::HashMap;
28use std::net::{Ipv4Addr, Ipv6Addr};
29use std::time::{Duration, SystemTime, UNIX_EPOCH};
30
31use std::sync::Arc;
32
33/// IPv6-based node identity that binds node ID to actual network location
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct IPv6NodeID {
36    /// Derived node ID (SHA256 of ipv6_addr + public_key + salt)
37    pub node_id: Vec<u8>,
38    /// IPv6 address this node ID is bound to
39    pub ipv6_addr: Ipv6Addr,
40    /// ML-DSA public key for signatures
41    pub public_key: Vec<u8>,
42    /// Signature proving ownership of the IPv6 address and keys
43    pub signature: Vec<u8>,
44    /// Timestamp when this ID was generated (seconds since epoch)
45    pub timestamp_secs: u64,
46    /// Salt used in node ID generation (for freshness)
47    pub salt: Vec<u8>,
48}
49
50/// Configuration for IP diversity enforcement at multiple subnet levels
51#[derive(Debug, Clone)]
52pub struct IPDiversityConfig {
53    /// Maximum nodes per /64 subnet (default: 1)
54    pub max_nodes_per_64: usize,
55    /// Maximum nodes per /48 allocation (default: 3)
56    pub max_nodes_per_48: usize,
57    /// Maximum nodes per /32 region (default: 10)
58    pub max_nodes_per_32: usize,
59    /// Maximum nodes per AS number (default: 20)
60    pub max_nodes_per_asn: usize,
61    /// Enable GeoIP-based diversity checks
62    pub enable_geolocation_check: bool,
63    /// Minimum number of different countries required
64    pub min_geographic_diversity: usize,
65}
66
67/// Analysis of an IPv6 address for diversity enforcement
68#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
69pub struct IPAnalysis {
70    /// /64 subnet (host allocation)
71    pub subnet_64: Ipv6Addr,
72    /// /48 subnet (site allocation)
73    pub subnet_48: Ipv6Addr,
74    /// /32 subnet (ISP allocation)
75    pub subnet_32: Ipv6Addr,
76    /// Autonomous System Number (if available)
77    pub asn: Option<u32>,
78    /// Country code from GeoIP lookup
79    pub country: Option<String>,
80    /// Whether this is a known hosting/VPS provider
81    pub is_hosting_provider: bool,
82    /// Whether this is a known VPN provider
83    pub is_vpn_provider: bool,
84    /// Historical reputation score for this IP range
85    pub reputation_score: f64,
86}
87
88/// Node reputation tracking for security-aware routing
89#[derive(Debug, Clone)]
90pub struct NodeReputation {
91    /// Peer ID
92    pub peer_id: PeerId,
93    /// Fraction of queries answered successfully
94    pub response_rate: f64,
95    /// Average response time
96    pub response_time: Duration,
97    /// Consistency of provided data (0.0-1.0)
98    pub consistency_score: f64,
99    /// Estimated continuous uptime
100    pub uptime_estimate: Duration,
101    /// Accuracy of routing information provided
102    pub routing_accuracy: f64,
103    /// Last time this node was seen
104    pub last_seen: SystemTime,
105    /// Total number of interactions
106    pub interaction_count: u64,
107}
108
109impl Default for IPDiversityConfig {
110    fn default() -> Self {
111        Self {
112            max_nodes_per_64: 1,
113            max_nodes_per_48: 3,
114            max_nodes_per_32: 10,
115            max_nodes_per_asn: 20,
116            enable_geolocation_check: true,
117            min_geographic_diversity: 3,
118        }
119    }
120}
121
122impl IPDiversityConfig {
123    /// Create a testnet configuration with relaxed diversity requirements.
124    ///
125    /// This is useful for testing environments like Digital Ocean where all nodes
126    /// share the same ASN (AS14061). The relaxed limits allow many nodes from the
127    /// same provider while still maintaining some diversity tracking.
128    ///
129    /// # Warning
130    ///
131    /// This configuration should NEVER be used in production as it significantly
132    /// weakens Sybil attack protection.
133    #[must_use]
134    pub fn testnet() -> Self {
135        Self {
136            max_nodes_per_64: 100,           // Allow many nodes per /64 subnet
137            max_nodes_per_48: 500,           // Allow many nodes per /48 allocation
138            max_nodes_per_32: 1000,          // Allow many nodes per /32 region
139            max_nodes_per_asn: 5000,         // Allow many nodes from same ASN (e.g., Digital Ocean)
140            enable_geolocation_check: false, // Disable geo checks for testing
141            min_geographic_diversity: 1,     // Single region is acceptable for testing
142        }
143    }
144
145    /// Create a permissive configuration that effectively disables diversity checks.
146    ///
147    /// This is useful for local development and unit testing where all nodes
148    /// run on localhost or the same machine.
149    #[must_use]
150    pub fn permissive() -> Self {
151        Self {
152            max_nodes_per_64: usize::MAX,
153            max_nodes_per_48: usize::MAX,
154            max_nodes_per_32: usize::MAX,
155            max_nodes_per_asn: usize::MAX,
156            enable_geolocation_check: false,
157            min_geographic_diversity: 0,
158        }
159    }
160
161    /// Check if this is a testnet or permissive configuration.
162    #[must_use]
163    pub fn is_relaxed(&self) -> bool {
164        self.max_nodes_per_asn > 100 || !self.enable_geolocation_check
165    }
166}
167
168impl IPv6NodeID {
169    /// Generate a new IPv6-based node ID
170    pub fn generate(
171        ipv6_addr: Ipv6Addr,
172        secret: &MlDsaSecretKey,
173        public: &MlDsaPublicKey,
174    ) -> Result<Self> {
175        let mut rng = rand::thread_rng();
176        let mut salt = vec![0u8; 16];
177        rand::RngCore::fill_bytes(&mut rng, &mut salt);
178
179        let timestamp = SystemTime::now();
180        let timestamp_secs = timestamp.duration_since(UNIX_EPOCH)?.as_secs();
181        let public_key = public.as_bytes().to_vec();
182
183        // Generate node ID: SHA256(ipv6_address || public_key || salt || timestamp)
184        let mut hasher = Sha256::new();
185        hasher.update(ipv6_addr.octets());
186        hasher.update(&public_key);
187        hasher.update(&salt);
188        hasher.update(timestamp_secs.to_le_bytes());
189        let node_id = hasher.finalize().to_vec();
190
191        // Create signature proving ownership
192        let mut message_to_sign = Vec::new();
193        message_to_sign.extend_from_slice(&ipv6_addr.octets());
194        message_to_sign.extend_from_slice(&public_key);
195        message_to_sign.extend_from_slice(&salt);
196        message_to_sign.extend_from_slice(&timestamp_secs.to_le_bytes());
197
198        let sig = ml_dsa_sign(secret, &message_to_sign)
199            .map_err(|e| anyhow!("ML-DSA sign failed: {:?}", e))?;
200        let signature = sig.0.to_vec();
201
202        Ok(IPv6NodeID {
203            node_id,
204            ipv6_addr,
205            public_key,
206            signature,
207            timestamp_secs,
208            salt,
209        })
210    }
211
212    /// Verify that this node ID is valid and properly signed
213    pub fn verify(&self) -> Result<bool> {
214        // Reconstruct the node ID
215        let mut hasher = Sha256::new();
216        hasher.update(self.ipv6_addr.octets());
217        hasher.update(&self.public_key);
218        hasher.update(&self.salt);
219        hasher.update(self.timestamp_secs.to_le_bytes());
220        let expected_node_id = hasher.finalize();
221
222        // Verify node ID matches
223        if expected_node_id.as_slice() != self.node_id {
224            return Ok(false);
225        }
226
227        let public_key = MlDsaPublicKey::from_bytes(&self.public_key)
228            .map_err(|e| anyhow!("Invalid ML-DSA public key: {:?}", e))?;
229
230        // Convert Vec<u8> to fixed-size array for MlDsaSignature
231        if self.signature.len() != 3309 {
232            return Ok(false); // Invalid signature length
233        }
234        let mut sig_bytes = [0u8; 3309];
235        sig_bytes.copy_from_slice(&self.signature);
236        let signature = MlDsaSignature(Box::new(sig_bytes));
237
238        let mut message_to_verify = Vec::new();
239        message_to_verify.extend_from_slice(&self.ipv6_addr.octets());
240        message_to_verify.extend_from_slice(&self.public_key);
241        message_to_verify.extend_from_slice(&self.salt);
242        message_to_verify.extend_from_slice(&self.timestamp_secs.to_le_bytes());
243
244        let ok = ml_dsa_verify(&public_key, &message_to_verify, &signature)
245            .map_err(|e| anyhow!("ML-DSA verify error: {:?}", e))?;
246        Ok(ok)
247    }
248
249    /// Extract /64 subnet from IPv6 address
250    pub fn extract_subnet_64(&self) -> Ipv6Addr {
251        let octets = self.ipv6_addr.octets();
252        let mut subnet = [0u8; 16];
253        subnet[..8].copy_from_slice(&octets[..8]); // Keep first 64 bits, zero the rest
254        Ipv6Addr::from(subnet)
255    }
256
257    /// Extract /48 subnet from IPv6 address
258    pub fn extract_subnet_48(&self) -> Ipv6Addr {
259        let octets = self.ipv6_addr.octets();
260        let mut subnet = [0u8; 16];
261        subnet[..6].copy_from_slice(&octets[..6]); // Keep first 48 bits, zero the rest
262        Ipv6Addr::from(subnet)
263    }
264
265    /// Extract /32 subnet from IPv6 address
266    pub fn extract_subnet_32(&self) -> Ipv6Addr {
267        let octets = self.ipv6_addr.octets();
268        let mut subnet = [0u8; 16];
269        subnet[..4].copy_from_slice(&octets[..4]); // Keep first 32 bits, zero the rest
270        Ipv6Addr::from(subnet)
271    }
272}
273
274/// IPv4-based node identity that binds node ID to actual network location
275/// Mirrors IPv6NodeID for security parity on IPv4 networks
276#[derive(Debug, Clone, Serialize, Deserialize)]
277pub struct IPv4NodeID {
278    /// Derived node ID (SHA256 of ipv4_addr + public_key + salt + timestamp)
279    pub node_id: Vec<u8>,
280    /// IPv4 address this node ID is bound to
281    pub ipv4_addr: Ipv4Addr,
282    /// ML-DSA public key for signatures
283    pub public_key: Vec<u8>,
284    /// Signature proving ownership of the IPv4 address and keys
285    pub signature: Vec<u8>,
286    /// Timestamp when this ID was generated (seconds since epoch)
287    pub timestamp_secs: u64,
288    /// Salt used in node ID generation (for freshness)
289    pub salt: Vec<u8>,
290}
291
292impl IPv4NodeID {
293    /// Generate a new IPv4-based node ID
294    pub fn generate(
295        ipv4_addr: Ipv4Addr,
296        secret: &MlDsaSecretKey,
297        public: &MlDsaPublicKey,
298    ) -> Result<Self> {
299        let mut rng = rand::thread_rng();
300        let mut salt = vec![0u8; 16];
301        rand::RngCore::fill_bytes(&mut rng, &mut salt);
302
303        let timestamp = SystemTime::now();
304        let timestamp_secs = timestamp.duration_since(UNIX_EPOCH)?.as_secs();
305        let public_key = public.as_bytes().to_vec();
306
307        // Generate node ID: SHA256(ipv4_address || public_key || salt || timestamp)
308        let mut hasher = Sha256::new();
309        hasher.update(ipv4_addr.octets());
310        hasher.update(&public_key);
311        hasher.update(&salt);
312        hasher.update(timestamp_secs.to_le_bytes());
313        let node_id = hasher.finalize().to_vec();
314
315        // Create signature proving ownership
316        let mut message_to_sign = Vec::new();
317        message_to_sign.extend_from_slice(&ipv4_addr.octets());
318        message_to_sign.extend_from_slice(&public_key);
319        message_to_sign.extend_from_slice(&salt);
320        message_to_sign.extend_from_slice(&timestamp_secs.to_le_bytes());
321
322        let sig = ml_dsa_sign(secret, &message_to_sign)
323            .map_err(|e| anyhow!("ML-DSA sign failed: {:?}", e))?;
324        let signature = sig.0.to_vec();
325
326        Ok(IPv4NodeID {
327            node_id,
328            ipv4_addr,
329            public_key,
330            signature,
331            timestamp_secs,
332            salt,
333        })
334    }
335
336    /// Verify that this node ID is valid and properly signed
337    pub fn verify(&self) -> Result<bool> {
338        // Reconstruct the node ID
339        let mut hasher = Sha256::new();
340        hasher.update(self.ipv4_addr.octets());
341        hasher.update(&self.public_key);
342        hasher.update(&self.salt);
343        hasher.update(self.timestamp_secs.to_le_bytes());
344        let expected_node_id = hasher.finalize();
345
346        // Verify node ID matches
347        if expected_node_id.as_slice() != self.node_id {
348            return Ok(false);
349        }
350
351        let public_key = MlDsaPublicKey::from_bytes(&self.public_key)
352            .map_err(|e| anyhow!("Invalid ML-DSA public key: {:?}", e))?;
353
354        // Convert Vec<u8> to fixed-size array for MlDsaSignature
355        if self.signature.len() != 3309 {
356            return Ok(false); // Invalid signature length
357        }
358        let mut sig_bytes = [0u8; 3309];
359        sig_bytes.copy_from_slice(&self.signature);
360        let signature = MlDsaSignature(Box::new(sig_bytes));
361
362        let mut message_to_verify = Vec::new();
363        message_to_verify.extend_from_slice(&self.ipv4_addr.octets());
364        message_to_verify.extend_from_slice(&self.public_key);
365        message_to_verify.extend_from_slice(&self.salt);
366        message_to_verify.extend_from_slice(&self.timestamp_secs.to_le_bytes());
367
368        let ok = ml_dsa_verify(&public_key, &message_to_verify, &signature)
369            .map_err(|e| anyhow!("ML-DSA verify error: {:?}", e))?;
370        Ok(ok)
371    }
372
373    /// Extract /24 subnet from IPv4 address (Class C / most ISP allocations)
374    pub fn extract_subnet_24(&self) -> Ipv4Addr {
375        let octets = self.ipv4_addr.octets();
376        Ipv4Addr::new(octets[0], octets[1], octets[2], 0)
377    }
378
379    /// Extract /16 subnet from IPv4 address (Class B / large ISP allocations)
380    pub fn extract_subnet_16(&self) -> Ipv4Addr {
381        let octets = self.ipv4_addr.octets();
382        Ipv4Addr::new(octets[0], octets[1], 0, 0)
383    }
384
385    /// Extract /8 subnet from IPv4 address (Class A / regional allocations)
386    pub fn extract_subnet_8(&self) -> Ipv4Addr {
387        let octets = self.ipv4_addr.octets();
388        Ipv4Addr::new(octets[0], 0, 0, 0)
389    }
390
391    /// Get the age of this node ID in seconds
392    pub fn age_secs(&self) -> u64 {
393        let now = SystemTime::now()
394            .duration_since(UNIX_EPOCH)
395            .map(|d| d.as_secs())
396            .unwrap_or(0);
397        now.saturating_sub(self.timestamp_secs)
398    }
399
400    /// Check if the node ID has expired (older than max_age)
401    pub fn is_expired(&self, max_age: Duration) -> bool {
402        self.age_secs() > max_age.as_secs()
403    }
404}
405
406/// IP diversity enforcement system
407#[derive(Debug)]
408pub struct IPDiversityEnforcer {
409    config: IPDiversityConfig,
410    subnet_64_counts: HashMap<Ipv6Addr, usize>,
411    subnet_48_counts: HashMap<Ipv6Addr, usize>,
412    subnet_32_counts: HashMap<Ipv6Addr, usize>,
413    asn_counts: HashMap<u32, usize>,
414    country_counts: HashMap<String, usize>,
415    geo_provider: Option<Arc<dyn GeoProvider + Send + Sync>>,
416}
417
418impl IPDiversityEnforcer {
419    /// Create a new IP diversity enforcer
420    pub fn new(config: IPDiversityConfig) -> Self {
421        Self {
422            config,
423            subnet_64_counts: HashMap::new(),
424            subnet_48_counts: HashMap::new(),
425            subnet_32_counts: HashMap::new(),
426            asn_counts: HashMap::new(),
427            country_counts: HashMap::new(),
428            geo_provider: None,
429        }
430    }
431
432    /// Create a new IP diversity enforcer with a GeoIP/ASN provider
433    pub fn with_geo_provider(
434        config: IPDiversityConfig,
435        provider: Arc<dyn GeoProvider + Send + Sync>,
436    ) -> Self {
437        let mut s = Self::new(config);
438        s.geo_provider = Some(provider);
439        s
440    }
441
442    /// Analyze an IPv6 address for diversity enforcement
443    pub fn analyze_ip(&self, ipv6_addr: Ipv6Addr) -> Result<IPAnalysis> {
444        let subnet_64 = Self::extract_subnet_prefix(ipv6_addr, 64);
445        let subnet_48 = Self::extract_subnet_prefix(ipv6_addr, 48);
446        let subnet_32 = Self::extract_subnet_prefix(ipv6_addr, 32);
447
448        // GeoIP/ASN lookup via provider if available
449        let (asn, country, is_hosting_provider, is_vpn_provider) =
450            if let Some(p) = &self.geo_provider {
451                let info = p.lookup(ipv6_addr);
452                (
453                    info.asn,
454                    info.country,
455                    info.is_hosting_provider,
456                    info.is_vpn_provider,
457                )
458            } else {
459                (None, None, false, false)
460            };
461
462        // Default reputation for new IPs
463        let reputation_score = 0.5;
464
465        Ok(IPAnalysis {
466            subnet_64,
467            subnet_48,
468            subnet_32,
469            asn,
470            country,
471            is_hosting_provider,
472            is_vpn_provider,
473            reputation_score,
474        })
475    }
476
477    /// Check if a new node can be accepted based on IP diversity constraints
478    pub fn can_accept_node(&self, ip_analysis: &IPAnalysis) -> bool {
479        // Determine limits based on hosting provider status
480        let (limit_64, limit_48, limit_32, limit_asn) =
481            if ip_analysis.is_hosting_provider || ip_analysis.is_vpn_provider {
482                // Stricter limits for hosting providers (halved)
483                (
484                    std::cmp::max(1, self.config.max_nodes_per_64 / 2),
485                    std::cmp::max(1, self.config.max_nodes_per_48 / 2),
486                    std::cmp::max(1, self.config.max_nodes_per_32 / 2),
487                    std::cmp::max(1, self.config.max_nodes_per_asn / 2),
488                )
489            } else {
490                // Regular limits for normal nodes
491                (
492                    self.config.max_nodes_per_64,
493                    self.config.max_nodes_per_48,
494                    self.config.max_nodes_per_32,
495                    self.config.max_nodes_per_asn,
496                )
497            };
498
499        // Check /64 subnet limit
500        if let Some(&count) = self.subnet_64_counts.get(&ip_analysis.subnet_64)
501            && count >= limit_64
502        {
503            return false;
504        }
505
506        // Check /48 subnet limit
507        if let Some(&count) = self.subnet_48_counts.get(&ip_analysis.subnet_48)
508            && count >= limit_48
509        {
510            return false;
511        }
512
513        // Check /32 subnet limit
514        if let Some(&count) = self.subnet_32_counts.get(&ip_analysis.subnet_32)
515            && count >= limit_32
516        {
517            return false;
518        }
519
520        // Check ASN limit
521        if let Some(asn) = ip_analysis.asn
522            && let Some(&count) = self.asn_counts.get(&asn)
523            && count >= limit_asn
524        {
525            return false;
526        }
527
528        true
529    }
530
531    /// Add a node to the diversity tracking
532    pub fn add_node(&mut self, ip_analysis: &IPAnalysis) -> Result<()> {
533        if !self.can_accept_node(ip_analysis) {
534            return Err(anyhow!("IP diversity limits exceeded"));
535        }
536
537        // Update counts
538        *self
539            .subnet_64_counts
540            .entry(ip_analysis.subnet_64)
541            .or_insert(0) += 1;
542        *self
543            .subnet_48_counts
544            .entry(ip_analysis.subnet_48)
545            .or_insert(0) += 1;
546        *self
547            .subnet_32_counts
548            .entry(ip_analysis.subnet_32)
549            .or_insert(0) += 1;
550
551        if let Some(asn) = ip_analysis.asn {
552            *self.asn_counts.entry(asn).or_insert(0) += 1;
553        }
554
555        if let Some(ref country) = ip_analysis.country {
556            *self.country_counts.entry(country.clone()).or_insert(0) += 1;
557        }
558
559        Ok(())
560    }
561
562    /// Remove a node from diversity tracking
563    pub fn remove_node(&mut self, ip_analysis: &IPAnalysis) {
564        if let Some(count) = self.subnet_64_counts.get_mut(&ip_analysis.subnet_64) {
565            *count = count.saturating_sub(1);
566            if *count == 0 {
567                self.subnet_64_counts.remove(&ip_analysis.subnet_64);
568            }
569        }
570
571        if let Some(count) = self.subnet_48_counts.get_mut(&ip_analysis.subnet_48) {
572            *count = count.saturating_sub(1);
573            if *count == 0 {
574                self.subnet_48_counts.remove(&ip_analysis.subnet_48);
575            }
576        }
577
578        if let Some(count) = self.subnet_32_counts.get_mut(&ip_analysis.subnet_32) {
579            *count = count.saturating_sub(1);
580            if *count == 0 {
581                self.subnet_32_counts.remove(&ip_analysis.subnet_32);
582            }
583        }
584
585        if let Some(asn) = ip_analysis.asn
586            && let Some(count) = self.asn_counts.get_mut(&asn)
587        {
588            *count = count.saturating_sub(1);
589            if *count == 0 {
590                self.asn_counts.remove(&asn);
591            }
592        }
593
594        if let Some(ref country) = ip_analysis.country
595            && let Some(count) = self.country_counts.get_mut(country)
596        {
597            *count = count.saturating_sub(1);
598            if *count == 0 {
599                self.country_counts.remove(country);
600            }
601        }
602    }
603
604    /// Extract network prefix of specified length from IPv6 address
605    pub fn extract_subnet_prefix(addr: Ipv6Addr, prefix_len: u8) -> Ipv6Addr {
606        let octets = addr.octets();
607        let mut subnet = [0u8; 16];
608
609        let bytes_to_copy = (prefix_len / 8) as usize;
610        let remaining_bits = prefix_len % 8;
611
612        // Copy full bytes
613        if bytes_to_copy < 16 {
614            subnet[..bytes_to_copy].copy_from_slice(&octets[..bytes_to_copy]);
615        } else {
616            subnet.copy_from_slice(&octets);
617        }
618
619        // Handle partial byte
620        if remaining_bits > 0 && bytes_to_copy < 16 {
621            let mask = 0xFF << (8 - remaining_bits);
622            subnet[bytes_to_copy] = octets[bytes_to_copy] & mask;
623        }
624
625        Ipv6Addr::from(subnet)
626    }
627
628    /// Get diversity statistics
629    pub fn get_diversity_stats(&self) -> DiversityStats {
630        DiversityStats {
631            total_64_subnets: self.subnet_64_counts.len(),
632            total_48_subnets: self.subnet_48_counts.len(),
633            total_32_subnets: self.subnet_32_counts.len(),
634            total_asns: self.asn_counts.len(),
635            total_countries: self.country_counts.len(),
636            max_nodes_per_64: self.subnet_64_counts.values().max().copied().unwrap_or(0),
637            max_nodes_per_48: self.subnet_48_counts.values().max().copied().unwrap_or(0),
638            max_nodes_per_32: self.subnet_32_counts.values().max().copied().unwrap_or(0),
639        }
640    }
641}
642
643/// GeoIP/ASN provider trait
644pub trait GeoProvider: std::fmt::Debug {
645    fn lookup(&self, ip: Ipv6Addr) -> GeoInfo;
646}
647
648/// Geo information
649#[derive(Debug, Clone)]
650pub struct GeoInfo {
651    pub asn: Option<u32>,
652    pub country: Option<String>,
653    pub is_hosting_provider: bool,
654    pub is_vpn_provider: bool,
655}
656
657/// A simple in-memory caching wrapper for a GeoProvider
658#[derive(Debug)]
659pub struct CachedGeoProvider<P: GeoProvider> {
660    inner: P,
661    cache: parking_lot::RwLock<HashMap<Ipv6Addr, GeoInfo>>,
662}
663
664impl<P: GeoProvider> CachedGeoProvider<P> {
665    pub fn new(inner: P) -> Self {
666        Self {
667            inner,
668            cache: parking_lot::RwLock::new(HashMap::new()),
669        }
670    }
671}
672
673impl<P: GeoProvider> GeoProvider for CachedGeoProvider<P> {
674    fn lookup(&self, ip: Ipv6Addr) -> GeoInfo {
675        if let Some(info) = self.cache.read().get(&ip).cloned() {
676            return info;
677        }
678        let info = self.inner.lookup(ip);
679        self.cache.write().insert(ip, info.clone());
680        info
681    }
682}
683
684/// Stub provider returning no ASN/GeoIP info
685#[derive(Debug)]
686pub struct StubGeoProvider;
687impl GeoProvider for StubGeoProvider {
688    fn lookup(&self, _ip: Ipv6Addr) -> GeoInfo {
689        GeoInfo {
690            asn: None,
691            country: None,
692            is_hosting_provider: false,
693            is_vpn_provider: false,
694        }
695    }
696}
697
698/// Diversity statistics for monitoring
699#[derive(Debug, Clone, Serialize, Deserialize)]
700pub struct DiversityStats {
701    /// Number of unique /64 subnets represented
702    pub total_64_subnets: usize,
703    /// Number of unique /48 subnets represented
704    pub total_48_subnets: usize,
705    /// Number of unique /32 subnets represented
706    pub total_32_subnets: usize,
707    /// Number of unique ASNs represented
708    pub total_asns: usize,
709    /// Number of unique countries represented
710    pub total_countries: usize,
711    /// Maximum nodes in any single /64 subnet
712    pub max_nodes_per_64: usize,
713    /// Maximum nodes in any single /48 subnet
714    pub max_nodes_per_48: usize,
715    /// Maximum nodes in any single /32 subnet
716    pub max_nodes_per_32: usize,
717}
718
719/// Reputation manager for tracking node behavior
720#[derive(Debug)]
721pub struct ReputationManager {
722    reputations: HashMap<PeerId, NodeReputation>,
723    reputation_decay: f64,
724    min_reputation: f64,
725}
726
727impl ReputationManager {
728    /// Create a new reputation manager
729    pub fn new(reputation_decay: f64, min_reputation: f64) -> Self {
730        Self {
731            reputations: HashMap::new(),
732            reputation_decay,
733            min_reputation,
734        }
735    }
736
737    /// Get reputation for a peer
738    pub fn get_reputation(&self, peer_id: &PeerId) -> Option<&NodeReputation> {
739        self.reputations.get(peer_id)
740    }
741
742    /// Update reputation based on interaction
743    pub fn update_reputation(&mut self, peer_id: &PeerId, success: bool, response_time: Duration) {
744        let reputation =
745            self.reputations
746                .entry(peer_id.clone())
747                .or_insert_with(|| NodeReputation {
748                    peer_id: peer_id.clone(),
749                    response_rate: 0.5,
750                    response_time: Duration::from_millis(500),
751                    consistency_score: 0.5,
752                    uptime_estimate: Duration::from_secs(0),
753                    routing_accuracy: 0.5,
754                    last_seen: SystemTime::now(),
755                    interaction_count: 0,
756                });
757
758        // Use higher learning rate for faster convergence in tests
759        let alpha = 0.3; // Increased from 0.1 for better test convergence
760
761        if success {
762            reputation.response_rate = reputation.response_rate * (1.0 - alpha) + alpha;
763        } else {
764            reputation.response_rate *= 1.0 - alpha;
765        }
766
767        // Update response time
768        let response_time_ms = response_time.as_millis() as f64;
769        let current_response_ms = reputation.response_time.as_millis() as f64;
770        let new_response_ms = current_response_ms * (1.0 - alpha) + response_time_ms * alpha;
771        reputation.response_time = Duration::from_millis(new_response_ms as u64);
772
773        reputation.last_seen = SystemTime::now();
774        reputation.interaction_count += 1;
775    }
776
777    /// Apply time-based reputation decay
778    pub fn apply_decay(&mut self) {
779        let now = SystemTime::now();
780
781        self.reputations.retain(|_, reputation| {
782            if let Ok(elapsed) = now.duration_since(reputation.last_seen) {
783                // Decay reputation over time
784                let decay_factor = (-elapsed.as_secs_f64() / 3600.0 * self.reputation_decay).exp();
785                reputation.response_rate *= decay_factor;
786                reputation.consistency_score *= decay_factor;
787                reputation.routing_accuracy *= decay_factor;
788
789                // Remove nodes with very low reputation
790                reputation.response_rate > self.min_reputation / 10.0
791            } else {
792                true
793            }
794        });
795    }
796}
797
798// Ed25519 compatibility removed
799
800#[cfg(test)]
801mod tests {
802    use super::*;
803    use crate::quantum_crypto::generate_ml_dsa_keypair;
804
805    fn create_test_keypair() -> (MlDsaPublicKey, MlDsaSecretKey) {
806        generate_ml_dsa_keypair().expect("Failed to generate test keypair")
807    }
808
809    fn create_test_ipv6() -> Ipv6Addr {
810        Ipv6Addr::new(
811            0x2001, 0xdb8, 0x85a3, 0x0000, 0x0000, 0x8a2e, 0x0370, 0x7334,
812        )
813    }
814
815    fn create_test_diversity_config() -> IPDiversityConfig {
816        IPDiversityConfig {
817            max_nodes_per_64: 1,
818            max_nodes_per_48: 3,
819            max_nodes_per_32: 10,
820            max_nodes_per_asn: 20,
821            enable_geolocation_check: true,
822            min_geographic_diversity: 3,
823        }
824    }
825
826    #[test]
827    fn test_ipv6_node_id_generation() -> Result<()> {
828        let (public_key, secret_key) = create_test_keypair();
829        let ipv6_addr = create_test_ipv6();
830
831        let node_id = IPv6NodeID::generate(ipv6_addr, &secret_key, &public_key)?;
832
833        assert_eq!(node_id.ipv6_addr, ipv6_addr);
834        assert_eq!(node_id.public_key.len(), 1952); // ML-DSA-65 public key size
835        assert_eq!(node_id.signature.len(), 3309); // ML-DSA-65 signature size
836        assert_eq!(node_id.node_id.len(), 32); // SHA256 output
837        assert_eq!(node_id.salt.len(), 16);
838        assert!(node_id.timestamp_secs > 0);
839
840        Ok(())
841    }
842
843    #[test]
844    fn test_ipv6_node_id_verification() -> Result<()> {
845        let (public_key, secret_key) = create_test_keypair();
846        let ipv6_addr = create_test_ipv6();
847
848        let node_id = IPv6NodeID::generate(ipv6_addr, &secret_key, &public_key)?;
849        let is_valid = node_id.verify()?;
850
851        assert!(is_valid);
852
853        Ok(())
854    }
855
856    #[test]
857    fn test_ipv6_node_id_verification_fails_with_wrong_data() -> Result<()> {
858        let (public_key, secret_key) = create_test_keypair();
859        let ipv6_addr = create_test_ipv6();
860
861        let mut node_id = IPv6NodeID::generate(ipv6_addr, &secret_key, &public_key)?;
862
863        // Tamper with the node ID
864        node_id.node_id[0] ^= 0xFF;
865        let is_valid = node_id.verify()?;
866        assert!(!is_valid);
867
868        // Test with wrong signature length
869        let mut node_id2 = IPv6NodeID::generate(ipv6_addr, &secret_key, &public_key)?;
870        node_id2.signature = vec![0u8; 32]; // Wrong length (should be 3309 for ML-DSA)
871        let is_valid2 = node_id2.verify()?;
872        assert!(!is_valid2);
873
874        // Test with wrong public key length
875        let mut node_id3 = IPv6NodeID::generate(ipv6_addr, &secret_key, &public_key)?;
876        node_id3.public_key = vec![0u8; 16]; // Wrong length (should be 1952 for ML-DSA-65)
877        let is_valid3 = node_id3.verify()?;
878        assert!(!is_valid3);
879
880        Ok(())
881    }
882
883    #[test]
884    fn test_ipv6_subnet_extraction() -> Result<()> {
885        let (public_key, secret_key) = create_test_keypair();
886        let ipv6_addr = Ipv6Addr::new(
887            0x2001, 0xdb8, 0x85a3, 0x1234, 0x5678, 0x8a2e, 0x0370, 0x7334,
888        );
889
890        let node_id = IPv6NodeID::generate(ipv6_addr, &secret_key, &public_key)?;
891
892        // Test /64 subnet extraction
893        let subnet_64 = node_id.extract_subnet_64();
894        let expected_64 = Ipv6Addr::new(0x2001, 0xdb8, 0x85a3, 0x1234, 0, 0, 0, 0);
895        assert_eq!(subnet_64, expected_64);
896
897        // Test /48 subnet extraction
898        let subnet_48 = node_id.extract_subnet_48();
899        let expected_48 = Ipv6Addr::new(0x2001, 0xdb8, 0x85a3, 0, 0, 0, 0, 0);
900        assert_eq!(subnet_48, expected_48);
901
902        // Test /32 subnet extraction
903        let subnet_32 = node_id.extract_subnet_32();
904        let expected_32 = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0);
905        assert_eq!(subnet_32, expected_32);
906
907        Ok(())
908    }
909
910    // =========== IPv4 Node ID Tests ===========
911
912    fn create_test_ipv4() -> Ipv4Addr {
913        Ipv4Addr::new(192, 168, 1, 100)
914    }
915
916    #[test]
917    fn test_ipv4_node_id_generation() -> Result<()> {
918        let (public_key, secret_key) = create_test_keypair();
919        let ipv4_addr = create_test_ipv4();
920
921        let node_id = IPv4NodeID::generate(ipv4_addr, &secret_key, &public_key)?;
922
923        assert_eq!(node_id.ipv4_addr, ipv4_addr);
924        assert_eq!(node_id.public_key.len(), 1952); // ML-DSA-65 public key size
925        assert_eq!(node_id.signature.len(), 3309); // ML-DSA-65 signature size
926        assert_eq!(node_id.node_id.len(), 32); // SHA256 output
927        assert_eq!(node_id.salt.len(), 16);
928        assert!(node_id.timestamp_secs > 0);
929
930        Ok(())
931    }
932
933    #[test]
934    fn test_ipv4_node_id_verification() -> Result<()> {
935        let (public_key, secret_key) = create_test_keypair();
936        let ipv4_addr = create_test_ipv4();
937
938        let node_id = IPv4NodeID::generate(ipv4_addr, &secret_key, &public_key)?;
939        let is_valid = node_id.verify()?;
940
941        assert!(is_valid);
942
943        Ok(())
944    }
945
946    #[test]
947    fn test_ipv4_node_id_verification_fails_with_wrong_data() -> Result<()> {
948        let (public_key, secret_key) = create_test_keypair();
949        let ipv4_addr = create_test_ipv4();
950
951        let mut node_id = IPv4NodeID::generate(ipv4_addr, &secret_key, &public_key)?;
952
953        // Tamper with the node ID
954        node_id.node_id[0] ^= 0xFF;
955        let is_valid = node_id.verify()?;
956        assert!(!is_valid);
957
958        // Test with wrong signature length
959        let mut node_id2 = IPv4NodeID::generate(ipv4_addr, &secret_key, &public_key)?;
960        node_id2.signature = vec![0u8; 32]; // Wrong length (should be 3309 for ML-DSA)
961        let is_valid2 = node_id2.verify()?;
962        assert!(!is_valid2);
963
964        // Test with wrong public key length
965        let mut node_id3 = IPv4NodeID::generate(ipv4_addr, &secret_key, &public_key)?;
966        node_id3.public_key = vec![0u8; 16]; // Wrong length (should be 1952 for ML-DSA-65)
967        let is_valid3 = node_id3.verify()?;
968        assert!(!is_valid3);
969
970        Ok(())
971    }
972
973    #[test]
974    fn test_ipv4_subnet_extraction() -> Result<()> {
975        let (public_key, secret_key) = create_test_keypair();
976        let ipv4_addr = Ipv4Addr::new(192, 168, 42, 100);
977
978        let node_id = IPv4NodeID::generate(ipv4_addr, &secret_key, &public_key)?;
979
980        // Test /24 subnet extraction
981        let subnet_24 = node_id.extract_subnet_24();
982        let expected_24 = Ipv4Addr::new(192, 168, 42, 0);
983        assert_eq!(subnet_24, expected_24);
984
985        // Test /16 subnet extraction
986        let subnet_16 = node_id.extract_subnet_16();
987        let expected_16 = Ipv4Addr::new(192, 168, 0, 0);
988        assert_eq!(subnet_16, expected_16);
989
990        // Test /8 subnet extraction
991        let subnet_8 = node_id.extract_subnet_8();
992        let expected_8 = Ipv4Addr::new(192, 0, 0, 0);
993        assert_eq!(subnet_8, expected_8);
994
995        Ok(())
996    }
997
998    #[test]
999    fn test_ipv4_node_id_age() -> Result<()> {
1000        let (public_key, secret_key) = create_test_keypair();
1001        let ipv4_addr = create_test_ipv4();
1002
1003        let node_id = IPv4NodeID::generate(ipv4_addr, &secret_key, &public_key)?;
1004
1005        // Age should be very small (just created)
1006        assert!(node_id.age_secs() < 5);
1007
1008        // Not expired with a 1 hour max age
1009        assert!(!node_id.is_expired(Duration::from_secs(3600)));
1010
1011        // A freshly created node with 0 age is NOT expired (0 > 0 is false)
1012        // This is correct behavior - a 0-second-old node is not older than 0 seconds
1013        assert!(!node_id.is_expired(Duration::from_secs(0)));
1014
1015        Ok(())
1016    }
1017
1018    #[test]
1019    fn test_ipv4_different_addresses_different_node_ids() -> Result<()> {
1020        let (public_key, secret_key) = create_test_keypair();
1021        let addr1 = Ipv4Addr::new(192, 168, 1, 1);
1022        let addr2 = Ipv4Addr::new(192, 168, 1, 2);
1023
1024        let node_id1 = IPv4NodeID::generate(addr1, &secret_key, &public_key)?;
1025        let node_id2 = IPv4NodeID::generate(addr2, &secret_key, &public_key)?;
1026
1027        // Different addresses should produce different node IDs
1028        assert_ne!(node_id1.node_id, node_id2.node_id);
1029
1030        // Both should verify successfully
1031        assert!(node_id1.verify()?);
1032        assert!(node_id2.verify()?);
1033
1034        Ok(())
1035    }
1036
1037    // =========== End IPv4 Tests ===========
1038
1039    #[test]
1040    fn test_ip_diversity_config_default() {
1041        let config = IPDiversityConfig::default();
1042
1043        assert_eq!(config.max_nodes_per_64, 1);
1044        assert_eq!(config.max_nodes_per_48, 3);
1045        assert_eq!(config.max_nodes_per_32, 10);
1046        assert_eq!(config.max_nodes_per_asn, 20);
1047        assert!(config.enable_geolocation_check);
1048        assert_eq!(config.min_geographic_diversity, 3);
1049    }
1050
1051    #[test]
1052    fn test_ip_diversity_enforcer_creation() {
1053        let config = create_test_diversity_config();
1054        let enforcer = IPDiversityEnforcer::new(config.clone());
1055
1056        assert_eq!(enforcer.config.max_nodes_per_64, config.max_nodes_per_64);
1057        assert_eq!(enforcer.subnet_64_counts.len(), 0);
1058        assert_eq!(enforcer.subnet_48_counts.len(), 0);
1059        assert_eq!(enforcer.subnet_32_counts.len(), 0);
1060    }
1061
1062    #[test]
1063    fn test_ip_analysis() -> Result<()> {
1064        let config = create_test_diversity_config();
1065        let enforcer = IPDiversityEnforcer::new(config);
1066
1067        let ipv6_addr = create_test_ipv6();
1068        let analysis = enforcer.analyze_ip(ipv6_addr)?;
1069
1070        assert_eq!(
1071            analysis.subnet_64,
1072            IPDiversityEnforcer::extract_subnet_prefix(ipv6_addr, 64)
1073        );
1074        assert_eq!(
1075            analysis.subnet_48,
1076            IPDiversityEnforcer::extract_subnet_prefix(ipv6_addr, 48)
1077        );
1078        assert_eq!(
1079            analysis.subnet_32,
1080            IPDiversityEnforcer::extract_subnet_prefix(ipv6_addr, 32)
1081        );
1082        assert!(analysis.asn.is_none()); // Not implemented in test
1083        assert!(analysis.country.is_none()); // Not implemented in test
1084        assert!(!analysis.is_hosting_provider);
1085        assert!(!analysis.is_vpn_provider);
1086        assert_eq!(analysis.reputation_score, 0.5);
1087
1088        Ok(())
1089    }
1090
1091    #[test]
1092    fn test_can_accept_node_basic() -> Result<()> {
1093        let config = create_test_diversity_config();
1094        let enforcer = IPDiversityEnforcer::new(config);
1095
1096        let ipv6_addr = create_test_ipv6();
1097        let analysis = enforcer.analyze_ip(ipv6_addr)?;
1098
1099        // Should accept first node
1100        assert!(enforcer.can_accept_node(&analysis));
1101
1102        Ok(())
1103    }
1104
1105    #[test]
1106    fn test_add_and_remove_node() -> Result<()> {
1107        let config = create_test_diversity_config();
1108        let mut enforcer = IPDiversityEnforcer::new(config);
1109
1110        let ipv6_addr = create_test_ipv6();
1111        let analysis = enforcer.analyze_ip(ipv6_addr)?;
1112
1113        // Add node
1114        enforcer.add_node(&analysis)?;
1115        assert_eq!(enforcer.subnet_64_counts.get(&analysis.subnet_64), Some(&1));
1116        assert_eq!(enforcer.subnet_48_counts.get(&analysis.subnet_48), Some(&1));
1117        assert_eq!(enforcer.subnet_32_counts.get(&analysis.subnet_32), Some(&1));
1118
1119        // Remove node
1120        enforcer.remove_node(&analysis);
1121        assert_eq!(enforcer.subnet_64_counts.get(&analysis.subnet_64), None);
1122        assert_eq!(enforcer.subnet_48_counts.get(&analysis.subnet_48), None);
1123        assert_eq!(enforcer.subnet_32_counts.get(&analysis.subnet_32), None);
1124
1125        Ok(())
1126    }
1127
1128    #[test]
1129    fn test_diversity_limits_enforcement() -> Result<()> {
1130        let config = create_test_diversity_config();
1131        let mut enforcer = IPDiversityEnforcer::new(config);
1132
1133        let ipv6_addr1 = Ipv6Addr::new(
1134            0x2001, 0xdb8, 0x85a3, 0x1234, 0x5678, 0x8a2e, 0x0370, 0x7334,
1135        );
1136        let ipv6_addr2 = Ipv6Addr::new(
1137            0x2001, 0xdb8, 0x85a3, 0x1234, 0x5678, 0x8a2e, 0x0370, 0x7335,
1138        ); // Same /64
1139
1140        let analysis1 = enforcer.analyze_ip(ipv6_addr1)?;
1141        let analysis2 = enforcer.analyze_ip(ipv6_addr2)?;
1142
1143        // First node should be accepted
1144        assert!(enforcer.can_accept_node(&analysis1));
1145        enforcer.add_node(&analysis1)?;
1146
1147        // Second node in same /64 should be rejected (max_nodes_per_64 = 1)
1148        assert!(!enforcer.can_accept_node(&analysis2));
1149
1150        // But adding should fail
1151        let result = enforcer.add_node(&analysis2);
1152        assert!(result.is_err());
1153        assert!(
1154            result
1155                .unwrap_err()
1156                .to_string()
1157                .contains("IP diversity limits exceeded")
1158        );
1159
1160        Ok(())
1161    }
1162
1163    #[test]
1164    fn test_hosting_provider_stricter_limits() -> Result<()> {
1165        let config = IPDiversityConfig {
1166            max_nodes_per_64: 4, // Set higher limit for regular nodes
1167            max_nodes_per_48: 8,
1168            ..create_test_diversity_config()
1169        };
1170        let mut enforcer = IPDiversityEnforcer::new(config);
1171
1172        let ipv6_addr = create_test_ipv6();
1173        let mut analysis = enforcer.analyze_ip(ipv6_addr)?;
1174        analysis.is_hosting_provider = true;
1175
1176        // Should accept first hosting provider node
1177        assert!(enforcer.can_accept_node(&analysis));
1178        enforcer.add_node(&analysis)?;
1179
1180        // Add second hosting provider node in same /64 (should be accepted with limit=2)
1181        let ipv6_addr2 = Ipv6Addr::new(
1182            0x2001, 0xdb8, 0x85a3, 0x0000, 0x0000, 0x8a2e, 0x0370, 0x7335,
1183        );
1184        let mut analysis2 = enforcer.analyze_ip(ipv6_addr2)?;
1185        analysis2.is_hosting_provider = true;
1186        analysis2.subnet_64 = analysis.subnet_64; // Force same subnet
1187
1188        assert!(enforcer.can_accept_node(&analysis2));
1189        enforcer.add_node(&analysis2)?;
1190
1191        // Should reject third hosting provider node in same /64 (exceeds limit=2)
1192        let ipv6_addr3 = Ipv6Addr::new(
1193            0x2001, 0xdb8, 0x85a3, 0x0000, 0x0000, 0x8a2e, 0x0370, 0x7336,
1194        );
1195        let mut analysis3 = enforcer.analyze_ip(ipv6_addr3)?;
1196        analysis3.is_hosting_provider = true;
1197        analysis3.subnet_64 = analysis.subnet_64; // Force same subnet
1198
1199        assert!(!enforcer.can_accept_node(&analysis3));
1200
1201        Ok(())
1202    }
1203
1204    #[test]
1205    fn test_diversity_stats() -> Result<()> {
1206        let config = create_test_diversity_config();
1207        let mut enforcer = IPDiversityEnforcer::new(config);
1208
1209        // Add some nodes with different subnets
1210        let addresses = [
1211            Ipv6Addr::new(
1212                0x2001, 0xdb8, 0x85a3, 0x1234, 0x5678, 0x8a2e, 0x0370, 0x7334,
1213            ),
1214            Ipv6Addr::new(
1215                0x2001, 0xdb8, 0x85a4, 0x1234, 0x5678, 0x8a2e, 0x0370, 0x7334,
1216            ), // Different /48
1217            Ipv6Addr::new(
1218                0x2002, 0xdb8, 0x85a3, 0x1234, 0x5678, 0x8a2e, 0x0370, 0x7334,
1219            ), // Different /32
1220        ];
1221
1222        for addr in addresses {
1223            let analysis = enforcer.analyze_ip(addr)?;
1224            enforcer.add_node(&analysis)?;
1225        }
1226
1227        let stats = enforcer.get_diversity_stats();
1228        assert_eq!(stats.total_64_subnets, 3);
1229        assert_eq!(stats.total_48_subnets, 3);
1230        assert_eq!(stats.total_32_subnets, 2); // Two /32 prefixes
1231        assert_eq!(stats.max_nodes_per_64, 1);
1232        assert_eq!(stats.max_nodes_per_48, 1);
1233        assert_eq!(stats.max_nodes_per_32, 2); // 2001:db8 has 2 nodes
1234
1235        Ok(())
1236    }
1237
1238    #[test]
1239    fn test_extract_subnet_prefix() {
1240        let addr = Ipv6Addr::new(
1241            0x2001, 0xdb8, 0x85a3, 0x1234, 0x5678, 0x8a2e, 0x0370, 0x7334,
1242        );
1243
1244        // Test /64 prefix
1245        let prefix_64 = IPDiversityEnforcer::extract_subnet_prefix(addr, 64);
1246        let expected_64 = Ipv6Addr::new(0x2001, 0xdb8, 0x85a3, 0x1234, 0, 0, 0, 0);
1247        assert_eq!(prefix_64, expected_64);
1248
1249        // Test /48 prefix
1250        let prefix_48 = IPDiversityEnforcer::extract_subnet_prefix(addr, 48);
1251        let expected_48 = Ipv6Addr::new(0x2001, 0xdb8, 0x85a3, 0, 0, 0, 0, 0);
1252        assert_eq!(prefix_48, expected_48);
1253
1254        // Test /32 prefix
1255        let prefix_32 = IPDiversityEnforcer::extract_subnet_prefix(addr, 32);
1256        let expected_32 = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0);
1257        assert_eq!(prefix_32, expected_32);
1258
1259        // Test /56 prefix (partial byte)
1260        let prefix_56 = IPDiversityEnforcer::extract_subnet_prefix(addr, 56);
1261        let expected_56 = Ipv6Addr::new(0x2001, 0xdb8, 0x85a3, 0x1200, 0, 0, 0, 0);
1262        assert_eq!(prefix_56, expected_56);
1263
1264        // Test /128 prefix (full address)
1265        let prefix_128 = IPDiversityEnforcer::extract_subnet_prefix(addr, 128);
1266        assert_eq!(prefix_128, addr);
1267    }
1268
1269    #[test]
1270    fn test_reputation_manager_creation() {
1271        let manager = ReputationManager::new(0.1, 0.1);
1272        assert_eq!(manager.reputation_decay, 0.1);
1273        assert_eq!(manager.min_reputation, 0.1);
1274        assert_eq!(manager.reputations.len(), 0);
1275    }
1276
1277    #[test]
1278    fn test_reputation_get_nonexistent() {
1279        let manager = ReputationManager::new(0.1, 0.1);
1280        let peer_id = "test_peer".to_string();
1281
1282        let reputation = manager.get_reputation(&peer_id);
1283        assert!(reputation.is_none());
1284    }
1285
1286    #[test]
1287    fn test_reputation_update_creates_entry() {
1288        let mut manager = ReputationManager::new(0.1, 0.1);
1289        let peer_id = "test_peer".to_string();
1290
1291        manager.update_reputation(&peer_id, true, Duration::from_millis(100));
1292
1293        let reputation = manager.get_reputation(&peer_id);
1294        assert!(reputation.is_some());
1295
1296        let rep = reputation.unwrap();
1297        assert_eq!(rep.peer_id, peer_id);
1298        assert!(rep.response_rate > 0.5); // Should increase from initial 0.5
1299        assert_eq!(rep.interaction_count, 1);
1300    }
1301
1302    #[test]
1303    fn test_reputation_update_success_improves_rate() {
1304        let mut manager = ReputationManager::new(0.1, 0.1);
1305        let peer_id = "test_peer".to_string();
1306
1307        // Multiple successful interactions
1308        for _ in 0..15 {
1309            manager.update_reputation(&peer_id, true, Duration::from_millis(100));
1310        }
1311
1312        let reputation = manager.get_reputation(&peer_id).unwrap();
1313        assert!(reputation.response_rate > 0.85); // Should be very high with higher learning rate
1314        assert_eq!(reputation.interaction_count, 15);
1315    }
1316
1317    #[test]
1318    fn test_reputation_update_failure_decreases_rate() {
1319        let mut manager = ReputationManager::new(0.1, 0.1);
1320        let peer_id = "test_peer".to_string();
1321
1322        // Multiple failed interactions
1323        for _ in 0..15 {
1324            manager.update_reputation(&peer_id, false, Duration::from_millis(1000));
1325        }
1326
1327        let reputation = manager.get_reputation(&peer_id).unwrap();
1328        assert!(reputation.response_rate < 0.15); // Should be very low with higher learning rate
1329        assert_eq!(reputation.interaction_count, 15);
1330    }
1331
1332    #[test]
1333    fn test_reputation_response_time_tracking() {
1334        let mut manager = ReputationManager::new(0.1, 0.1);
1335        let peer_id = "test_peer".to_string();
1336
1337        // Update with specific response time
1338        manager.update_reputation(&peer_id, true, Duration::from_millis(200));
1339
1340        let reputation = manager.get_reputation(&peer_id).unwrap();
1341        // Response time should be between initial 500ms and new 200ms
1342        assert!(reputation.response_time.as_millis() > 200);
1343        assert!(reputation.response_time.as_millis() < 500);
1344    }
1345
1346    #[test]
1347    fn test_reputation_decay() {
1348        let mut manager = ReputationManager::new(1.0, 0.01); // High decay rate
1349        let peer_id = "test_peer".to_string();
1350
1351        // Create a reputation entry
1352        manager.update_reputation(&peer_id, true, Duration::from_millis(100));
1353
1354        // Manually set last_seen to past
1355        if let Some(reputation) = manager.reputations.get_mut(&peer_id) {
1356            reputation.last_seen = SystemTime::now() - Duration::from_secs(7200); // 2 hours ago
1357        }
1358
1359        let original_rate = manager.get_reputation(&peer_id).unwrap().response_rate;
1360
1361        // Apply decay
1362        manager.apply_decay();
1363
1364        let reputation = manager.get_reputation(&peer_id);
1365        if let Some(rep) = reputation {
1366            // Should have decayed
1367            assert!(rep.response_rate < original_rate);
1368        } // else the reputation was removed due to low score
1369    }
1370
1371    #[test]
1372    fn test_reputation_decay_removes_low_reputation() {
1373        let mut manager = ReputationManager::new(0.1, 0.5); // High min reputation
1374        let peer_id = "test_peer".to_string();
1375
1376        // Create a low reputation entry
1377        for _ in 0..10 {
1378            manager.update_reputation(&peer_id, false, Duration::from_millis(1000));
1379        }
1380
1381        // Manually set last_seen to past
1382        if let Some(reputation) = manager.reputations.get_mut(&peer_id) {
1383            reputation.last_seen = SystemTime::now() - Duration::from_secs(3600); // 1 hour ago
1384            reputation.response_rate = 0.01; // Very low
1385        }
1386
1387        // Apply decay
1388        manager.apply_decay();
1389
1390        // Should be removed
1391        assert!(manager.get_reputation(&peer_id).is_none());
1392    }
1393
1394    #[test]
1395    fn test_security_types_keypair() {
1396        let (public_key, secret_key) =
1397            generate_ml_dsa_keypair().expect("Failed to generate keypair");
1398
1399        let public_key_bytes = public_key.as_bytes();
1400        assert_eq!(public_key_bytes.len(), 1952); // ML-DSA-65 public key size
1401
1402        let message = b"test message";
1403        let signature = ml_dsa_sign(&secret_key, message).expect("Failed to sign message");
1404        assert_eq!(signature.as_bytes().len(), 3309); // ML-DSA-65 signature size
1405
1406        // Verify the signature
1407        assert!(ml_dsa_verify(&public_key, message, &signature).is_ok());
1408    }
1409
1410    #[test]
1411    fn test_node_reputation_structure() {
1412        let peer_id = "test_peer".to_string();
1413        let reputation = NodeReputation {
1414            peer_id: peer_id.clone(),
1415            response_rate: 0.85,
1416            response_time: Duration::from_millis(150),
1417            consistency_score: 0.9,
1418            uptime_estimate: Duration::from_secs(86400),
1419            routing_accuracy: 0.8,
1420            last_seen: SystemTime::now(),
1421            interaction_count: 42,
1422        };
1423
1424        assert_eq!(reputation.peer_id, peer_id);
1425        assert_eq!(reputation.response_rate, 0.85);
1426        assert_eq!(reputation.response_time, Duration::from_millis(150));
1427        assert_eq!(reputation.consistency_score, 0.9);
1428        assert_eq!(reputation.uptime_estimate, Duration::from_secs(86400));
1429        assert_eq!(reputation.routing_accuracy, 0.8);
1430        assert_eq!(reputation.interaction_count, 42);
1431    }
1432
1433    #[test]
1434    fn test_ip_analysis_structure() {
1435        let analysis = IPAnalysis {
1436            subnet_64: Ipv6Addr::new(0x2001, 0xdb8, 0x85a3, 0x1234, 0, 0, 0, 0),
1437            subnet_48: Ipv6Addr::new(0x2001, 0xdb8, 0x85a3, 0, 0, 0, 0, 0),
1438            subnet_32: Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0),
1439            asn: Some(64512),
1440            country: Some("US".to_string()),
1441            is_hosting_provider: true,
1442            is_vpn_provider: false,
1443            reputation_score: 0.75,
1444        };
1445
1446        assert_eq!(analysis.asn, Some(64512));
1447        assert_eq!(analysis.country, Some("US".to_string()));
1448        assert!(analysis.is_hosting_provider);
1449        assert!(!analysis.is_vpn_provider);
1450        assert_eq!(analysis.reputation_score, 0.75);
1451    }
1452
1453    #[test]
1454    fn test_diversity_stats_structure() {
1455        let stats = DiversityStats {
1456            total_64_subnets: 100,
1457            total_48_subnets: 50,
1458            total_32_subnets: 25,
1459            total_asns: 15,
1460            total_countries: 8,
1461            max_nodes_per_64: 1,
1462            max_nodes_per_48: 3,
1463            max_nodes_per_32: 10,
1464        };
1465
1466        assert_eq!(stats.total_64_subnets, 100);
1467        assert_eq!(stats.total_48_subnets, 50);
1468        assert_eq!(stats.total_32_subnets, 25);
1469        assert_eq!(stats.total_asns, 15);
1470        assert_eq!(stats.total_countries, 8);
1471        assert_eq!(stats.max_nodes_per_64, 1);
1472        assert_eq!(stats.max_nodes_per_48, 3);
1473        assert_eq!(stats.max_nodes_per_32, 10);
1474    }
1475
1476    #[test]
1477    fn test_multiple_same_subnet_nodes() -> Result<()> {
1478        let config = IPDiversityConfig {
1479            max_nodes_per_64: 3, // Allow more nodes in same /64
1480            max_nodes_per_48: 5,
1481            max_nodes_per_32: 10,
1482            ..create_test_diversity_config()
1483        };
1484        let mut enforcer = IPDiversityEnforcer::new(config);
1485
1486        let _base_addr = Ipv6Addr::new(
1487            0x2001, 0xdb8, 0x85a3, 0x1234, 0x5678, 0x8a2e, 0x0370, 0x0000,
1488        );
1489
1490        // Add 3 nodes in same /64 subnet
1491        for i in 1..=3 {
1492            let addr = Ipv6Addr::new(0x2001, 0xdb8, 0x85a3, 0x1234, 0x5678, 0x8a2e, 0x0370, i);
1493            let analysis = enforcer.analyze_ip(addr)?;
1494            assert!(enforcer.can_accept_node(&analysis));
1495            enforcer.add_node(&analysis)?;
1496        }
1497
1498        // 4th node should be rejected
1499        let addr4 = Ipv6Addr::new(0x2001, 0xdb8, 0x85a3, 0x1234, 0x5678, 0x8a2e, 0x0370, 4);
1500        let analysis4 = enforcer.analyze_ip(addr4)?;
1501        assert!(!enforcer.can_accept_node(&analysis4));
1502
1503        let stats = enforcer.get_diversity_stats();
1504        assert_eq!(stats.total_64_subnets, 1);
1505        assert_eq!(stats.max_nodes_per_64, 3);
1506
1507        Ok(())
1508    }
1509
1510    #[test]
1511    fn test_asn_and_country_tracking() -> Result<()> {
1512        let config = create_test_diversity_config();
1513        let mut enforcer = IPDiversityEnforcer::new(config);
1514
1515        // Create analysis with ASN and country
1516        let ipv6_addr = create_test_ipv6();
1517        let mut analysis = enforcer.analyze_ip(ipv6_addr)?;
1518        analysis.asn = Some(64512);
1519        analysis.country = Some("US".to_string());
1520
1521        enforcer.add_node(&analysis)?;
1522
1523        assert_eq!(enforcer.asn_counts.get(&64512), Some(&1));
1524        assert_eq!(enforcer.country_counts.get("US"), Some(&1));
1525
1526        // Remove and check cleanup
1527        enforcer.remove_node(&analysis);
1528        assert!(!enforcer.asn_counts.contains_key(&64512));
1529        assert!(!enforcer.country_counts.contains_key("US"));
1530
1531        Ok(())
1532    }
1533
1534    #[test]
1535    fn test_reputation_mixed_interactions() {
1536        let mut manager = ReputationManager::new(0.1, 0.1);
1537        let peer_id = "test_peer".to_string();
1538
1539        // Mix of successful and failed interactions
1540        for i in 0..15 {
1541            let success = i % 3 != 0; // 2/3 success rate
1542            manager.update_reputation(&peer_id, success, Duration::from_millis(100 + i * 10));
1543        }
1544
1545        let reputation = manager.get_reputation(&peer_id).unwrap();
1546        // Should converge closer to 2/3 with more iterations and higher learning rate
1547        // With alpha=0.3 and 2/3 success rate, convergence may be higher
1548        assert!(reputation.response_rate > 0.55);
1549        assert!(reputation.response_rate < 0.85);
1550        assert_eq!(reputation.interaction_count, 15);
1551    }
1552}