laantungir/fips
1
0
Fork 0
mirror of https://github.com/jmcorgan/fips.git synced 2026-07-30 19:46:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

Module reorganization and clippy cleanup

Browse Source 
Move single-consumer modules into node/:
- rate_limit.rs, wire.rs, dns.rs — exclusively used by node subsystem
- Reduces top-level lib.rs from 16 to 13 modules

Split large files into focused subdirectories:
- noise.rs (1475 lines) → noise/{mod, handshake, session, replay, tests}.rs
- tree.rs (1479 lines) → tree/{mod, coordinate, declaration, state, tests}.rs
- bloom.rs (849 lines) → bloom/{mod, filter, state, tests}.rs
- All public APIs re-exported from mod.rs, no external import changes

Remove unused rate_limit defaults:
- HANDSHAKE_TIMEOUT_SECS, MAX_PENDING_INBOUND constants
- Default constructor eliminated in favor of with_params() taking config values

Fix all clippy warnings across codebase:
- Remove .clone() on Copy types, collapse nested ifs, replace match-return-None
  with ?, remove/gate unused code, fix loop indexing, remove unnecessary casts
- Box large PeerSlot enum variants to reduce size disparity
- cargo clippy --all-targets now reports zero warnings
This commit is contained in:
Johnathan Corgan
2026-02-15 15:07:42 +00:00
parent 89bc9cc4b0
commit b8a1f322c2
43 changed files with 3997 additions and 3981 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/bloom.rs
-849
View File
@@ -1,849 +0,0 @@
//! Bloom Filter Implementation
//!
//! 1KB Bloom filters for reachability in FIPS routing. Each node
//! maintains filters that summarize which destinations are reachable
//! through each peer, enabling efficient routing decisions without
//! global network knowledge.
//!
//! ## v1 Parameters
//!
//! - Size: 1 KB (8,192 bits) - sized for actual ~400-800 entry occupancy
//! - Hash functions: k=5 - optimal for 800-1,600 entries at 1KB
//! - Bandwidth: 1 KB/announce (75% reduction from original 4KB design)
//!
//! These parameters are right-sized for typical network occupancy of
//! ~250-800 entries per node.
use crate::NodeAddr;
use std::collections::{HashMap, HashSet};
use std::fmt;
use thiserror::Error;
/// Default filter size in bits (1KB = 8,192 bits).
///
/// Sized for ~800-1,600 entries with <5% FPR at typical occupancy (~400 entries).
/// This is v1 protocol default (size_class=1).
pub const DEFAULT_FILTER_SIZE_BITS: usize = 8192;
/// Default filter size in bytes (1KB).
pub const DEFAULT_FILTER_SIZE_BYTES: usize = DEFAULT_FILTER_SIZE_BITS / 8;
/// Default number of hash functions.
///
/// k=5 is optimal for 800-1,600 entries at 1KB filter size.
/// At 400 entries: FPR ~0.3%. At 800 entries: FPR ~2.4%.
pub const DEFAULT_HASH_COUNT: u8 = 5;
/// Size class for v1 protocol (1 KB filters).
pub const V1_SIZE_CLASS: u8 = 1;
/// Filter sizes by size_class: bytes = 512 << size_class
pub const SIZE_CLASS_BYTES: [usize; 4] = [512, 1024, 2048, 4096];
/// Errors related to Bloom filter operations.
#[derive(Debug, Error)]
pub enum BloomError {
#[error("invalid filter size: expected {expected} bits, got {got}")]
InvalidSize { expected: usize, got: usize },
#[error("filter size must be a multiple of 8, got {0}")]
SizeNotByteAligned(usize),
#[error("hash count must be positive")]
ZeroHashCount,
}
/// A Bloom filter for probabilistic set membership.
///
/// Used in FIPS to track which destinations are reachable through a peer.
/// The filter uses double hashing to generate k hash functions from two
/// base hashes derived from the input.
#[derive(Clone)]
pub struct BloomFilter {
/// Bit array storage (packed as bytes).
bits: Vec<u8>,
/// Number of bits in the filter.
num_bits: usize,
/// Number of hash functions to use.
hash_count: u8,
}
impl BloomFilter {
/// Create a new empty Bloom filter with default parameters.
pub fn new() -> Self {
Self::with_params(DEFAULT_FILTER_SIZE_BITS, DEFAULT_HASH_COUNT)
.expect("default params are valid")
}
/// Create a Bloom filter with custom parameters.
pub fn with_params(num_bits: usize, hash_count: u8) -> Result<Self, BloomError> {
if num_bits == 0 || !num_bits.is_multiple_of(8) {
return Err(BloomError::SizeNotByteAligned(num_bits));
}
if hash_count == 0 {
return Err(BloomError::ZeroHashCount);
}
let num_bytes = num_bits / 8;
Ok(Self {
bits: vec![0u8; num_bytes],
num_bits,
hash_count,
})
}
/// Create a Bloom filter from raw bytes.
pub fn from_bytes(bytes: Vec<u8>, hash_count: u8) -> Result<Self, BloomError> {
if hash_count == 0 {
return Err(BloomError::ZeroHashCount);
}
if bytes.is_empty() {
return Err(BloomError::SizeNotByteAligned(0));
}
let num_bits = bytes.len() * 8;
Ok(Self {
bits: bytes,
num_bits,
hash_count,
})
}
/// Create a Bloom filter from a byte slice.
pub fn from_slice(bytes: &[u8], hash_count: u8) -> Result<Self, BloomError> {
Self::from_bytes(bytes.to_vec(), hash_count)
}
/// Insert a NodeAddr into the filter.
pub fn insert(&mut self, node_addr: &NodeAddr) {
for i in 0..self.hash_count {
let bit_index = self.hash(node_addr.as_bytes(), i);
self.set_bit(bit_index);
}
}
/// Insert raw bytes into the filter.
pub fn insert_bytes(&mut self, data: &[u8]) {
for i in 0..self.hash_count {
let bit_index = self.hash(data, i);
self.set_bit(bit_index);
}
}
/// Check if the filter might contain a NodeAddr.
///
/// Returns `true` if the item might be in the set (possible false positive).
/// Returns `false` if the item is definitely not in the set.
pub fn contains(&self, node_addr: &NodeAddr) -> bool {
self.contains_bytes(node_addr.as_bytes())
}
/// Check if the filter might contain raw bytes.
pub fn contains_bytes(&self, data: &[u8]) -> bool {
for i in 0..self.hash_count {
let bit_index = self.hash(data, i);
if !self.get_bit(bit_index) {
return false;
}
}
true
}
/// Merge another filter into this one (OR operation).
///
/// After merge, this filter contains all elements from both filters.
pub fn merge(&mut self, other: &BloomFilter) -> Result<(), BloomError> {
if self.num_bits != other.num_bits {
return Err(BloomError::InvalidSize {
expected: self.num_bits,
got: other.num_bits,
});
}
for (a, b) in self.bits.iter_mut().zip(other.bits.iter()) {
*a |= b;
}
Ok(())
}
/// Create a new filter that is the union of this and another.
pub fn union(&self, other: &BloomFilter) -> Result<Self, BloomError> {
let mut result = self.clone();
result.merge(other)?;
Ok(result)
}
/// Clear all bits in the filter.
pub fn clear(&mut self) {
self.bits.fill(0);
}
/// Count the number of set bits (population count).
pub fn count_ones(&self) -> usize {
self.bits.iter().map(|b| b.count_ones() as usize).sum()
}
/// Estimate the fill ratio (set bits / total bits).
pub fn fill_ratio(&self) -> f64 {
self.count_ones() as f64 / self.num_bits as f64
}
/// Estimate the number of elements in the filter.
///
/// Uses the formula: n = -(m/k) * ln(1 - X/m)
/// where m = num_bits, k = hash_count, X = count_ones
pub fn estimated_count(&self) -> f64 {
let m = self.num_bits as f64;
let k = self.hash_count as f64;
let x = self.count_ones() as f64;
if x >= m {
return f64::INFINITY;
}
-(m / k) * (1.0 - x / m).ln()
}
/// Check if the filter is empty.
pub fn is_empty(&self) -> bool {
self.bits.iter().all(|&b| b == 0)
}
/// Get the raw bytes.
pub fn as_bytes(&self) -> &[u8] {
&self.bits
}
/// Get the filter size in bits.
pub fn num_bits(&self) -> usize {
self.num_bits
}
/// Get the filter size in bytes.
pub fn num_bytes(&self) -> usize {
self.bits.len()
}
/// Get the number of hash functions.
pub fn hash_count(&self) -> u8 {
self.hash_count
}
/// Compute a hash index for the given data and hash function number.
///
/// Uses double hashing: h(x,i) = (h1(x) + i*h2(x)) mod m
fn hash(&self, data: &[u8], k: u8) -> usize {
// Use first 16 bytes of SHA-256 for h1 and h2
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
hasher.update(data);
let hash = hasher.finalize();
// h1 from first 8 bytes
let h1 = u64::from_le_bytes(hash[0..8].try_into().unwrap());
// h2 from next 8 bytes
let h2 = u64::from_le_bytes(hash[8..16].try_into().unwrap());
let combined = h1.wrapping_add((k as u64).wrapping_mul(h2));
(combined as usize) % self.num_bits
}
fn set_bit(&mut self, index: usize) {
let byte_index = index / 8;
let bit_offset = index % 8;
self.bits[byte_index] |= 1 << bit_offset;
}
fn get_bit(&self, index: usize) -> bool {
let byte_index = index / 8;
let bit_offset = index % 8;
(self.bits[byte_index] >> bit_offset) & 1 == 1
}
}
impl Default for BloomFilter {
fn default() -> Self {
Self::new()
}
}
impl PartialEq for BloomFilter {
fn eq(&self, other: &Self) -> bool {
self.num_bits == other.num_bits
&& self.hash_count == other.hash_count
&& self.bits == other.bits
}
}
impl Eq for BloomFilter {}
impl fmt::Debug for BloomFilter {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("BloomFilter")
.field("bits", &self.num_bits)
.field("hash_count", &self.hash_count)
.field("fill_ratio", &format!("{:.2}%", self.fill_ratio() * 100.0))
.field("est_count", &format!("{:.0}", self.estimated_count()))
.finish()
}
}
/// State for managing Bloom filter announcements.
///
/// Tracks local filter state and what needs to be sent to peers.
#[derive(Clone, Debug)]
pub struct BloomState {
/// This node's NodeAddr (always included in outgoing filters).
own_node_addr: NodeAddr,
/// Leaf-only nodes we speak for (included in our filter).
leaf_dependents: HashSet<NodeAddr>,
/// Whether this node operates in leaf-only mode.
is_leaf_only: bool,
/// Rate limiting: minimum interval between outgoing updates (milliseconds).
update_debounce_ms: u64,
/// Timestamp of last update sent (per peer, in milliseconds).
last_update_sent: HashMap<NodeAddr, u64>,
/// Peers that need a filter update.
pending_updates: HashSet<NodeAddr>,
/// Current sequence number for outgoing filters.
sequence: u64,
/// Last outgoing filter sent to each peer (for change detection).
last_sent_filters: HashMap<NodeAddr, BloomFilter>,
}
impl BloomState {
/// Create new Bloom state for a node.
pub fn new(own_node_addr: NodeAddr) -> Self {
Self {
own_node_addr,
leaf_dependents: HashSet::new(),
is_leaf_only: false,
update_debounce_ms: 500,
last_update_sent: HashMap::new(),
pending_updates: HashSet::new(),
sequence: 0,
last_sent_filters: HashMap::new(),
}
}
/// Create state for a leaf-only node.
pub fn leaf_only(own_node_addr: NodeAddr) -> Self {
let mut state = Self::new(own_node_addr);
state.is_leaf_only = true;
state
}
/// Get the node's own ID.
pub fn own_node_addr(&self) -> &NodeAddr {
&self.own_node_addr
}
/// Check if this is a leaf-only node.
pub fn is_leaf_only(&self) -> bool {
self.is_leaf_only
}
/// Get the current sequence number.
pub fn sequence(&self) -> u64 {
self.sequence
}
/// Increment and return the next sequence number.
pub fn next_sequence(&mut self) -> u64 {
self.sequence += 1;
self.sequence
}
/// Get the update debounce interval in milliseconds.
pub fn update_debounce_ms(&self) -> u64 {
self.update_debounce_ms
}
/// Set the update debounce interval.
pub fn set_update_debounce_ms(&mut self, ms: u64) {
self.update_debounce_ms = ms;
}
/// Add a leaf dependent that we'll include in our filter.
pub fn add_leaf_dependent(&mut self, node_addr: NodeAddr) {
self.leaf_dependents.insert(node_addr);
}
/// Remove a leaf dependent.
pub fn remove_leaf_dependent(&mut self, node_addr: &NodeAddr) -> bool {
self.leaf_dependents.remove(node_addr)
}
/// Get the set of leaf dependents.
pub fn leaf_dependents(&self) -> &HashSet<NodeAddr> {
&self.leaf_dependents
}
/// Number of leaf dependents.
pub fn leaf_dependent_count(&self) -> usize {
self.leaf_dependents.len()
}
/// Mark that a peer needs an update.
pub fn mark_update_needed(&mut self, peer_id: NodeAddr) {
self.pending_updates.insert(peer_id);
}
/// Mark all peers as needing updates.
pub fn mark_all_updates_needed(&mut self, peer_ids: impl IntoIterator<Item = NodeAddr>) {
self.pending_updates.extend(peer_ids);
}
/// Check if a peer needs an update.
pub fn needs_update(&self, peer_id: &NodeAddr) -> bool {
self.pending_updates.contains(peer_id)
}
/// Check if we should send an update to a peer (respecting debounce).
pub fn should_send_update(&self, peer_id: &NodeAddr, current_time_ms: u64) -> bool {
if !self.pending_updates.contains(peer_id) {
return false;
}
match self.last_update_sent.get(peer_id) {
Some(&last_time) => current_time_ms >= last_time + self.update_debounce_ms,
None => true,
}
}
/// Record that we sent an update to a peer.
pub fn record_update_sent(&mut self, peer_id: NodeAddr, current_time_ms: u64) {
self.last_update_sent.insert(peer_id, current_time_ms);
self.pending_updates.remove(&peer_id);
}
/// Clear all pending updates.
pub fn clear_pending_updates(&mut self) {
self.pending_updates.clear();
}
/// Record the outgoing filter that was sent to a peer.
pub fn record_sent_filter(&mut self, peer_id: NodeAddr, filter: BloomFilter) {
self.last_sent_filters.insert(peer_id, filter);
}
/// Remove stored filter state for a peer that was removed.
pub fn remove_peer_state(&mut self, peer_id: &NodeAddr) {
self.last_sent_filters.remove(peer_id);
self.last_update_sent.remove(peer_id);
self.pending_updates.remove(peer_id);
}
/// Mark only peers whose outgoing filter has actually changed.
///
/// Computes the outgoing filter for each peer and compares it
/// against what was last sent. Only marks peers where the filter
/// differs. This prevents cascading update loops in steady state.
pub fn mark_changed_peers(
&mut self,
exclude_from: &NodeAddr,
peer_addrs: &[NodeAddr],
peer_filters: &HashMap<NodeAddr, BloomFilter>,
) {
for peer_addr in peer_addrs {
if peer_addr == exclude_from {
continue;
}
let new_filter = self.compute_outgoing_filter(peer_addr, peer_filters);
let changed = match self.last_sent_filters.get(peer_addr) {
Some(last) => *last != new_filter,
None => true, // never sent → must send
};
if changed {
self.pending_updates.insert(*peer_addr);
}
}
}
/// Compute the outgoing filter for a specific peer.
///
/// The filter includes:
/// - This node's own ID
/// - All leaf dependents
/// - Entries from other peers' inbound filters (excluding the destination peer)
///
/// The `peer_filters` map contains inbound filters from each peer.
/// The filter for `exclude_peer` is excluded to prevent routing loops.
pub fn compute_outgoing_filter(
&self,
exclude_peer: &NodeAddr,
peer_filters: &HashMap<NodeAddr, BloomFilter>,
) -> BloomFilter {
let mut filter = BloomFilter::new();
// Always include ourselves
filter.insert(&self.own_node_addr);
// Include leaf dependents
for dep in &self.leaf_dependents {
filter.insert(dep);
}
// Merge filters from other peers
for (peer_id, peer_filter) in peer_filters {
if peer_id != exclude_peer {
// Ignore merge errors (size mismatches) - just skip that filter
let _ = filter.merge(peer_filter);
}
}
filter
}
/// Create a base filter containing just this node and its dependents.
pub fn base_filter(&self) -> BloomFilter {
let mut filter = BloomFilter::new();
filter.insert(&self.own_node_addr);
for dep in &self.leaf_dependents {
filter.insert(dep);
}
filter
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_node_addr(val: u8) -> NodeAddr {
let mut bytes = [0u8; 16];
bytes[0] = val;
NodeAddr::from_bytes(bytes)
}
// ===== BloomFilter Tests =====
#[test]
fn test_bloom_filter_new() {
let filter = BloomFilter::new();
assert_eq!(filter.num_bits(), DEFAULT_FILTER_SIZE_BITS);
assert_eq!(filter.hash_count(), DEFAULT_HASH_COUNT);
assert_eq!(filter.count_ones(), 0);
assert!(filter.is_empty());
}
#[test]
fn test_bloom_filter_insert_contains() {
let mut filter = BloomFilter::new();
let node1 = make_node_addr(1);
let node2 = make_node_addr(2);
assert!(!filter.contains(&node1));
assert!(!filter.contains(&node2));
filter.insert(&node1);
assert!(filter.contains(&node1));
// node2 might have false positive, but very unlikely with single insert
assert!(!filter.is_empty());
}
#[test]
fn test_bloom_filter_multiple_inserts() {
let mut filter = BloomFilter::new();
for i in 0..100 {
let node = make_node_addr(i);
filter.insert(&node);
}
// All inserted items should be found
for i in 0..100 {
let node = make_node_addr(i);
assert!(filter.contains(&node), "Node {} not found", i);
}
// Fill ratio should be reasonable
let fill = filter.fill_ratio();
assert!(fill > 0.0 && fill < 0.5, "Unexpected fill ratio: {}", fill);
}
#[test]
fn test_bloom_filter_merge() {
let mut filter1 = BloomFilter::new();
let mut filter2 = BloomFilter::new();
let node1 = make_node_addr(1);
let node2 = make_node_addr(2);
filter1.insert(&node1);
filter2.insert(&node2);
filter1.merge(&filter2).unwrap();
assert!(filter1.contains(&node1));
assert!(filter1.contains(&node2));
}
#[test]
fn test_bloom_filter_union() {
let mut filter1 = BloomFilter::new();
let mut filter2 = BloomFilter::new();
let node1 = make_node_addr(1);
let node2 = make_node_addr(2);
filter1.insert(&node1);
filter2.insert(&node2);
let union = filter1.union(&filter2).unwrap();
assert!(union.contains(&node1));
assert!(union.contains(&node2));
// Original filters unchanged
assert!(!filter1.contains(&node2));
assert!(!filter2.contains(&node1));
}
#[test]
fn test_bloom_filter_clear() {
let mut filter = BloomFilter::new();
let node = make_node_addr(1);
filter.insert(&node);
assert!(!filter.is_empty());
filter.clear();
assert!(filter.is_empty());
assert_eq!(filter.count_ones(), 0);
assert!(!filter.contains(&node));
}
#[test]
fn test_bloom_filter_merge_size_mismatch() {
let mut filter1 = BloomFilter::with_params(1024, 7).unwrap();
let filter2 = BloomFilter::with_params(2048, 7).unwrap();
let result = filter1.merge(&filter2);
assert!(matches!(result, Err(BloomError::InvalidSize { .. })));
}
#[test]
fn test_bloom_filter_custom_params() {
let filter = BloomFilter::with_params(1024, 5).unwrap();
assert_eq!(filter.num_bits(), 1024);
assert_eq!(filter.num_bytes(), 128);
assert_eq!(filter.hash_count(), 5);
}
#[test]
fn test_bloom_filter_invalid_params() {
// Not byte-aligned (1001 is not divisible by 8)
assert!(matches!(
BloomFilter::with_params(1001, 7),
Err(BloomError::SizeNotByteAligned(1001))
));
// Zero size
assert!(matches!(
BloomFilter::with_params(0, 7),
Err(BloomError::SizeNotByteAligned(0))
));
// Zero hash count
assert!(matches!(
BloomFilter::with_params(1024, 0),
Err(BloomError::ZeroHashCount)
));
}
#[test]
fn test_bloom_filter_from_bytes() {
let original = BloomFilter::new();
let bytes = original.as_bytes().to_vec();
let restored =
BloomFilter::from_bytes(bytes, original.hash_count()).unwrap();
assert_eq!(original, restored);
}
#[test]
fn test_bloom_filter_estimated_count() {
let mut filter = BloomFilter::new();
// Empty filter
assert_eq!(filter.estimated_count(), 0.0);
// Insert some items
for i in 0..50 {
filter.insert(&make_node_addr(i));
}
// Estimate should be reasonably close to 50
let estimate = filter.estimated_count();
assert!(
estimate > 30.0 && estimate < 100.0,
"Unexpected estimate: {}",
estimate
);
}
#[test]
fn test_bloom_filter_equality() {
let mut filter1 = BloomFilter::new();
let mut filter2 = BloomFilter::new();
assert_eq!(filter1, filter2);
filter1.insert(&make_node_addr(1));
assert_ne!(filter1, filter2);
filter2.insert(&make_node_addr(1));
assert_eq!(filter1, filter2);
}
// ===== BloomState Tests =====
#[test]
fn test_bloom_state_new() {
let node = make_node_addr(0);
let state = BloomState::new(node);
assert_eq!(state.own_node_addr(), &node);
assert!(!state.is_leaf_only());
assert_eq!(state.sequence(), 0);
assert_eq!(state.leaf_dependent_count(), 0);
}
#[test]
fn test_bloom_state_leaf_only() {
let node = make_node_addr(0);
let state = BloomState::leaf_only(node);
assert!(state.is_leaf_only());
}
#[test]
fn test_bloom_state_leaf_dependents() {
let node = make_node_addr(0);
let mut state = BloomState::new(node);
let leaf1 = make_node_addr(1);
let leaf2 = make_node_addr(2);
state.add_leaf_dependent(leaf1);
state.add_leaf_dependent(leaf2);
assert_eq!(state.leaf_dependent_count(), 2);
assert!(state.remove_leaf_dependent(&leaf1));
assert_eq!(state.leaf_dependent_count(), 1);
assert!(!state.remove_leaf_dependent(&leaf1)); // already removed
}
#[test]
fn test_bloom_state_debounce() {
let node = make_node_addr(0);
let peer = make_node_addr(1);
let mut state = BloomState::new(node);
state.set_update_debounce_ms(500);
state.mark_update_needed(peer);
// Should send initially
assert!(state.should_send_update(&peer, 1000));
// Record send
state.record_update_sent(peer, 1000);
state.mark_update_needed(peer);
// Should not send immediately (within debounce)
assert!(!state.should_send_update(&peer, 1200));
// Should send after debounce period
assert!(state.should_send_update(&peer, 1600));
}
#[test]
fn test_bloom_state_sequence() {
let node = make_node_addr(0);
let mut state = BloomState::new(node);
assert_eq!(state.sequence(), 0);
assert_eq!(state.next_sequence(), 1);
assert_eq!(state.next_sequence(), 2);
assert_eq!(state.sequence(), 2);
}
#[test]
fn test_bloom_state_pending_updates() {
let node = make_node_addr(0);
let mut state = BloomState::new(node);
let peer1 = make_node_addr(1);
let peer2 = make_node_addr(2);
assert!(!state.needs_update(&peer1));
state.mark_update_needed(peer1);
assert!(state.needs_update(&peer1));
assert!(!state.needs_update(&peer2));
state.mark_all_updates_needed(vec![peer1, peer2]);
assert!(state.needs_update(&peer1));
assert!(state.needs_update(&peer2));
state.clear_pending_updates();
assert!(!state.needs_update(&peer1));
assert!(!state.needs_update(&peer2));
}
#[test]
fn test_bloom_state_base_filter() {
let node = make_node_addr(0);
let mut state = BloomState::new(node);
let leaf = make_node_addr(1);
state.add_leaf_dependent(leaf);
let filter = state.base_filter();
assert!(filter.contains(&node));
assert!(filter.contains(&leaf));
assert!(!filter.contains(&make_node_addr(99)));
}
#[test]
fn test_bloom_state_compute_outgoing_filter() {
let my_node = make_node_addr(0);
let mut state = BloomState::new(my_node);
let leaf = make_node_addr(1);
state.add_leaf_dependent(leaf);
let peer1 = make_node_addr(10);
let peer2 = make_node_addr(20);
// Create peer filters
let mut filter1 = BloomFilter::new();
filter1.insert(&make_node_addr(100));
filter1.insert(&make_node_addr(101));
let mut filter2 = BloomFilter::new();
filter2.insert(&make_node_addr(200));
let mut peer_filters = HashMap::new();
peer_filters.insert(peer1, filter1);
peer_filters.insert(peer2, filter2);
// Filter for peer1 should exclude peer1's contributions
let outgoing1 = state.compute_outgoing_filter(&peer1, &peer_filters);
assert!(outgoing1.contains(&my_node)); // self
assert!(outgoing1.contains(&leaf)); // leaf dependent
assert!(outgoing1.contains(&make_node_addr(200))); // from peer2
// peer1's nodes may or may not be present (depends on split brain)
// Filter for peer2 should exclude peer2's contributions
let outgoing2 = state.compute_outgoing_filter(&peer2, &peer_filters);
assert!(outgoing2.contains(&my_node));
assert!(outgoing2.contains(&leaf));
assert!(outgoing2.contains(&make_node_addr(100))); // from peer1
assert!(outgoing2.contains(&make_node_addr(101))); // from peer1
}
}
src/bloom/filter.rs
+240
View File
@@ -0,0 +1,240 @@
//! Generic Bloom filter data structure.
use std::fmt;
use super::{BloomError, DEFAULT_FILTER_SIZE_BITS, DEFAULT_HASH_COUNT};
use crate::NodeAddr;
/// A Bloom filter for probabilistic set membership.
///
/// Used in FIPS to track which destinations are reachable through a peer.
/// The filter uses double hashing to generate k hash functions from two
/// base hashes derived from the input.
#[derive(Clone)]
pub struct BloomFilter {
/// Bit array storage (packed as bytes).
bits: Vec<u8>,
/// Number of bits in the filter.
num_bits: usize,
/// Number of hash functions to use.
hash_count: u8,
}
impl BloomFilter {
/// Create a new empty Bloom filter with default parameters.
pub fn new() -> Self {
Self::with_params(DEFAULT_FILTER_SIZE_BITS, DEFAULT_HASH_COUNT)
.expect("default params are valid")
}
/// Create a Bloom filter with custom parameters.
pub fn with_params(num_bits: usize, hash_count: u8) -> Result<Self, BloomError> {
if num_bits == 0 || !num_bits.is_multiple_of(8) {
return Err(BloomError::SizeNotByteAligned(num_bits));
}
if hash_count == 0 {
return Err(BloomError::ZeroHashCount);
}
let num_bytes = num_bits / 8;
Ok(Self {
bits: vec![0u8; num_bytes],
num_bits,
hash_count,
})
}
/// Create a Bloom filter from raw bytes.
pub fn from_bytes(bytes: Vec<u8>, hash_count: u8) -> Result<Self, BloomError> {
if hash_count == 0 {
return Err(BloomError::ZeroHashCount);
}
if bytes.is_empty() {
return Err(BloomError::SizeNotByteAligned(0));
}
let num_bits = bytes.len() * 8;
Ok(Self {
bits: bytes,
num_bits,
hash_count,
})
}
/// Create a Bloom filter from a byte slice.
pub fn from_slice(bytes: &[u8], hash_count: u8) -> Result<Self, BloomError> {
Self::from_bytes(bytes.to_vec(), hash_count)
}
/// Insert a NodeAddr into the filter.
pub fn insert(&mut self, node_addr: &NodeAddr) {
for i in 0..self.hash_count {
let bit_index = self.hash(node_addr.as_bytes(), i);
self.set_bit(bit_index);
}
}
/// Insert raw bytes into the filter.
pub fn insert_bytes(&mut self, data: &[u8]) {
for i in 0..self.hash_count {
let bit_index = self.hash(data, i);
self.set_bit(bit_index);
}
}
/// Check if the filter might contain a NodeAddr.
///
/// Returns `true` if the item might be in the set (possible false positive).
/// Returns `false` if the item is definitely not in the set.
pub fn contains(&self, node_addr: &NodeAddr) -> bool {
self.contains_bytes(node_addr.as_bytes())
}
/// Check if the filter might contain raw bytes.
pub fn contains_bytes(&self, data: &[u8]) -> bool {
for i in 0..self.hash_count {
let bit_index = self.hash(data, i);
if !self.get_bit(bit_index) {
return false;
}
}
true
}
/// Merge another filter into this one (OR operation).
///
/// After merge, this filter contains all elements from both filters.
pub fn merge(&mut self, other: &BloomFilter) -> Result<(), BloomError> {
if self.num_bits != other.num_bits {
return Err(BloomError::InvalidSize {
expected: self.num_bits,
got: other.num_bits,
});
}
for (a, b) in self.bits.iter_mut().zip(other.bits.iter()) {
*a |= b;
}
Ok(())
}
/// Create a new filter that is the union of this and another.
pub fn union(&self, other: &BloomFilter) -> Result<Self, BloomError> {
let mut result = self.clone();
result.merge(other)?;
Ok(result)
}
/// Clear all bits in the filter.
pub fn clear(&mut self) {
self.bits.fill(0);
}
/// Count the number of set bits (population count).
pub fn count_ones(&self) -> usize {
self.bits.iter().map(|b| b.count_ones() as usize).sum()
}
/// Estimate the fill ratio (set bits / total bits).
pub fn fill_ratio(&self) -> f64 {
self.count_ones() as f64 / self.num_bits as f64
}
/// Estimate the number of elements in the filter.
///
/// Uses the formula: n = -(m/k) * ln(1 - X/m)
/// where m = num_bits, k = hash_count, X = count_ones
pub fn estimated_count(&self) -> f64 {
let m = self.num_bits as f64;
let k = self.hash_count as f64;
let x = self.count_ones() as f64;
if x >= m {
return f64::INFINITY;
}
-(m / k) * (1.0 - x / m).ln()
}
/// Check if the filter is empty.
pub fn is_empty(&self) -> bool {
self.bits.iter().all(|&b| b == 0)
}
/// Get the raw bytes.
pub fn as_bytes(&self) -> &[u8] {
&self.bits
}
/// Get the filter size in bits.
pub fn num_bits(&self) -> usize {
self.num_bits
}
/// Get the filter size in bytes.
pub fn num_bytes(&self) -> usize {
self.bits.len()
}
/// Get the number of hash functions.
pub fn hash_count(&self) -> u8 {
self.hash_count
}
/// Compute a hash index for the given data and hash function number.
///
/// Uses double hashing: h(x,i) = (h1(x) + i*h2(x)) mod m
fn hash(&self, data: &[u8], k: u8) -> usize {
// Use first 16 bytes of SHA-256 for h1 and h2
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
hasher.update(data);
let hash = hasher.finalize();
// h1 from first 8 bytes
let h1 = u64::from_le_bytes(hash[0..8].try_into().unwrap());
// h2 from next 8 bytes
let h2 = u64::from_le_bytes(hash[8..16].try_into().unwrap());
let combined = h1.wrapping_add((k as u64).wrapping_mul(h2));
(combined as usize) % self.num_bits
}
fn set_bit(&mut self, index: usize) {
let byte_index = index / 8;
let bit_offset = index % 8;
self.bits[byte_index] |= 1 << bit_offset;
}
fn get_bit(&self, index: usize) -> bool {
let byte_index = index / 8;
let bit_offset = index % 8;
(self.bits[byte_index] >> bit_offset) & 1 == 1
}
}
impl Default for BloomFilter {
fn default() -> Self {
Self::new()
}
}
impl PartialEq for BloomFilter {
fn eq(&self, other: &Self) -> bool {
self.num_bits == other.num_bits
&& self.hash_count == other.hash_count
&& self.bits == other.bits
}
}
impl Eq for BloomFilter {}
impl fmt::Debug for BloomFilter {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("BloomFilter")
.field("bits", &self.num_bits)
.field("hash_count", &self.hash_count)
.field("fill_ratio", &format!("{:.2}%", self.fill_ratio() * 100.0))
.field("est_count", &format!("{:.0}", self.estimated_count()))
.finish()
}
}
src/bloom/mod.rs
+60
View File
@@ -0,0 +1,60 @@
//! Bloom Filter Implementation
//!
//! 1KB Bloom filters for reachability in FIPS routing. Each node
//! maintains filters that summarize which destinations are reachable
//! through each peer, enabling efficient routing decisions without
//! global network knowledge.
//!
//! ## v1 Parameters
//!
//! - Size: 1 KB (8,192 bits) - sized for actual ~400-800 entry occupancy
//! - Hash functions: k=5 - optimal for 800-1,600 entries at 1KB
//! - Bandwidth: 1 KB/announce (75% reduction from original 4KB design)
//!
//! These parameters are right-sized for typical network occupancy of
//! ~250-800 entries per node.
mod filter;
mod state;
use thiserror::Error;
pub use filter::BloomFilter;
pub use state::BloomState;
/// Default filter size in bits (1KB = 8,192 bits).
///
/// Sized for ~800-1,600 entries with <5% FPR at typical occupancy (~400 entries).
/// This is v1 protocol default (size_class=1).
pub const DEFAULT_FILTER_SIZE_BITS: usize = 8192;
/// Default filter size in bytes (1KB).
pub const DEFAULT_FILTER_SIZE_BYTES: usize = DEFAULT_FILTER_SIZE_BITS / 8;
/// Default number of hash functions.
///
/// k=5 is optimal for 800-1,600 entries at 1KB filter size.
/// At 400 entries: FPR ~0.3%. At 800 entries: FPR ~2.4%.
pub const DEFAULT_HASH_COUNT: u8 = 5;
/// Size class for v1 protocol (1 KB filters).
pub const V1_SIZE_CLASS: u8 = 1;
/// Filter sizes by size_class: bytes = 512 << size_class
pub const SIZE_CLASS_BYTES: [usize; 4] = [512, 1024, 2048, 4096];
/// Errors related to Bloom filter operations.
#[derive(Debug, Error)]
pub enum BloomError {
#[error("invalid filter size: expected {expected} bits, got {got}")]
InvalidSize { expected: usize, got: usize },
#[error("filter size must be a multiple of 8, got {0}")]
SizeNotByteAligned(usize),
#[error("hash count must be positive")]
ZeroHashCount,
}
#[cfg(test)]
mod tests;
src/bloom/state.rs
+224
View File
@@ -0,0 +1,224 @@
//! FIPS-specific Bloom filter announcement state management.
use std::collections::{HashMap, HashSet};
use super::BloomFilter;
use crate::NodeAddr;
/// State for managing Bloom filter announcements.
///
/// Tracks local filter state and what needs to be sent to peers.
#[derive(Clone, Debug)]
pub struct BloomState {
/// This node's NodeAddr (always included in outgoing filters).
own_node_addr: NodeAddr,
/// Leaf-only nodes we speak for (included in our filter).
leaf_dependents: HashSet<NodeAddr>,
/// Whether this node operates in leaf-only mode.
is_leaf_only: bool,
/// Rate limiting: minimum interval between outgoing updates (milliseconds).
update_debounce_ms: u64,
/// Timestamp of last update sent (per peer, in milliseconds).
last_update_sent: HashMap<NodeAddr, u64>,
/// Peers that need a filter update.
pending_updates: HashSet<NodeAddr>,
/// Current sequence number for outgoing filters.
sequence: u64,
/// Last outgoing filter sent to each peer (for change detection).
last_sent_filters: HashMap<NodeAddr, BloomFilter>,
}
impl BloomState {
/// Create new Bloom state for a node.
pub fn new(own_node_addr: NodeAddr) -> Self {
Self {
own_node_addr,
leaf_dependents: HashSet::new(),
is_leaf_only: false,
update_debounce_ms: 500,
last_update_sent: HashMap::new(),
pending_updates: HashSet::new(),
sequence: 0,
last_sent_filters: HashMap::new(),
}
}
/// Create state for a leaf-only node.
pub fn leaf_only(own_node_addr: NodeAddr) -> Self {
let mut state = Self::new(own_node_addr);
state.is_leaf_only = true;
state
}
/// Get the node's own ID.
pub fn own_node_addr(&self) -> &NodeAddr {
&self.own_node_addr
}
/// Check if this is a leaf-only node.
pub fn is_leaf_only(&self) -> bool {
self.is_leaf_only
}
/// Get the current sequence number.
pub fn sequence(&self) -> u64 {
self.sequence
}
/// Increment and return the next sequence number.
pub fn next_sequence(&mut self) -> u64 {
self.sequence += 1;
self.sequence
}
/// Get the update debounce interval in milliseconds.
pub fn update_debounce_ms(&self) -> u64 {
self.update_debounce_ms
}
/// Set the update debounce interval.
pub fn set_update_debounce_ms(&mut self, ms: u64) {
self.update_debounce_ms = ms;
}
/// Add a leaf dependent that we'll include in our filter.
pub fn add_leaf_dependent(&mut self, node_addr: NodeAddr) {
self.leaf_dependents.insert(node_addr);
}
/// Remove a leaf dependent.
pub fn remove_leaf_dependent(&mut self, node_addr: &NodeAddr) -> bool {
self.leaf_dependents.remove(node_addr)
}
/// Get the set of leaf dependents.
pub fn leaf_dependents(&self) -> &HashSet<NodeAddr> {
&self.leaf_dependents
}
/// Number of leaf dependents.
pub fn leaf_dependent_count(&self) -> usize {
self.leaf_dependents.len()
}
/// Mark that a peer needs an update.
pub fn mark_update_needed(&mut self, peer_id: NodeAddr) {
self.pending_updates.insert(peer_id);
}
/// Mark all peers as needing updates.
pub fn mark_all_updates_needed(&mut self, peer_ids: impl IntoIterator<Item = NodeAddr>) {
self.pending_updates.extend(peer_ids);
}
/// Check if a peer needs an update.
pub fn needs_update(&self, peer_id: &NodeAddr) -> bool {
self.pending_updates.contains(peer_id)
}
/// Check if we should send an update to a peer (respecting debounce).
pub fn should_send_update(&self, peer_id: &NodeAddr, current_time_ms: u64) -> bool {
if !self.pending_updates.contains(peer_id) {
return false;
}
match self.last_update_sent.get(peer_id) {
Some(&last_time) => current_time_ms >= last_time + self.update_debounce_ms,
None => true,
}
}
/// Record that we sent an update to a peer.
pub fn record_update_sent(&mut self, peer_id: NodeAddr, current_time_ms: u64) {
self.last_update_sent.insert(peer_id, current_time_ms);
self.pending_updates.remove(&peer_id);
}
/// Clear all pending updates.
pub fn clear_pending_updates(&mut self) {
self.pending_updates.clear();
}
/// Record the outgoing filter that was sent to a peer.
pub fn record_sent_filter(&mut self, peer_id: NodeAddr, filter: BloomFilter) {
self.last_sent_filters.insert(peer_id, filter);
}
/// Remove stored filter state for a peer that was removed.
pub fn remove_peer_state(&mut self, peer_id: &NodeAddr) {
self.last_sent_filters.remove(peer_id);
self.last_update_sent.remove(peer_id);
self.pending_updates.remove(peer_id);
}
/// Mark only peers whose outgoing filter has actually changed.
///
/// Computes the outgoing filter for each peer and compares it
/// against what was last sent. Only marks peers where the filter
/// differs. This prevents cascading update loops in steady state.
pub fn mark_changed_peers(
&mut self,
exclude_from: &NodeAddr,
peer_addrs: &[NodeAddr],
peer_filters: &HashMap<NodeAddr, BloomFilter>,
) {
for peer_addr in peer_addrs {
if peer_addr == exclude_from {
continue;
}
let new_filter = self.compute_outgoing_filter(peer_addr, peer_filters);
let changed = match self.last_sent_filters.get(peer_addr) {
Some(last) => *last != new_filter,
None => true, // never sent → must send
};
if changed {
self.pending_updates.insert(*peer_addr);
}
}
}
/// Compute the outgoing filter for a specific peer.
///
/// The filter includes:
/// - This node's own ID
/// - All leaf dependents
/// - Entries from other peers' inbound filters (excluding the destination peer)
///
/// The `peer_filters` map contains inbound filters from each peer.
/// The filter for `exclude_peer` is excluded to prevent routing loops.
pub fn compute_outgoing_filter(
&self,
exclude_peer: &NodeAddr,
peer_filters: &HashMap<NodeAddr, BloomFilter>,
) -> BloomFilter {
let mut filter = BloomFilter::new();
// Always include ourselves
filter.insert(&self.own_node_addr);
// Include leaf dependents
for dep in &self.leaf_dependents {
filter.insert(dep);
}
// Merge filters from other peers
for (peer_id, peer_filter) in peer_filters {
if peer_id != exclude_peer {
// Ignore merge errors (size mismatches) - just skip that filter
let _ = filter.merge(peer_filter);
}
}
filter
}
/// Create a base filter containing just this node and its dependents.
pub fn base_filter(&self) -> BloomFilter {
let mut filter = BloomFilter::new();
filter.insert(&self.own_node_addr);
for dep in &self.leaf_dependents {
filter.insert(dep);
}
filter
}
}
src/bloom/tests.rs
+340
View File
@@ -0,0 +1,340 @@
use super::*;
use crate::NodeAddr;
use std::collections::HashMap;
fn make_node_addr(val: u8) -> NodeAddr {
let mut bytes = [0u8; 16];
bytes[0] = val;
NodeAddr::from_bytes(bytes)
}
// ===== BloomFilter Tests =====
#[test]
fn test_bloom_filter_new() {
let filter = BloomFilter::new();
assert_eq!(filter.num_bits(), DEFAULT_FILTER_SIZE_BITS);
assert_eq!(filter.hash_count(), DEFAULT_HASH_COUNT);
assert_eq!(filter.count_ones(), 0);
assert!(filter.is_empty());
}
#[test]
fn test_bloom_filter_insert_contains() {
let mut filter = BloomFilter::new();
let node1 = make_node_addr(1);
let node2 = make_node_addr(2);
assert!(!filter.contains(&node1));
assert!(!filter.contains(&node2));
filter.insert(&node1);
assert!(filter.contains(&node1));
// node2 might have false positive, but very unlikely with single insert
assert!(!filter.is_empty());
}
#[test]
fn test_bloom_filter_multiple_inserts() {
let mut filter = BloomFilter::new();
for i in 0..100 {
let node = make_node_addr(i);
filter.insert(&node);
}
// All inserted items should be found
for i in 0..100 {
let node = make_node_addr(i);
assert!(filter.contains(&node), "Node {} not found", i);
}
// Fill ratio should be reasonable
let fill = filter.fill_ratio();
assert!(fill > 0.0 && fill < 0.5, "Unexpected fill ratio: {}", fill);
}
#[test]
fn test_bloom_filter_merge() {
let mut filter1 = BloomFilter::new();
let mut filter2 = BloomFilter::new();
let node1 = make_node_addr(1);
let node2 = make_node_addr(2);
filter1.insert(&node1);
filter2.insert(&node2);
filter1.merge(&filter2).unwrap();
assert!(filter1.contains(&node1));
assert!(filter1.contains(&node2));
}
#[test]
fn test_bloom_filter_union() {
let mut filter1 = BloomFilter::new();
let mut filter2 = BloomFilter::new();
let node1 = make_node_addr(1);
let node2 = make_node_addr(2);
filter1.insert(&node1);
filter2.insert(&node2);
let union = filter1.union(&filter2).unwrap();
assert!(union.contains(&node1));
assert!(union.contains(&node2));
// Original filters unchanged
assert!(!filter1.contains(&node2));
assert!(!filter2.contains(&node1));
}
#[test]
fn test_bloom_filter_clear() {
let mut filter = BloomFilter::new();
let node = make_node_addr(1);
filter.insert(&node);
assert!(!filter.is_empty());
filter.clear();
assert!(filter.is_empty());
assert_eq!(filter.count_ones(), 0);
assert!(!filter.contains(&node));
}
#[test]
fn test_bloom_filter_merge_size_mismatch() {
let mut filter1 = BloomFilter::with_params(1024, 7).unwrap();
let filter2 = BloomFilter::with_params(2048, 7).unwrap();
let result = filter1.merge(&filter2);
assert!(matches!(result, Err(BloomError::InvalidSize { .. })));
}
#[test]
fn test_bloom_filter_custom_params() {
let filter = BloomFilter::with_params(1024, 5).unwrap();
assert_eq!(filter.num_bits(), 1024);
assert_eq!(filter.num_bytes(), 128);
assert_eq!(filter.hash_count(), 5);
}
#[test]
fn test_bloom_filter_invalid_params() {
// Not byte-aligned (1001 is not divisible by 8)
assert!(matches!(
BloomFilter::with_params(1001, 7),
Err(BloomError::SizeNotByteAligned(1001))
));
// Zero size
assert!(matches!(
BloomFilter::with_params(0, 7),
Err(BloomError::SizeNotByteAligned(0))
));
// Zero hash count
assert!(matches!(
BloomFilter::with_params(1024, 0),
Err(BloomError::ZeroHashCount)
));
}
#[test]
fn test_bloom_filter_from_bytes() {
let original = BloomFilter::new();
let bytes = original.as_bytes().to_vec();
let restored =
BloomFilter::from_bytes(bytes, original.hash_count()).unwrap();
assert_eq!(original, restored);
}
#[test]
fn test_bloom_filter_estimated_count() {
let mut filter = BloomFilter::new();
// Empty filter
assert_eq!(filter.estimated_count(), 0.0);
// Insert some items
for i in 0..50 {
filter.insert(&make_node_addr(i));
}
// Estimate should be reasonably close to 50
let estimate = filter.estimated_count();
assert!(
estimate > 30.0 && estimate < 100.0,
"Unexpected estimate: {}",
estimate
);
}
#[test]
fn test_bloom_filter_equality() {
let mut filter1 = BloomFilter::new();
let mut filter2 = BloomFilter::new();
assert_eq!(filter1, filter2);
filter1.insert(&make_node_addr(1));
assert_ne!(filter1, filter2);
filter2.insert(&make_node_addr(1));
assert_eq!(filter1, filter2);
}
// ===== BloomState Tests =====
#[test]
fn test_bloom_state_new() {
let node = make_node_addr(0);
let state = BloomState::new(node);
assert_eq!(state.own_node_addr(), &node);
assert!(!state.is_leaf_only());
assert_eq!(state.sequence(), 0);
assert_eq!(state.leaf_dependent_count(), 0);
}
#[test]
fn test_bloom_state_leaf_only() {
let node = make_node_addr(0);
let state = BloomState::leaf_only(node);
assert!(state.is_leaf_only());
}
#[test]
fn test_bloom_state_leaf_dependents() {
let node = make_node_addr(0);
let mut state = BloomState::new(node);
let leaf1 = make_node_addr(1);
let leaf2 = make_node_addr(2);
state.add_leaf_dependent(leaf1);
state.add_leaf_dependent(leaf2);
assert_eq!(state.leaf_dependent_count(), 2);
assert!(state.remove_leaf_dependent(&leaf1));
assert_eq!(state.leaf_dependent_count(), 1);
assert!(!state.remove_leaf_dependent(&leaf1)); // already removed
}
#[test]
fn test_bloom_state_debounce() {
let node = make_node_addr(0);
let peer = make_node_addr(1);
let mut state = BloomState::new(node);
state.set_update_debounce_ms(500);
state.mark_update_needed(peer);
// Should send initially
assert!(state.should_send_update(&peer, 1000));
// Record send
state.record_update_sent(peer, 1000);
state.mark_update_needed(peer);
// Should not send immediately (within debounce)
assert!(!state.should_send_update(&peer, 1200));
// Should send after debounce period
assert!(state.should_send_update(&peer, 1600));
}
#[test]
fn test_bloom_state_sequence() {
let node = make_node_addr(0);
let mut state = BloomState::new(node);
assert_eq!(state.sequence(), 0);
assert_eq!(state.next_sequence(), 1);
assert_eq!(state.next_sequence(), 2);
assert_eq!(state.sequence(), 2);
}
#[test]
fn test_bloom_state_pending_updates() {
let node = make_node_addr(0);
let mut state = BloomState::new(node);
let peer1 = make_node_addr(1);
let peer2 = make_node_addr(2);
assert!(!state.needs_update(&peer1));
state.mark_update_needed(peer1);
assert!(state.needs_update(&peer1));
assert!(!state.needs_update(&peer2));
state.mark_all_updates_needed(vec![peer1, peer2]);
assert!(state.needs_update(&peer1));
assert!(state.needs_update(&peer2));
state.clear_pending_updates();
assert!(!state.needs_update(&peer1));
assert!(!state.needs_update(&peer2));
}
#[test]
fn test_bloom_state_base_filter() {
let node = make_node_addr(0);
let mut state = BloomState::new(node);
let leaf = make_node_addr(1);
state.add_leaf_dependent(leaf);
let filter = state.base_filter();
assert!(filter.contains(&node));
assert!(filter.contains(&leaf));
assert!(!filter.contains(&make_node_addr(99)));
}
#[test]
fn test_bloom_state_compute_outgoing_filter() {
let my_node = make_node_addr(0);
let mut state = BloomState::new(my_node);
let leaf = make_node_addr(1);
state.add_leaf_dependent(leaf);
let peer1 = make_node_addr(10);
let peer2 = make_node_addr(20);
// Create peer filters
let mut filter1 = BloomFilter::new();
filter1.insert(&make_node_addr(100));
filter1.insert(&make_node_addr(101));
let mut filter2 = BloomFilter::new();
filter2.insert(&make_node_addr(200));
let mut peer_filters = HashMap::new();
peer_filters.insert(peer1, filter1);
peer_filters.insert(peer2, filter2);
// Filter for peer1 should exclude peer1's contributions
let outgoing1 = state.compute_outgoing_filter(&peer1, &peer_filters);
assert!(outgoing1.contains(&my_node)); // self
assert!(outgoing1.contains(&leaf)); // leaf dependent
assert!(outgoing1.contains(&make_node_addr(200))); // from peer2
// peer1's nodes may or may not be present (depends on split brain)
// Filter for peer2 should exclude peer2's contributions
let outgoing2 = state.compute_outgoing_filter(&peer2, &peer_filters);
assert!(outgoing2.contains(&my_node));
assert!(outgoing2.contains(&leaf));
assert!(outgoing2.contains(&make_node_addr(100))); // from peer1
assert!(outgoing2.contains(&make_node_addr(101))); // from peer1
}
src/config.rs
+1 -1
View File
@@ -1189,7 +1189,7 @@ transports:
assert_eq!(config.transports.udp.len(), 2);
let instances: std::collections::HashMap<_, _> =
config.transports.udp.iter().map(|(k, v)| (k, v)).collect();
config.transports.udp.iter().collect();
// Named instances have Some(name)
assert!(instances.contains_key(&Some("main")));
src/icmp.rs
+2 -2
View File
@@ -211,8 +211,8 @@ fn icmpv6_checksum(icmpv6_message: &[u8], src: &Ipv6Addr, dst: &Ipv6Addr) -> u16
// Pseudo-header: upper-layer packet length (4 bytes, as u32)
let len = icmpv6_message.len() as u32;
sum += (len >> 16) as u32;
sum += (len & 0xffff) as u32;
sum += len >> 16;
sum += len & 0xffff;
// Pseudo-header: next header (padded to 4 bytes)
sum += IPPROTO_ICMPV6 as u32;
src/lib.rs
-6
View File
@@ -6,7 +6,6 @@
pub mod bloom;
pub mod cache;
pub mod config;
pub mod dns;
pub mod icmp;
pub mod identity;
pub mod index;
@@ -14,11 +13,9 @@ pub mod noise;
pub mod node;
pub mod peer;
pub mod protocol;
pub mod rate_limit;
pub mod transport;
pub mod tree;
pub mod tun;
pub mod wire;
// Re-export identity types
pub use identity::{
@@ -29,9 +26,6 @@ pub use identity::{
// Re-export config types
pub use config::{Config, ConfigError, DnsConfig, IdentityConfig, TunConfig, UdpConfig};
// Re-export DNS types
pub use dns::{DnsIdentityRx, DnsIdentityTx, DnsResolvedIdentity};
// Re-export tree types
pub use tree::{CoordEntry, ParentDeclaration, TreeCoordinate, TreeError, TreeState};
src/dns.rs → src/node/dns.rs
View File
src/node/handlers/discovery.rs
+4 -4
View File
@@ -300,10 +300,10 @@ impl Node {
pub(in crate::node) async fn maybe_initiate_lookup(&mut self, dest: &NodeAddr) {
let now_ms = Self::now_ms();
let lookup_timeout_ms = self.config.node.discovery.timeout_secs * 1000;
if let Some(&initiated_at) = self.pending_lookups.get(dest) {
if now_ms.saturating_sub(initiated_at) < lookup_timeout_ms {
return;
}
if let Some(&initiated_at) = self.pending_lookups.get(dest)
&& now_ms.saturating_sub(initiated_at) < lookup_timeout_ms
{
return;
}
self.pending_lookups.insert(*dest, now_ms);
let ttl = self.config.node.discovery.ttl;
src/node/handlers/encrypted.rs
+1 -1
View File
@@ -2,7 +2,7 @@
use crate::node::Node;
use crate::transport::ReceivedPacket;
use crate::wire::EncryptedHeader;
use crate::node::wire::EncryptedHeader;
use tracing::{debug, warn};
impl Node {
src/node/handlers/handshake.rs
+18 -18
View File
@@ -5,7 +5,7 @@ use crate::peer::{
cross_connection_winner, ActivePeer, PeerConnection, PromotionResult,
};
use crate::transport::{Link, LinkDirection, LinkId, ReceivedPacket};
use crate::wire::{build_msg2, Msg1Header, Msg2Header};
use crate::node::wire::{build_msg2, Msg1Header, Msg2Header};
use crate::PeerIdentity;
use std::time::Duration;
use tracing::{debug, info, warn};
@@ -43,25 +43,25 @@ impl Node {
// (we initiated to them AND they initiated to us), this is a cross-connection.
// Allow it to proceed — promote_connection() will resolve via tie-breaker.
let addr_key = (packet.transport_id, packet.remote_addr.clone());
if let Some(&existing_link_id) = self.addr_to_link.get(&addr_key) {
if let Some(link) = self.links.get(&existing_link_id) {
if link.direction() == LinkDirection::Inbound {
self.msg1_rate_limiter.complete_handshake();
debug!(
transport_id = %packet.transport_id,
remote_addr = %packet.remote_addr,
"Already have inbound connection from this address"
);
return;
}
// Outbound link to this address — cross-connection, allow msg1
if let Some(&existing_link_id) = self.addr_to_link.get(&addr_key)
&& let Some(link) = self.links.get(&existing_link_id)
{
if link.direction() == LinkDirection::Inbound {
self.msg1_rate_limiter.complete_handshake();
debug!(
transport_id = %packet.transport_id,
remote_addr = %packet.remote_addr,
existing_link_id = %existing_link_id,
"Cross-connection detected: have outbound, received inbound msg1"
"Already have inbound connection from this address"
);
return;
}
// Outbound link to this address — cross-connection, allow msg1
debug!(
transport_id = %packet.transport_id,
remote_addr = %packet.remote_addr,
existing_link_id = %existing_link_id,
"Cross-connection detected: have outbound, received inbound msg1"
);
}
// === CRYPTO COST PAID HERE ===
@@ -89,7 +89,7 @@ impl Node {
// Learn peer identity from msg1
let peer_identity = match conn.expected_identity() {
Some(id) => id.clone(),
Some(id) => *id,
None => {
self.msg1_rate_limiter.complete_handshake();
warn!("Identity not learned from msg1");
@@ -275,7 +275,7 @@ impl Node {
// Get peer identity for promotion
let peer_identity = match conn.expected_identity() {
Some(id) => id.clone(),
Some(id) => *id,
None => {
warn!(link_id = %link_id, "No identity after handshake");
return;
@@ -398,7 +398,7 @@ impl Node {
}
// Normal path: promote to active peer
match self.promote_connection(link_id, peer_identity.clone(), packet.timestamp_ms) {
match self.promote_connection(link_id, peer_identity, packet.timestamp_ms) {
Ok(result) => {
// Clean up pending_outbound
self.pending_outbound.remove(&key);
src/node/handlers/rx_loop.rs
+1 -1
View File
@@ -2,7 +2,7 @@
use crate::node::{Node, NodeError};
use crate::transport::ReceivedPacket;
use crate::wire::{DISCRIMINATOR_ENCRYPTED, DISCRIMINATOR_MSG1, DISCRIMINATOR_MSG2};
use crate::node::wire::{DISCRIMINATOR_ENCRYPTED, DISCRIMINATOR_MSG1, DISCRIMINATOR_MSG2};
use std::time::Duration;
use tracing::{debug, info};
src/node/handlers/session.rs
+8 -8
View File
@@ -379,10 +379,10 @@ impl Node {
dest_pubkey: PublicKey,
) -> Result<(), NodeError> {
// Check for existing session
if let Some(existing) = self.sessions.get(&dest_addr) {
if existing.state().is_established() || existing.state().is_initiating() {
return Ok(());
}
if let Some(existing) = self.sessions.get(&dest_addr)
&& (existing.state().is_established() || existing.state().is_initiating())
{
return Ok(());
}
// Create Noise IK initiator handshake
@@ -638,10 +638,10 @@ impl Node {
};
// Skip if a session already exists
if let Some(existing) = self.sessions.get(&dest_addr) {
if existing.state().is_established() || existing.state().is_initiating() {
return;
}
if let Some(existing) = self.sessions.get(&dest_addr)
&& (existing.state().is_established() || existing.state().is_initiating())
{
return;
}
match self.initiate_session(dest_addr, dest_pubkey).await {
src/node/handlers/timeout.rs
+4 -4
View File
@@ -46,10 +46,10 @@ impl Node {
}
// Schedule retry for failed outbound auto-connect peers
if conn.is_outbound() {
if let Some(identity) = conn.expected_identity() {
self.schedule_retry(*identity.node_addr(), now_ms);
}
if conn.is_outbound()
&& let Some(identity) = conn.expected_identity()
{
self.schedule_retry(*identity.node_addr(), now_ms);
}
}
self.cleanup_stale_connection(link_id, now_ms);
src/node/lifecycle.rs
+3 -3
View File
@@ -5,7 +5,7 @@ use crate::peer::PeerConnection;
use crate::protocol::{Disconnect, DisconnectReason};
use crate::transport::{packet_channel, Link, LinkDirection, TransportAddr};
use crate::tun::{run_tun_reader, shutdown_tun_interface, TunDevice, TunState};
use crate::wire::build_msg1;
use crate::node::wire::build_msg1;
use crate::{NodeAddr, PeerIdentity};
use std::thread;
use std::time::Duration;
@@ -116,7 +116,7 @@ impl Node {
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0);
let mut connection = PeerConnection::outbound(link_id, peer_identity.clone(), current_time_ms);
let mut connection = PeerConnection::outbound(link_id, peer_identity, current_time_ms);
// Allocate a session index for this handshake
let our_index = match self.index_allocator.allocate() {
@@ -319,7 +319,7 @@ impl Node {
let dns_channel_size = self.config.node.buffers.dns_channel;
let (identity_tx, identity_rx) = tokio::sync::mpsc::channel(dns_channel_size);
let dns_ttl = self.config.dns.ttl();
let handle = tokio::spawn(crate::dns::run_dns_responder(socket, identity_tx, dns_ttl));
let handle = tokio::spawn(crate::node::dns::run_dns_responder(socket, identity_tx, dns_ttl));
self.dns_identity_rx = Some(identity_rx);
self.dns_task = Some(handle);
info!(bind = %bind, "DNS responder started for .fips domain");
src/node/mod.rs
+12 -11
View File
@@ -5,10 +5,13 @@
//! Bloom filters, coordinate caches, transports, links, and peers.
mod bloom;
pub(crate) mod dns;
mod handlers;
mod lifecycle;
mod retry;
mod rate_limit;
pub(crate) mod session;
pub(crate) mod wire;
mod tree;
#[cfg(test)]
mod tests;
@@ -18,14 +21,14 @@ use crate::cache::{CoordCache, RouteCache};
use crate::index::IndexAllocator;
use crate::node::session::SessionEntry;
use crate::peer::{ActivePeer, PeerConnection};
use crate::rate_limit::HandshakeRateLimiter;
use self::rate_limit::HandshakeRateLimiter;
use crate::transport::{
Link, LinkId, PacketRx, PacketTx, TransportAddr, TransportHandle, TransportId,
};
use crate::transport::udp::UdpTransport;
use crate::tree::TreeState;
use crate::tun::{TunError, TunOutboundRx, TunState, TunTx};
use crate::wire::build_encrypted;
use self::wire::build_encrypted;
use crate::{Config, ConfigError, Identity, IdentityError, NodeAddr};
use std::collections::{HashMap, VecDeque};
use std::fmt;
@@ -186,9 +189,7 @@ type AddrKey = (TransportId, TransportAddr);
///
/// The `addr_to_link` map enables dispatching incoming packets to the right
/// connection before authentication completes.
///
// Discovery lookup constants moved to config: node.discovery.timeout_secs, node.discovery.ttl
pub struct Node {
// === Identity ===
/// This node's cryptographic identity.
@@ -296,7 +297,7 @@ pub struct Node {
// === DNS Responder ===
/// Receiver for resolved identities from the DNS responder.
dns_identity_rx: Option<crate::dns::DnsIdentityRx>,
dns_identity_rx: Option<dns::DnsIdentityRx>,
/// DNS responder task handle.
dns_task: Option<tokio::task::JoinHandle<()>>,
@@ -360,7 +361,7 @@ impl Node {
let route_cache = RouteCache::new(config.node.cache.route_size);
let rl = &config.node.rate_limit;
let msg1_rate_limiter = HandshakeRateLimiter::with_params(
crate::rate_limit::TokenBucket::with_params(rl.handshake_burst, rl.handshake_rate),
rate_limit::TokenBucket::with_params(rl.handshake_burst, rl.handshake_rate),
config.node.limits.max_pending_inbound,
);
@@ -437,7 +438,7 @@ impl Node {
let route_cache = RouteCache::new(config.node.cache.route_size);
let rl = &config.node.rate_limit;
let msg1_rate_limiter = HandshakeRateLimiter::with_params(
crate::rate_limit::TokenBucket::with_params(rl.handshake_burst, rl.handshake_rate),
rate_limit::TokenBucket::with_params(rl.handshake_burst, rl.handshake_rate),
config.node.limits.max_pending_inbound,
);
@@ -903,10 +904,10 @@ impl Node {
}
// 2. Direct peer
if let Some(peer) = self.peers.get(dest_node_addr) {
if peer.can_send() {
return Some(peer);
}
if let Some(peer) = self.peers.get(dest_node_addr)
&& peer.can_send()
{
return Some(peer);
}
// Look up destination coords (required by both bloom and tree paths).
src/rate_limit.rs → src/node/rate_limit.rs
+17 -61
View File
@@ -16,7 +16,7 @@
//! - Refill rate: 10 tokens/second (sustained handshake rate)
//! - This allows handling burst traffic while limiting sustained attack impact
use std::time::{Duration, Instant};
use std::time::Instant;
/// Default burst capacity (max tokens).
pub const DEFAULT_BURST_CAPACITY: u32 = 100;
@@ -24,12 +24,6 @@ pub const DEFAULT_BURST_CAPACITY: u32 = 100;
/// Default refill rate (tokens per second).
pub const DEFAULT_REFILL_RATE: f64 = 10.0;
/// Maximum pending inbound connections.
pub const MAX_PENDING_INBOUND: usize = 1000;
/// Handshake timeout in seconds.
pub const HANDSHAKE_TIMEOUT_SECS: u64 = 30;
/// Token bucket rate limiter.
///
/// Uses a classic token bucket algorithm where tokens are consumed for each
@@ -95,27 +89,25 @@ impl TokenBucket {
}
/// Check if tokens are available without consuming them.
#[cfg(test)]
pub fn available(&mut self) -> bool {
self.refill();
self.tokens >= 1.0
}
/// Get the current number of available tokens.
#[cfg(test)]
pub fn tokens(&mut self) -> f64 {
self.refill();
self.tokens
}
/// Get the capacity (max tokens).
#[cfg(test)]
pub fn capacity(&self) -> u32 {
self.capacity
}
/// Get the refill rate (tokens per second).
pub fn refill_rate(&self) -> f64 {
self.refill_rate
}
/// Refill tokens based on elapsed time.
fn refill(&mut self) {
let now = Instant::now();
@@ -134,37 +126,26 @@ impl TokenBucket {
}
/// Reset to full capacity.
#[cfg(test)]
pub fn reset(&mut self) {
self.tokens = self.capacity as f64;
self.last_refill = Instant::now();
}
/// Set a new capacity.
pub fn set_capacity(&mut self, capacity: u32) {
self.capacity = capacity;
if self.tokens > capacity as f64 {
self.tokens = capacity as f64;
}
}
/// Set a new refill rate.
pub fn set_refill_rate(&mut self, rate: f64) {
self.refill_rate = rate;
}
/// Time until the next token is available.
///
/// Returns `Duration::ZERO` if tokens are available, otherwise the
/// estimated time until one token will be available.
pub fn time_until_available(&mut self) -> Duration {
#[cfg(test)]
pub fn time_until_available(&mut self) -> std::time::Duration {
self.refill();
if self.tokens >= 1.0 {
Duration::ZERO
std::time::Duration::ZERO
} else {
let needed = 1.0 - self.tokens;
let secs = needed / self.refill_rate;
Duration::from_secs_f64(secs)
std::time::Duration::from_secs_f64(secs)
}
}
}
@@ -190,16 +171,7 @@ pub struct HandshakeRateLimiter {
}
impl HandshakeRateLimiter {
/// Create a new handshake rate limiter with default parameters.
pub fn new() -> Self {
Self {
bucket: TokenBucket::new(),
pending_count: 0,
max_pending: MAX_PENDING_INBOUND,
}
}
/// Create with custom parameters.
/// Create a handshake rate limiter with the given parameters.
pub fn with_params(bucket: TokenBucket, max_pending: usize) -> Self {
Self {
bucket,
@@ -215,6 +187,7 @@ impl HandshakeRateLimiter {
/// - Pending connection count is below maximum
///
/// Does NOT consume a token - call `start_handshake` for that.
#[cfg(test)]
pub fn can_start_handshake(&mut self) -> bool {
self.bucket.available() && self.pending_count < self.max_pending
}
@@ -245,42 +218,25 @@ impl HandshakeRateLimiter {
}
/// Get the current pending connection count.
#[cfg(test)]
pub fn pending_count(&self) -> usize {
self.pending_count
}
/// Get the maximum pending connections.
pub fn max_pending(&self) -> usize {
self.max_pending
}
/// Set the maximum pending connections.
pub fn set_max_pending(&mut self, max: usize) {
self.max_pending = max;
}
/// Get a reference to the token bucket.
#[cfg(test)]
pub fn bucket(&self) -> &TokenBucket {
&self.bucket
}
/// Get a mutable reference to the token bucket.
pub fn bucket_mut(&mut self) -> &mut TokenBucket {
&mut self.bucket
}
/// Reset the rate limiter.
#[cfg(test)]
pub fn reset(&mut self) {
self.bucket.reset();
self.pending_count = 0;
}
}
impl Default for HandshakeRateLimiter {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
@@ -321,7 +277,7 @@ mod tests {
// Should have tokens now
let tokens = bucket.tokens();
assert!(tokens >= 4.0 && tokens <= 6.0, "tokens: {}", tokens);
assert!((4.0..=6.0).contains(&tokens), "tokens: {}", tokens);
}
#[test]
@@ -374,7 +330,7 @@ mod tests {
#[test]
fn test_handshake_rate_limiter_basic() {
let mut limiter = HandshakeRateLimiter::new();
let mut limiter = HandshakeRateLimiter::with_params(TokenBucket::new(), 100);
assert!(limiter.can_start_handshake());
assert_eq!(limiter.pending_count(), 0);
@@ -432,7 +388,7 @@ mod tests {
#[test]
fn test_handshake_rate_limiter_reset() {
let mut limiter = HandshakeRateLimiter::new();
let mut limiter = HandshakeRateLimiter::with_params(TokenBucket::new(), 100);
// Start some handshakes
limiter.start_handshake();
src/node/session.rs
+1 -1
View File
@@ -103,7 +103,7 @@ impl SessionEntry {
/// Check if the session is established.
pub(crate) fn is_established(&self) -> bool {
self.state.as_ref().map_or(false, |s| s.is_established())
self.state.as_ref().is_some_and(|s| s.is_established())
}
/// Get creation time.
src/node/tests/bloom.rs
+6 -5
View File
@@ -120,11 +120,12 @@ async fn test_bloom_filter_star() {
let filter = peer.inbound_filter().unwrap();
// Filter from hub should contain all OTHER spokes
for other in 1..5 {
for (other, other_node) in nodes[1..5].iter().enumerate() {
let other = other + 1; // adjust for slice offset
if other == spoke {
continue;
}
let other_addr = *nodes[other].node.node_addr();
let other_addr = *other_node.node.node_addr();
assert!(
filter.contains(&other_addr),
"Spoke {}'s filter from hub should contain spoke {} (addr={})",
@@ -168,12 +169,12 @@ async fn test_bloom_filter_chain_propagation() {
// Entries propagate through the full chain because each
// intermediate node merges its peer's filter into its outgoing
// filter. Verify all nodes are reachable from the endpoints.
for i in 2..8 {
for (i, addr) in addrs[2..8].iter().enumerate() {
assert!(
filter.contains(&addrs[i]),
filter.contains(addr),
"Node 0's filter from node 1 should contain node {} \
(chain merge propagation)",
i
i + 2
);
}
src/node/tests/disconnect.rs
+5 -4
View File
@@ -100,20 +100,21 @@ async fn test_disconnect_star_hub_departs() {
process_available_packets(&mut nodes).await;
// All spokes should have removed the hub
for spoke_idx in 1..4 {
for (spoke_idx, spoke) in nodes[1..4].iter().enumerate() {
let spoke_idx = spoke_idx + 1; // adjust for slice offset
assert!(
nodes[spoke_idx].node.get_peer(&hub_addr).is_none(),
spoke.node.get_peer(&hub_addr).is_none(),
"Spoke {} should have removed hub",
spoke_idx
);
assert_eq!(
nodes[spoke_idx].node.peer_count(),
spoke.node.peer_count(),
0,
"Spoke {} should have no peers (no spoke-spoke links)",
spoke_idx
);
assert!(
nodes[spoke_idx].node.tree_state().is_root(),
spoke.node.tree_state().is_root(),
"Isolated spoke {} should become root",
spoke_idx
);
src/node/tests/handshake.rs
+10 -10
View File
@@ -6,7 +6,7 @@ use super::*;
async fn test_two_node_handshake_udp() {
use crate::config::UdpConfig;
use crate::transport::udp::UdpTransport;
use crate::wire::{build_encrypted, build_msg1};
use crate::node::wire::{build_encrypted, build_msg1};
use tokio::time::{timeout, Duration};
// === Setup: Two nodes with UDP transports on localhost ===
@@ -55,7 +55,7 @@ async fn test_two_node_handshake_udp() {
let link_id_a = node_a.allocate_link_id();
let mut conn_a = PeerConnection::outbound(
link_id_a,
peer_b_identity.clone(),
peer_b_identity,
1000,
);
@@ -233,7 +233,7 @@ async fn test_two_node_handshake_udp() {
async fn test_run_rx_loop_handshake() {
use crate::config::UdpConfig;
use crate::transport::udp::UdpTransport;
use crate::wire::build_msg1;
use crate::node::wire::build_msg1;
use tokio::time::Duration;
// === Setup: Two nodes with UDP transports on localhost ===
@@ -287,7 +287,7 @@ async fn test_run_rx_loop_handshake() {
let link_id_a = node_a.allocate_link_id();
let mut conn_a = PeerConnection::outbound(
link_id_a,
peer_b_identity.clone(),
peer_b_identity,
1000,
);
@@ -423,7 +423,7 @@ async fn test_run_rx_loop_handshake() {
async fn test_cross_connection_both_initiate() {
use crate::config::UdpConfig;
use crate::transport::udp::UdpTransport;
use crate::wire::build_msg1;
use crate::node::wire::build_msg1;
use tokio::time::{timeout, Duration};
// === Setup: Two nodes with UDP transports on localhost ===
@@ -474,7 +474,7 @@ async fn test_cross_connection_both_initiate() {
// Node A initiates to Node B
let link_id_a_out = node_a.allocate_link_id();
let mut conn_a = PeerConnection::outbound(link_id_a_out, peer_b_identity.clone(), 1000);
let mut conn_a = PeerConnection::outbound(link_id_a_out, peer_b_identity, 1000);
let our_index_a = node_a.index_allocator.allocate().unwrap();
let our_keypair_a = node_a.identity.keypair();
let noise_msg1_a = conn_a.start_handshake(our_keypair_a, 1000).unwrap();
@@ -495,7 +495,7 @@ async fn test_cross_connection_both_initiate() {
// Node B initiates to Node A
let link_id_b_out = node_b.allocate_link_id();
let mut conn_b = PeerConnection::outbound(link_id_b_out, peer_a_identity.clone(), 1000);
let mut conn_b = PeerConnection::outbound(link_id_b_out, peer_a_identity, 1000);
let our_index_b = node_b.index_allocator.allocate().unwrap();
let our_keypair_b = node_b.identity.keypair();
let noise_msg1_b = conn_b.start_handshake(our_keypair_b, 1000).unwrap();
@@ -595,7 +595,7 @@ async fn test_stale_connection_cleanup() {
// Create outbound connection with a timestamp far in the past
let past_time_ms = 1000; // A very early timestamp
let link_id = node.allocate_link_id();
let mut conn = PeerConnection::outbound(link_id, peer_identity.clone(), past_time_ms);
let mut conn = PeerConnection::outbound(link_id, peer_identity, past_time_ms);
// Allocate session index and set transport info
let our_index = node.index_allocator.allocate().unwrap();
@@ -631,7 +631,7 @@ async fn test_stale_connection_cleanup() {
assert!(!node.pending_outbound.contains_key(&(transport_id, our_index.as_u32())),
"pending_outbound should be cleaned up");
assert_eq!(node.index_allocator.count(), 0, "Session index should be freed");
assert!(node.addr_to_link.get(&(transport_id, remote_addr)).is_none(),
assert!(!node.addr_to_link.contains_key(&(transport_id, remote_addr)),
"addr_to_link should be cleaned up");
}
@@ -650,7 +650,7 @@ async fn test_failed_connection_cleanup() {
.map(|d| d.as_millis() as u64)
.unwrap_or(0);
let link_id = node.allocate_link_id();
let mut conn = PeerConnection::outbound(link_id, peer_identity.clone(), now_ms);
let mut conn = PeerConnection::outbound(link_id, peer_identity, now_ms);
let our_index = node.index_allocator.allocate().unwrap();
let our_keypair = node.identity.keypair();
src/node/tests/mod.rs
+1 -1
View File
@@ -46,7 +46,7 @@ pub(super) fn make_completed_connection(
let peer_identity = PeerIdentity::from_pubkey_full(peer_identity_full.pubkey_full());
// Create outbound connection
let mut conn = PeerConnection::outbound(link_id, peer_identity.clone(), current_time_ms);
let mut conn = PeerConnection::outbound(link_id, peer_identity, current_time_ms);
// Run initiator side of handshake
let our_keypair = node.identity.keypair();
src/node/tests/routing.rs
+3 -3
View File
@@ -562,7 +562,7 @@ async fn test_routing_reachability_100_nodes() {
.collect();
for node in &mut nodes {
for &(ref addr, ref coords) in &all_coords {
for (addr, coords) in &all_coords {
if addr != node.node.node_addr() {
node.node.coord_cache_mut().insert(*addr, coords.clone(), now_ms);
}
@@ -696,7 +696,7 @@ async fn test_routing_stops_after_peer_removal() {
.collect();
for node in &mut nodes {
for &(ref addr, ref coords) in &all_coords {
for (addr, coords) in &all_coords {
if addr != node.node.node_addr() {
node.node.coord_cache_mut().insert(*addr, coords.clone(), now_ms);
}
@@ -931,7 +931,7 @@ async fn test_routing_source_only_coords_100_nodes() {
// Now compare: inject coords at ALL nodes (full cache) and verify 100%
for node in &mut nodes {
for &(ref addr, ref coords) in &all_coords {
for (addr, coords) in &all_coords {
if addr != node.node.node_addr() {
node.node.coord_cache_mut().insert(*addr, coords.clone(), now_ms);
}
src/node/tests/session.rs
+2 -2
View File
@@ -639,13 +639,13 @@ async fn test_session_100_nodes() {
let fwd_payload = format!("fwd-{}", pair_idx).into_bytes();
let rev_payload = format!("rev-{}", pair_idx).into_bytes();
if delivered_per_node[dst].iter().any(|p| *p == fwd_payload) {
if delivered_per_node[dst].contains(&fwd_payload) {
fwd_delivered += 1;
} else if fwd_missing.len() < 20 {
fwd_missing.push((src, dst));
}
if delivered_per_node[src].iter().any(|p| *p == rev_payload) {
if delivered_per_node[src].contains(&rev_payload) {
rev_delivered += 1;
} else if rev_missing.len() < 20 {
rev_missing.push((src, dst));
src/node/tests/spanning_tree.rs
+7 -7
View File
@@ -48,7 +48,7 @@ pub(super) async fn make_test_node() -> TestNode {
/// Sends msg1 over UDP. The drain loop will handle msg1 processing,
/// msg2 response, and subsequent TreeAnnounce exchange.
pub(super) async fn initiate_handshake(nodes: &mut [TestNode], i: usize, j: usize) {
use crate::wire::build_msg1;
use crate::node::wire::build_msg1;
// Extract responder info before mutably borrowing initiator
let responder_addr = nodes[j].addr.clone();
@@ -203,19 +203,19 @@ pub(super) fn print_tree_snapshot(label: &str, nodes: &[TestNode]) {
///
/// Returns the number of packets processed.
pub(super) async fn process_available_packets(nodes: &mut [TestNode]) -> usize {
use crate::wire::{DISCRIMINATOR_ENCRYPTED, DISCRIMINATOR_MSG1, DISCRIMINATOR_MSG2};
use crate::node::wire::{DISCRIMINATOR_ENCRYPTED, DISCRIMINATOR_MSG1, DISCRIMINATOR_MSG2};
let mut count = 0;
for i in 0..nodes.len() {
while let Ok(packet) = nodes[i].packet_rx.try_recv() {
for node in nodes.iter_mut() {
while let Ok(packet) = node.packet_rx.try_recv() {
if packet.data.is_empty() {
continue;
}
match packet.data[0] {
DISCRIMINATOR_MSG1 => nodes[i].node.handle_msg1(packet).await,
DISCRIMINATOR_MSG2 => nodes[i].node.handle_msg2(packet).await,
DISCRIMINATOR_MSG1 => node.node.handle_msg1(packet).await,
DISCRIMINATOR_MSG2 => node.node.handle_msg2(packet).await,
DISCRIMINATOR_ENCRYPTED => {
nodes[i].node.handle_encrypted_frame(packet).await
node.node.handle_encrypted_frame(packet).await
}
_ => {}
}
src/node/tests/unit.rs
+5 -5
View File
@@ -152,7 +152,7 @@ fn test_node_connection_duplicate() {
let identity = make_peer_identity();
let link_id = LinkId::new(1);
let conn1 = PeerConnection::outbound(link_id, identity.clone(), 1000);
let conn1 = PeerConnection::outbound(link_id, identity, 1000);
let conn2 = PeerConnection::outbound(link_id, identity, 2000);
node.add_connection(conn1).unwrap();
@@ -206,7 +206,7 @@ fn test_node_cross_connection_resolution() {
let node_addr = *identity.node_addr();
node.add_connection(conn1).unwrap();
node.promote_connection(link_id1, identity.clone(), 1500).unwrap();
node.promote_connection(link_id1, identity, 1500).unwrap();
assert_eq!(node.peer_count(), 1);
assert_eq!(node.get_peer(&node_addr).unwrap().link_id(), link_id1);
@@ -447,7 +447,7 @@ fn test_promote_cleans_up_pending_outbound_to_same_peer() {
let pending_link_id = LinkId::new(1);
let pending_time_ms = 1000;
let mut pending_conn =
PeerConnection::outbound(pending_link_id, peer_b_identity.clone(), pending_time_ms);
PeerConnection::outbound(pending_link_id, peer_b_identity, pending_time_ms);
let our_keypair = node.identity.keypair();
let _msg1 = pending_conn.start_handshake(our_keypair, pending_time_ms).unwrap();
@@ -485,7 +485,7 @@ fn test_promote_cleans_up_pending_outbound_to_same_peer() {
let mut completing_conn = PeerConnection::outbound(
completing_link_id,
peer_b_identity.clone(),
peer_b_identity,
completing_time_ms,
);
@@ -519,7 +519,7 @@ fn test_promote_cleans_up_pending_outbound_to_same_peer() {
// --- Promote the completing connection ---
let result = node
.promote_connection(completing_link_id, peer_b_identity.clone(), completing_time_ms)
.promote_connection(completing_link_id, peer_b_identity, completing_time_ms)
.unwrap();
assert!(matches!(result, PromotionResult::Promoted(_)));
src/wire.rs → src/node/wire.rs
+3 -3
View File
@@ -36,9 +36,6 @@ pub const MSG2_WIRE_SIZE: usize = 1 + 4 + 4 + HANDSHAKE_MSG2_SIZE; // 42 bytes
/// Minimum size for encrypted frame: discriminator + receiver_idx + counter + tag.
pub const ENCRYPTED_MIN_SIZE: usize = 1 + 4 + 8 + TAG_SIZE; // 29 bytes
/// Overhead added by encrypted frame wrapper.
pub const ENCRYPTED_OVERHEAD: usize = ENCRYPTED_MIN_SIZE;
// ============================================================================
// Encrypted Frame Header
// ============================================================================
@@ -85,6 +82,7 @@ impl EncryptedHeader {
}
/// Get the ciphertext slice from the original packet.
#[cfg(test)]
pub fn ciphertext<'a>(&self, data: &'a [u8]) -> &'a [u8] {
&data[self.ciphertext_offset..]
}
@@ -130,6 +128,7 @@ impl Msg1Header {
}
/// Get the Noise msg1 payload from the original packet.
#[cfg(test)]
pub fn noise_msg1<'a>(&self, data: &'a [u8]) -> &'a [u8] {
&data[self.noise_msg1_offset..]
}
@@ -179,6 +178,7 @@ impl Msg2Header {
}
/// Get the Noise msg2 payload from the original packet.
#[cfg(test)]
pub fn noise_msg2<'a>(&self, data: &'a [u8]) -> &'a [u8] {
&data[self.noise_msg2_offset..]
}
src/noise.rs
-1475
View File
File diff suppressed because it is too large Load Diff
src/noise/handshake.rs
+478
View File
@@ -0,0 +1,478 @@
use super::{
CipherState, HandshakeProgress, HandshakeRole, NoiseError, NoiseSession,
HANDSHAKE_MSG1_SIZE, HANDSHAKE_MSG2_SIZE, PROTOCOL_NAME, PUBKEY_SIZE,
};
use hkdf::Hkdf;
use rand::RngCore;
use secp256k1::{ecdh::shared_secret_point, Keypair, PublicKey, Secp256k1, SecretKey};
use sha2::{Digest, Sha256};
use std::fmt;
/// Symmetric state during handshake.
///
/// Maintains the chaining key (ck), handshake hash (h), and current cipher.
struct SymmetricState {
/// Chaining key for key derivation.
ck: [u8; 32],
/// Handshake hash for transcript binding.
h: [u8; 32],
/// Current cipher state for encrypting handshake payloads.
cipher: CipherState,
}
impl SymmetricState {
/// Initialize with protocol name.
fn initialize() -> Self {
// If protocol name <= 32 bytes, pad with zeros
// If > 32 bytes, hash it
let h = if PROTOCOL_NAME.len() <= 32 {
let mut h = [0u8; 32];
h[..PROTOCOL_NAME.len()].copy_from_slice(PROTOCOL_NAME);
h
} else {
let mut hasher = Sha256::new();
hasher.update(PROTOCOL_NAME);
hasher.finalize().into()
};
Self {
ck: h,
h,
cipher: CipherState::empty(),
}
}
/// Mix data into the handshake hash.
fn mix_hash(&mut self, data: &[u8]) {
let mut hasher = Sha256::new();
hasher.update(self.h);
hasher.update(data);
self.h = hasher.finalize().into();
}
/// Mix key material into the chaining key.
fn mix_key(&mut self, input_key_material: &[u8]) {
let hk = Hkdf::<Sha256>::new(Some(&self.ck), input_key_material);
let mut output = [0u8; 64];
hk.expand(&[], &mut output)
.expect("64 bytes is valid output length");
self.ck.copy_from_slice(&output[..32]);
// Initialize cipher with derived key for handshake encryption
let mut key = [0u8; 32];
key.copy_from_slice(&output[32..64]);
self.cipher.initialize_key(key);
}
/// Encrypt and mix into hash.
fn encrypt_and_hash(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, NoiseError> {
let ciphertext = self.cipher.encrypt(plaintext)?;
self.mix_hash(&ciphertext);
Ok(ciphertext)
}
/// Decrypt and mix ciphertext into hash.
fn decrypt_and_hash(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>, NoiseError> {
let plaintext = self.cipher.decrypt(ciphertext)?;
self.mix_hash(ciphertext);
Ok(plaintext)
}
/// Split into two cipher states for transport.
fn split(&self) -> (CipherState, CipherState) {
let hk = Hkdf::<Sha256>::new(Some(&self.ck), &[]);
let mut output = [0u8; 64];
hk.expand(&[], &mut output)
.expect("64 bytes is valid output length");
let mut k1 = [0u8; 32];
let mut k2 = [0u8; 32];
k1.copy_from_slice(&output[..32]);
k2.copy_from_slice(&output[32..64]);
(CipherState::new(k1), CipherState::new(k2))
}
/// Get the handshake hash (for channel binding).
fn handshake_hash(&self) -> [u8; 32] {
self.h
}
}
/// Handshake state for Noise IK.
pub struct HandshakeState {
/// Our role in the handshake.
role: HandshakeRole,
/// Current progress.
progress: HandshakeProgress,
/// Symmetric state.
symmetric: SymmetricState,
/// Our static keypair.
static_keypair: Keypair,
/// Our ephemeral keypair (generated at handshake start).
ephemeral_keypair: Option<Keypair>,
/// Remote static public key.
/// For initiator: known before handshake (from config).
/// For responder: learned from message 1.
remote_static: Option<PublicKey>,
/// Remote ephemeral public key (learned during handshake).
remote_ephemeral: Option<PublicKey>,
/// Secp256k1 context.
secp: Secp256k1<secp256k1::All>,
}
impl HandshakeState {
/// Normalize a compressed public key to even parity for pre-message hashing.
///
/// Nostr npubs encode x-only keys (no parity). The Noise IK pre-message
/// mixes the responder's static key into the hash before any messages.
/// Both sides must mix identical bytes. Since the initiator may only have
/// the x-only key (from an npub), we normalize to even parity (0x02 prefix)
/// so the hash chain matches regardless of the key's actual parity.
///
/// This does NOT affect ECDH operations (which use x-coordinate-only output)
/// or the keys sent in handshake messages (which use actual parity).
fn normalize_for_premessage(pubkey: &PublicKey) -> [u8; PUBKEY_SIZE] {
let mut bytes = pubkey.serialize();
bytes[0] = 0x02; // Force even parity
bytes
}
/// Create a new handshake as initiator.
///
/// The initiator knows the responder's static key and will send first.
pub fn new_initiator(static_keypair: Keypair, remote_static: PublicKey) -> Self {
let secp = Secp256k1::new();
let mut state = Self {
role: HandshakeRole::Initiator,
progress: HandshakeProgress::Initial,
symmetric: SymmetricState::initialize(),
static_keypair,
ephemeral_keypair: None,
remote_static: Some(remote_static),
remote_ephemeral: None,
secp,
};
// Mix in pre-message: <- s (responder's static is known)
// Normalize to even parity so initiator and responder hash chains match
// even when the initiator only has the x-only key (from npub).
let normalized = Self::normalize_for_premessage(&remote_static);
state.symmetric.mix_hash(&normalized);
state
}
/// Create a new handshake as responder.
///
/// The responder does NOT know the initiator's static key - it will be
/// learned from message 1.
pub fn new_responder(static_keypair: Keypair) -> Self {
let secp = Secp256k1::new();
let mut state = Self {
role: HandshakeRole::Responder,
progress: HandshakeProgress::Initial,
symmetric: SymmetricState::initialize(),
static_keypair,
ephemeral_keypair: None,
remote_static: None, // Will learn from message 1
remote_ephemeral: None,
secp,
};
// Mix in pre-message: <- s (our static, since we're responder)
// Normalize to even parity to match initiator's hash chain.
let normalized = Self::normalize_for_premessage(&state.static_keypair.public_key());
state.symmetric.mix_hash(&normalized);
state
}
/// Get our role.
pub fn role(&self) -> HandshakeRole {
self.role
}
/// Get current progress.
pub fn progress(&self) -> HandshakeProgress {
self.progress
}
/// Check if handshake is complete.
pub fn is_complete(&self) -> bool {
self.progress == HandshakeProgress::Complete
}
/// Get the remote static key (available after message 1 for responder).
pub fn remote_static(&self) -> Option<&PublicKey> {
self.remote_static.as_ref()
}
/// Generate ephemeral keypair.
fn generate_ephemeral(&mut self) {
let mut rng = rand::thread_rng();
let mut secret_bytes = [0u8; 32];
rng.fill_bytes(&mut secret_bytes);
let secret_key =
SecretKey::from_slice(&secret_bytes).expect("32 random bytes is valid secret key");
self.ephemeral_keypair = Some(Keypair::from_secret_key(&self.secp, &secret_key));
}
/// Perform ECDH between our secret and their public key.
///
/// Uses x-only hashing (SHA-256 of just the x-coordinate) to produce
/// a parity-independent shared secret. This is necessary because Nostr
/// npubs encode x-only keys without parity information, so the initiator
/// may have the wrong parity for the responder's static key. Since P and
/// -P produce ECDH result points with the same x-coordinate, hashing
/// only x ensures both sides derive the same shared secret.
fn ecdh(&self, our_secret: &SecretKey, their_public: &PublicKey) -> [u8; 32] {
// Get raw (x, y) coordinates (64 bytes) without any hashing
let point = shared_secret_point(their_public, our_secret);
// Hash only the x-coordinate (first 32 bytes), ignoring y/parity
let mut hasher = Sha256::new();
hasher.update(&point[..32]);
let hash = hasher.finalize();
let mut result = [0u8; 32];
result.copy_from_slice(&hash);
result
}
/// Write message 1 (initiator only).
///
/// Message 1 contains:
/// - e: ephemeral public key (33 bytes)
/// - encrypted s: our static public key encrypted (33 + 16 = 49 bytes)
///
/// Total: 82 bytes
pub fn write_message_1(&mut self) -> Result<Vec<u8>, NoiseError> {
if self.role != HandshakeRole::Initiator {
return Err(NoiseError::WrongState {
expected: "initiator".to_string(),
got: "responder".to_string(),
});
}
if self.progress != HandshakeProgress::Initial {
return Err(NoiseError::WrongState {
expected: HandshakeProgress::Initial.to_string(),
got: self.progress.to_string(),
});
}
let remote_static = self.remote_static.expect("initiator must have remote static");
// Generate ephemeral keypair
self.generate_ephemeral();
let ephemeral = self.ephemeral_keypair.as_ref().unwrap();
let e_pub = ephemeral.public_key().serialize();
let mut message = Vec::with_capacity(HANDSHAKE_MSG1_SIZE);
// -> e: send ephemeral, mix into hash
message.extend_from_slice(&e_pub);
self.symmetric.mix_hash(&e_pub);
// -> es: DH(e, rs), mix into key
let es = self.ecdh(&ephemeral.secret_key(), &remote_static);
self.symmetric.mix_key(&es);
// -> s: encrypt our static and send
let our_static = self.static_keypair.public_key().serialize();
let encrypted_static = self.symmetric.encrypt_and_hash(&our_static)?;
message.extend_from_slice(&encrypted_static);
// -> ss: DH(s, rs), mix into key
let ss = self.ecdh(&self.static_keypair.secret_key(), &remote_static);
self.symmetric.mix_key(&ss);
self.progress = HandshakeProgress::Message1Done;
Ok(message)
}
/// Read message 1 (responder only).
///
/// Processes the initiator's first message and learns their identity.
pub fn read_message_1(&mut self, message: &[u8]) -> Result<(), NoiseError> {
if self.role != HandshakeRole::Responder {
return Err(NoiseError::WrongState {
expected: "responder".to_string(),
got: "initiator".to_string(),
});
}
if self.progress != HandshakeProgress::Initial {
return Err(NoiseError::WrongState {
expected: HandshakeProgress::Initial.to_string(),
got: self.progress.to_string(),
});
}
if message.len() != HANDSHAKE_MSG1_SIZE {
return Err(NoiseError::MessageTooShort {
expected: HANDSHAKE_MSG1_SIZE,
got: message.len(),
});
}
// -> e: parse remote ephemeral, mix into hash
let re = PublicKey::from_slice(&message[..PUBKEY_SIZE])
.map_err(|_| NoiseError::InvalidPublicKey)?;
self.remote_ephemeral = Some(re);
self.symmetric.mix_hash(&message[..PUBKEY_SIZE]);
// -> es: DH(s, re), mix into key
// (responder uses their static with initiator's ephemeral)
let es = self.ecdh(&self.static_keypair.secret_key(), &re);
self.symmetric.mix_key(&es);
// -> s: decrypt initiator's static
let encrypted_static = &message[PUBKEY_SIZE..];
let decrypted_static = self.symmetric.decrypt_and_hash(encrypted_static)?;
let rs =
PublicKey::from_slice(&decrypted_static).map_err(|_| NoiseError::InvalidPublicKey)?;
self.remote_static = Some(rs);
// -> ss: DH(s, rs), mix into key
let ss = self.ecdh(&self.static_keypair.secret_key(), &rs);
self.symmetric.mix_key(&ss);
self.progress = HandshakeProgress::Message1Done;
Ok(())
}
/// Write message 2 (responder only).
///
/// Message 2 contains:
/// - e: ephemeral public key (33 bytes)
///
/// Total: 33 bytes
pub fn write_message_2(&mut self) -> Result<Vec<u8>, NoiseError> {
if self.role != HandshakeRole::Responder {
return Err(NoiseError::WrongState {
expected: "responder".to_string(),
got: "initiator".to_string(),
});
}
if self.progress != HandshakeProgress::Message1Done {
return Err(NoiseError::WrongState {
expected: HandshakeProgress::Message1Done.to_string(),
got: self.progress.to_string(),
});
}
let re = self.remote_ephemeral.expect("should have remote ephemeral");
// Generate ephemeral keypair
self.generate_ephemeral();
let ephemeral = self.ephemeral_keypair.as_ref().unwrap();
let e_pub = ephemeral.public_key().serialize();
// <- e: send ephemeral, mix into hash
self.symmetric.mix_hash(&e_pub);
// <- ee: DH(e, re), mix into key
let ee = self.ecdh(&ephemeral.secret_key(), &re);
self.symmetric.mix_key(&ee);
// <- se: DH(s, re), mix into key
let se = self.ecdh(&self.static_keypair.secret_key(), &re);
self.symmetric.mix_key(&se);
self.progress = HandshakeProgress::Complete;
Ok(e_pub.to_vec())
}
/// Read message 2 (initiator only).
///
/// Processes the responder's message and completes the handshake.
pub fn read_message_2(&mut self, message: &[u8]) -> Result<(), NoiseError> {
if self.role != HandshakeRole::Initiator {
return Err(NoiseError::WrongState {
expected: "initiator".to_string(),
got: "responder".to_string(),
});
}
if self.progress != HandshakeProgress::Message1Done {
return Err(NoiseError::WrongState {
expected: HandshakeProgress::Message1Done.to_string(),
got: self.progress.to_string(),
});
}
if message.len() != HANDSHAKE_MSG2_SIZE {
return Err(NoiseError::MessageTooShort {
expected: HANDSHAKE_MSG2_SIZE,
got: message.len(),
});
}
// <- e: parse remote ephemeral, mix into hash
let re = PublicKey::from_slice(message).map_err(|_| NoiseError::InvalidPublicKey)?;
self.remote_ephemeral = Some(re);
self.symmetric.mix_hash(message);
// <- ee: DH(e, re), mix into key
let ephemeral = self.ephemeral_keypair.as_ref().unwrap();
let ee = self.ecdh(&ephemeral.secret_key(), &re);
self.symmetric.mix_key(&ee);
// <- se: DH(e, rs), mix into key
// (initiator uses their ephemeral with responder's static)
let rs = self.remote_static.expect("initiator has remote static");
let se = self.ecdh(&ephemeral.secret_key(), &rs);
self.symmetric.mix_key(&se);
self.progress = HandshakeProgress::Complete;
Ok(())
}
/// Complete the handshake and return a NoiseSession.
///
/// Must be called after the handshake is complete.
pub fn into_session(self) -> Result<NoiseSession, NoiseError> {
if !self.is_complete() {
return Err(NoiseError::HandshakeNotComplete);
}
let (c1, c2) = self.symmetric.split();
let handshake_hash = self.symmetric.handshake_hash();
let remote_static = self
.remote_static
.expect("remote static must be known after handshake");
// Initiator sends with c1, receives with c2
// Responder sends with c2, receives with c1
let (send_cipher, recv_cipher) = match self.role {
HandshakeRole::Initiator => (c1, c2),
HandshakeRole::Responder => (c2, c1),
};
Ok(NoiseSession::from_handshake(
self.role,
send_cipher,
recv_cipher,
handshake_hash,
remote_static,
))
}
/// Get the handshake hash (for channel binding, available after complete).
pub fn handshake_hash(&self) -> [u8; 32] {
self.symmetric.handshake_hash()
}
}
impl fmt::Debug for HandshakeState {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("HandshakeState")
.field("role", &self.role)
.field("progress", &self.progress)
.field("has_ephemeral", &self.ephemeral_keypair.is_some())
.field("has_remote_static", &self.remote_static.is_some())
.field("has_remote_ephemeral", &self.remote_ephemeral.is_some())
.finish()
}
}
src/noise/mod.rs
+314
View File
@@ -0,0 +1,314 @@
//! Noise IK Protocol for Peer Authentication
//!
//! Implements the Noise Protocol Framework IK pattern using secp256k1
//! for link-local peer authentication. This establishes encrypted
//! channels between direct peers over a transport.
//!
//! The IK pattern assumes the initiator knows the responder's static
//! public key before the handshake. The responder learns the initiator's
//! identity from the encrypted payload in message 1.
//!
//! ## Handshake Pattern
//!
//! Pre-message (key known before handshake):
//! ```text
//! <- s (responder's static known to initiator)
//! ```
//!
//! Messages:
//! ```text
//! -> e, es, s, ss (initiator sends ephemeral + encrypted static)
//! <- e, ee, se (responder sends ephemeral)
//! ```
//!
//! After handshake, both parties derive symmetric keys for bidirectional
//! encrypted communication over the peer link.
//!
//! ## Separation of Concerns
//!
//! This module handles **peer authentication** only - securing the direct
//! link between neighboring nodes. End-to-end FIPS session encryption
//! between arbitrary network addresses is a separate concern handled by
//! the session layer.
mod handshake;
mod replay;
mod session;
use chacha20poly1305::{
aead::{Aead, KeyInit},
ChaCha20Poly1305, Nonce,
};
use std::fmt;
use thiserror::Error;
pub use handshake::HandshakeState;
pub use replay::ReplayWindow;
pub use session::NoiseSession;
/// Protocol name for Noise IK with secp256k1.
/// Format: Noise_IK_secp256k1_ChaChaPoly_SHA256
pub(crate) const PROTOCOL_NAME: &[u8] = b"Noise_IK_secp256k1_ChaChaPoly_SHA256";
/// Maximum message size for noise transport messages.
pub const MAX_MESSAGE_SIZE: usize = 65535;
/// Size of the AEAD tag.
pub const TAG_SIZE: usize = 16;
/// Size of a public key (compressed secp256k1).
pub const PUBKEY_SIZE: usize = 33;
/// Size of handshake message 1: ephemeral (33) + encrypted static (33 + 16 tag).
pub const HANDSHAKE_MSG1_SIZE: usize = PUBKEY_SIZE + PUBKEY_SIZE + TAG_SIZE;
/// Size of handshake message 2: ephemeral only.
pub const HANDSHAKE_MSG2_SIZE: usize = PUBKEY_SIZE;
/// Replay window size in packets (matching WireGuard).
pub const REPLAY_WINDOW_SIZE: usize = 2048;
/// Errors from Noise protocol operations.
#[derive(Debug, Error)]
pub enum NoiseError {
#[error("handshake not complete")]
HandshakeNotComplete,
#[error("handshake already complete")]
HandshakeAlreadyComplete,
#[error("wrong handshake state: expected {expected}, got {got}")]
WrongState { expected: String, got: String },
#[error("invalid public key")]
InvalidPublicKey,
#[error("decryption failed")]
DecryptionFailed,
#[error("encryption failed")]
EncryptionFailed,
#[error("message too large: {size} > {max}")]
MessageTooLarge { size: usize, max: usize },
#[error("message too short: expected at least {expected}, got {got}")]
MessageTooShort { expected: usize, got: usize },
#[error("nonce overflow")]
NonceOverflow,
#[error("replay detected: counter {0} already seen or too old")]
ReplayDetected(u64),
#[error("secp256k1 error: {0}")]
Secp256k1(#[from] secp256k1::Error),
}
/// Role in the handshake.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HandshakeRole {
/// We initiated the connection.
Initiator,
/// They initiated the connection.
Responder,
}
impl fmt::Display for HandshakeRole {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
HandshakeRole::Initiator => write!(f, "initiator"),
HandshakeRole::Responder => write!(f, "responder"),
}
}
}
/// Handshake state machine states.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HandshakeProgress {
/// Initial state, ready to send/receive message 1.
Initial,
/// Message 1 sent/received, ready for message 2.
Message1Done,
/// Handshake complete, ready for transport.
Complete,
}
impl fmt::Display for HandshakeProgress {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
HandshakeProgress::Initial => write!(f, "initial"),
HandshakeProgress::Message1Done => write!(f, "message1_done"),
HandshakeProgress::Complete => write!(f, "complete"),
}
}
}
/// Symmetric cipher state for post-handshake encryption.
#[derive(Clone)]
pub struct CipherState {
/// Encryption key (32 bytes).
key: [u8; 32],
/// Nonce counter (8 bytes used, 4 bytes zero prefix).
pub(super) nonce: u64,
/// Whether this cipher has a valid key.
has_key: bool,
}
impl CipherState {
/// Create a new cipher state with the given key.
pub(crate) fn new(key: [u8; 32]) -> Self {
Self {
key,
nonce: 0,
has_key: true,
}
}
/// Create an empty cipher state (no key yet).
pub(super) fn empty() -> Self {
Self {
key: [0u8; 32],
nonce: 0,
has_key: false,
}
}
/// Initialize with a key.
pub(super) fn initialize_key(&mut self, key: [u8; 32]) {
self.key = key;
self.nonce = 0;
self.has_key = true;
}
/// Encrypt plaintext, returning ciphertext with appended tag.
pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, NoiseError> {
if !self.has_key {
// No key means no encryption (shouldn't happen in transport phase)
return Ok(plaintext.to_vec());
}
if plaintext.len() > MAX_MESSAGE_SIZE - TAG_SIZE {
return Err(NoiseError::MessageTooLarge {
size: plaintext.len(),
max: MAX_MESSAGE_SIZE - TAG_SIZE,
});
}
let cipher = ChaCha20Poly1305::new_from_slice(&self.key)
.map_err(|_| NoiseError::EncryptionFailed)?;
let nonce = self.next_nonce()?;
let ciphertext = cipher
.encrypt(&nonce, plaintext)
.map_err(|_| NoiseError::EncryptionFailed)?;
Ok(ciphertext)
}
/// Decrypt ciphertext (with appended tag), returning plaintext.
///
/// Uses the internal nonce counter. For transport phase with explicit
/// counters from the wire format, use `decrypt_with_counter` instead.
pub fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>, NoiseError> {
if !self.has_key {
// No key means no encryption
return Ok(ciphertext.to_vec());
}
if ciphertext.len() < TAG_SIZE {
return Err(NoiseError::MessageTooShort {
expected: TAG_SIZE,
got: ciphertext.len(),
});
}
let cipher = ChaCha20Poly1305::new_from_slice(&self.key)
.map_err(|_| NoiseError::DecryptionFailed)?;
let nonce = self.next_nonce()?;
let plaintext = cipher
.decrypt(&nonce, ciphertext)
.map_err(|_| NoiseError::DecryptionFailed)?;
Ok(plaintext)
}
/// Decrypt with an explicit counter value (for transport phase).
///
/// This is used when the counter comes from the wire format rather than
/// an internal counter. The counter must be validated by a replay window
/// before calling this method.
pub fn decrypt_with_counter(
&self,
ciphertext: &[u8],
counter: u64,
) -> Result<Vec<u8>, NoiseError> {
if !self.has_key {
return Ok(ciphertext.to_vec());
}
if ciphertext.len() < TAG_SIZE {
return Err(NoiseError::MessageTooShort {
expected: TAG_SIZE,
got: ciphertext.len(),
});
}
let cipher = ChaCha20Poly1305::new_from_slice(&self.key)
.map_err(|_| NoiseError::DecryptionFailed)?;
let nonce = Self::counter_to_nonce(counter);
let plaintext = cipher
.decrypt(&nonce, ciphertext)
.map_err(|_| NoiseError::DecryptionFailed)?;
Ok(plaintext)
}
/// Convert a counter value to a nonce.
fn counter_to_nonce(counter: u64) -> Nonce {
let mut nonce_bytes = [0u8; 12];
nonce_bytes[4..12].copy_from_slice(&counter.to_le_bytes());
*Nonce::from_slice(&nonce_bytes)
}
/// Get the next nonce, incrementing the counter.
fn next_nonce(&mut self) -> Result<Nonce, NoiseError> {
if self.nonce == u64::MAX {
return Err(NoiseError::NonceOverflow);
}
let n = self.nonce;
self.nonce += 1;
// Noise uses 8-byte counter with 4-byte zero prefix
let mut nonce_bytes = [0u8; 12];
nonce_bytes[4..12].copy_from_slice(&n.to_le_bytes());
Ok(*Nonce::from_slice(&nonce_bytes))
}
/// Get the current nonce value (for debugging/testing).
pub fn nonce(&self) -> u64 {
self.nonce
}
/// Check if cipher has a key.
pub fn has_key(&self) -> bool {
self.has_key
}
}
impl fmt::Debug for CipherState {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("CipherState")
.field("nonce", &self.nonce)
.field("has_key", &self.has_key)
.field("key", &"[redacted]")
.finish()
}
}
#[cfg(test)]
mod tests;
src/noise/replay.rs
+138
View File
@@ -0,0 +1,138 @@
use super::REPLAY_WINDOW_SIZE;
use std::fmt;
/// Sliding window for replay protection.
///
/// Tracks which packet counters have been received within a window of
/// REPLAY_WINDOW_SIZE. Packets with counters below the window or already
/// seen within the window are rejected.
///
/// Based on WireGuard's anti-replay mechanism (RFC 6479 style).
#[derive(Clone)]
pub struct ReplayWindow {
/// Highest counter value seen.
highest: u64,
/// Bitmap tracking which counters in the window have been seen.
/// Bit i corresponds to counter (highest - i).
bitmap: [u64; REPLAY_WINDOW_SIZE / 64],
}
impl ReplayWindow {
/// Create a new replay window.
pub fn new() -> Self {
Self {
highest: 0,
bitmap: [0; REPLAY_WINDOW_SIZE / 64],
}
}
/// Check if a counter is valid (not replayed, not too old).
///
/// Returns true if the counter is acceptable, false if it should be rejected.
/// Does NOT update the window - call `accept` after successful decryption.
pub fn check(&self, counter: u64) -> bool {
if counter > self.highest {
// New highest - always acceptable
return true;
}
// Counter is <= highest, check if it's within the window
let diff = self.highest - counter;
if diff as usize >= REPLAY_WINDOW_SIZE {
// Too old (outside window)
return false;
}
// Check bitmap - bit is set if counter was already seen
let word_idx = (diff as usize) / 64;
let bit_idx = (diff as usize) % 64;
(self.bitmap[word_idx] & (1u64 << bit_idx)) == 0
}
/// Accept a counter into the window.
///
/// Call this only after successful decryption to prevent
/// DoS attacks that exhaust the window.
pub fn accept(&mut self, counter: u64) {
if counter > self.highest {
// Shift the window
let shift = counter - self.highest;
if shift as usize >= REPLAY_WINDOW_SIZE {
// Complete reset
self.bitmap = [0; REPLAY_WINDOW_SIZE / 64];
} else {
// Shift bitmap
self.shift_bitmap(shift as usize);
}
self.highest = counter;
// Mark counter 0 (which is now the highest) as seen
self.bitmap[0] |= 1;
} else {
// Mark the counter as seen
let diff = self.highest - counter;
let word_idx = (diff as usize) / 64;
let bit_idx = (diff as usize) % 64;
self.bitmap[word_idx] |= 1u64 << bit_idx;
}
}
/// Shift the bitmap by the given number of positions.
///
/// This moves old counters to higher bit positions to make room for the
/// new highest counter at position 0.
fn shift_bitmap(&mut self, shift: usize) {
if shift >= REPLAY_WINDOW_SIZE {
self.bitmap = [0; REPLAY_WINDOW_SIZE / 64];
return;
}
let word_shift = shift / 64;
let bit_shift = shift % 64;
// Shift entire words first (from high to low to avoid overwriting)
if word_shift > 0 {
for i in (word_shift..self.bitmap.len()).rev() {
self.bitmap[i] = self.bitmap[i - word_shift];
}
for i in 0..word_shift {
self.bitmap[i] = 0;
}
}
// Shift bits within words (from low to high so carry propagates correctly)
if bit_shift > 0 {
let mut carry = 0u64;
for i in 0..self.bitmap.len() {
let new_carry = self.bitmap[i] >> (64 - bit_shift);
self.bitmap[i] = (self.bitmap[i] << bit_shift) | carry;
carry = new_carry;
}
}
}
/// Get the highest counter seen.
pub fn highest(&self) -> u64 {
self.highest
}
/// Reset the window (use when rekeying).
pub fn reset(&mut self) {
self.highest = 0;
self.bitmap = [0; REPLAY_WINDOW_SIZE / 64];
}
}
impl Default for ReplayWindow {
fn default() -> Self {
Self::new()
}
}
impl fmt::Debug for ReplayWindow {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ReplayWindow")
.field("highest", &self.highest)
.field("window_size", &REPLAY_WINDOW_SIZE)
.finish()
}
}
src/noise/session.rs
+157
View File
@@ -0,0 +1,157 @@
use super::{CipherState, HandshakeRole, NoiseError, ReplayWindow};
use secp256k1::{PublicKey, XOnlyPublicKey};
use std::fmt;
/// Completed Noise session for transport encryption.
///
/// Provides bidirectional authenticated encryption with replay protection.
/// The send counter is monotonically incremented; received counters are
/// validated against a sliding window to prevent replay attacks.
pub struct NoiseSession {
/// Our role in the original handshake.
role: HandshakeRole,
/// Cipher for sending.
send_cipher: CipherState,
/// Cipher for receiving.
recv_cipher: CipherState,
/// Handshake hash for channel binding.
handshake_hash: [u8; 32],
/// Remote peer's static public key.
remote_static: PublicKey,
/// Replay window for received packets.
replay_window: ReplayWindow,
}
impl NoiseSession {
/// Create a new session from completed handshake data.
pub(super) fn from_handshake(
role: HandshakeRole,
send_cipher: CipherState,
recv_cipher: CipherState,
handshake_hash: [u8; 32],
remote_static: PublicKey,
) -> Self {
Self {
role,
send_cipher,
recv_cipher,
handshake_hash,
remote_static,
replay_window: ReplayWindow::new(),
}
}
/// Encrypt a message for sending (using internal counter).
///
/// Returns the ciphertext. The current send counter should be included
/// in the wire format before calling this method.
pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, NoiseError> {
self.send_cipher.encrypt(plaintext)
}
/// Get the current send counter (before incrementing).
///
/// Use this to get the counter to include in the wire format.
/// The counter will be incremented when `encrypt` is called.
pub fn current_send_counter(&self) -> u64 {
self.send_cipher.nonce
}
/// Decrypt a received message (using internal counter).
///
/// This is for handshake-phase decryption. For transport phase with
/// explicit counters, use `decrypt_with_replay_check` instead.
pub fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>, NoiseError> {
self.recv_cipher.decrypt(ciphertext)
}
/// Check if a counter passes the replay window.
///
/// Returns Ok(()) if the counter is acceptable, Err if it should be rejected.
/// Call this before attempting decryption to avoid wasting CPU on replay attacks.
pub fn check_replay(&self, counter: u64) -> Result<(), NoiseError> {
if self.replay_window.check(counter) {
Ok(())
} else {
Err(NoiseError::ReplayDetected(counter))
}
}
/// Decrypt with explicit counter and replay protection.
///
/// This is the primary decryption method for transport phase.
/// The counter comes from the wire format and is validated against
/// the replay window before and after decryption.
///
/// On success, the counter is accepted into the replay window.
pub fn decrypt_with_replay_check(
&mut self,
ciphertext: &[u8],
counter: u64,
) -> Result<Vec<u8>, NoiseError> {
// Check replay window first (cheap)
if !self.replay_window.check(counter) {
return Err(NoiseError::ReplayDetected(counter));
}
// Attempt decryption (expensive)
let plaintext = self.recv_cipher.decrypt_with_counter(ciphertext, counter)?;
// Only accept into window after successful decryption
// This prevents DoS attacks that exhaust the window
self.replay_window.accept(counter);
Ok(plaintext)
}
/// Get the highest received counter.
pub fn highest_received_counter(&self) -> u64 {
self.replay_window.highest()
}
/// Reset the replay window (use when rekeying).
pub fn reset_replay_window(&mut self) {
self.replay_window.reset();
}
/// Get the handshake hash for channel binding.
pub fn handshake_hash(&self) -> &[u8; 32] {
&self.handshake_hash
}
/// Get the remote peer's static public key.
pub fn remote_static(&self) -> &PublicKey {
&self.remote_static
}
/// Get the remote peer's x-only public key.
pub fn remote_static_xonly(&self) -> XOnlyPublicKey {
self.remote_static.x_only_public_key().0
}
/// Get our role in the handshake.
pub fn role(&self) -> HandshakeRole {
self.role
}
/// Get the send nonce (for debugging).
pub fn send_nonce(&self) -> u64 {
self.send_cipher.nonce()
}
/// Get the receive nonce (for debugging).
pub fn recv_nonce(&self) -> u64 {
self.recv_cipher.nonce()
}
}
impl fmt::Debug for NoiseSession {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("NoiseSession")
.field("role", &self.role)
.field("send_nonce", &self.send_cipher.nonce())
.field("recv_nonce", &self.recv_cipher.nonce())
.field("handshake_hash", &hex::encode(&self.handshake_hash[..8]))
.finish()
}
}
src/noise/tests.rs
+418
View File
@@ -0,0 +1,418 @@
use super::*;
use secp256k1::Parity;
fn generate_keypair() -> secp256k1::Keypair {
let secp = secp256k1::Secp256k1::new();
let mut rng = rand::thread_rng();
let (secret_key, _) = secp.generate_keypair(&mut rng);
secp256k1::Keypair::from_secret_key(&secp, &secret_key)
}
#[test]
fn test_full_handshake() {
let initiator_keypair = generate_keypair();
let responder_keypair = generate_keypair();
let responder_pub = responder_keypair.public_key();
// Initiator knows responder's static key
// Responder does NOT know initiator's static key (IK pattern)
let mut initiator = HandshakeState::new_initiator(initiator_keypair, responder_pub);
let mut responder = HandshakeState::new_responder(responder_keypair);
assert_eq!(initiator.role(), HandshakeRole::Initiator);
assert_eq!(responder.role(), HandshakeRole::Responder);
// Initially, responder doesn't know initiator's identity
assert!(responder.remote_static().is_none());
// Message 1: Initiator -> Responder
let msg1 = initiator.write_message_1().unwrap();
assert_eq!(msg1.len(), HANDSHAKE_MSG1_SIZE);
responder.read_message_1(&msg1).unwrap();
// Now responder knows initiator's identity!
assert!(responder.remote_static().is_some());
assert_eq!(
responder.remote_static().unwrap(),
&initiator_keypair.public_key()
);
// Message 2: Responder -> Initiator
let msg2 = responder.write_message_2().unwrap();
assert_eq!(msg2.len(), HANDSHAKE_MSG2_SIZE);
initiator.read_message_2(&msg2).unwrap();
// Both should be complete
assert!(initiator.is_complete());
assert!(responder.is_complete());
// Handshake hashes should match
assert_eq!(initiator.handshake_hash(), responder.handshake_hash());
// Convert to sessions
let mut initiator_session = initiator.into_session().unwrap();
let mut responder_session = responder.into_session().unwrap();
// Test encryption/decryption
let plaintext = b"Hello, secure world!";
let ciphertext = initiator_session.encrypt(plaintext).unwrap();
let decrypted = responder_session.decrypt(&ciphertext).unwrap();
assert_eq!(decrypted, plaintext);
// Test reverse direction
let plaintext2 = b"Hello back!";
let ciphertext2 = responder_session.encrypt(plaintext2).unwrap();
let decrypted2 = initiator_session.decrypt(&ciphertext2).unwrap();
assert_eq!(decrypted2, plaintext2);
}
#[test]
fn test_multiple_messages() {
let initiator_keypair = generate_keypair();
let responder_keypair = generate_keypair();
let mut initiator =
HandshakeState::new_initiator(initiator_keypair, responder_keypair.public_key());
let mut responder = HandshakeState::new_responder(responder_keypair);
let msg1 = initiator.write_message_1().unwrap();
responder.read_message_1(&msg1).unwrap();
let msg2 = responder.write_message_2().unwrap();
initiator.read_message_2(&msg2).unwrap();
let mut initiator_session = initiator.into_session().unwrap();
let mut responder_session = responder.into_session().unwrap();
// Send many messages to test nonce increment
for i in 0..100 {
let msg = format!("Message {}", i);
let ct = initiator_session.encrypt(msg.as_bytes()).unwrap();
let pt = responder_session.decrypt(&ct).unwrap();
assert_eq!(pt, msg.as_bytes());
}
assert_eq!(initiator_session.send_nonce(), 100);
assert_eq!(responder_session.recv_nonce(), 100);
}
#[test]
fn test_wrong_role_errors() {
let keypair1 = generate_keypair();
let keypair2 = generate_keypair();
let mut initiator = HandshakeState::new_initiator(keypair1, keypair2.public_key());
// Initiator can't read message 1
assert!(initiator
.read_message_1(&[0u8; HANDSHAKE_MSG1_SIZE])
.is_err());
// Initiator can't write message 2 before message 1
assert!(initiator.write_message_2().is_err());
}
#[test]
fn test_invalid_pubkey_in_msg1() {
let keypair = generate_keypair();
let mut responder = HandshakeState::new_responder(keypair);
// Invalid pubkey bytes (first 33 bytes are zero)
let invalid_msg = [0u8; HANDSHAKE_MSG1_SIZE];
assert!(responder.read_message_1(&invalid_msg).is_err());
}
#[test]
fn test_decryption_failure_wrong_key() {
let keypair1 = generate_keypair();
let keypair2 = generate_keypair();
let keypair3 = generate_keypair();
// Session between 1 and 2
let mut init1 = HandshakeState::new_initiator(keypair1, keypair2.public_key());
let mut resp1 = HandshakeState::new_responder(keypair2);
let msg1 = init1.write_message_1().unwrap();
resp1.read_message_1(&msg1).unwrap();
let msg2 = resp1.write_message_2().unwrap();
init1.read_message_2(&msg2).unwrap();
let mut session1 = init1.into_session().unwrap();
// Session between 1 and 3
let mut init2 = HandshakeState::new_initiator(keypair1, keypair3.public_key());
let mut resp2 = HandshakeState::new_responder(keypair3);
let msg1 = init2.write_message_1().unwrap();
resp2.read_message_1(&msg1).unwrap();
let msg2 = resp2.write_message_2().unwrap();
init2.read_message_2(&msg2).unwrap();
let mut session2 = resp2.into_session().unwrap();
// Encrypt with session 1, try to decrypt with session 2
let ciphertext = session1.encrypt(b"test").unwrap();
assert!(session2.decrypt(&ciphertext).is_err());
}
#[test]
fn test_cipher_state_nonce_sequence() {
let key = [0u8; 32];
let mut cipher = CipherState::new(key);
assert_eq!(cipher.nonce(), 0);
let _ = cipher.encrypt(b"test").unwrap();
assert_eq!(cipher.nonce(), 1);
let _ = cipher.encrypt(b"test").unwrap();
assert_eq!(cipher.nonce(), 2);
}
#[test]
fn test_session_remote_static() {
let keypair1 = generate_keypair();
let keypair2 = generate_keypair();
let mut init = HandshakeState::new_initiator(keypair1, keypair2.public_key());
let mut resp = HandshakeState::new_responder(keypair2);
let msg1 = init.write_message_1().unwrap();
resp.read_message_1(&msg1).unwrap();
let msg2 = resp.write_message_2().unwrap();
init.read_message_2(&msg2).unwrap();
let session1 = init.into_session().unwrap();
let session2 = resp.into_session().unwrap();
// Each session should know the other's static key
assert_eq!(session1.remote_static(), &keypair2.public_key());
assert_eq!(session2.remote_static(), &keypair1.public_key());
}
#[test]
fn test_message_sizes() {
// Verify our size constants are correct
assert_eq!(HANDSHAKE_MSG1_SIZE, 33 + 33 + 16); // e + encrypted_s
assert_eq!(HANDSHAKE_MSG2_SIZE, 33); // e only
}
#[test]
fn test_responder_identity_discovery() {
// This test verifies the key IK property: responder learns initiator's identity
let initiator_keypair = generate_keypair();
let responder_keypair = generate_keypair();
let mut responder = HandshakeState::new_responder(responder_keypair);
// Before message 1: responder has no idea who's connecting
assert!(responder.remote_static().is_none());
let mut initiator =
HandshakeState::new_initiator(initiator_keypair, responder_keypair.public_key());
let msg1 = initiator.write_message_1().unwrap();
// After processing message 1: responder knows initiator's identity
responder.read_message_1(&msg1).unwrap();
let discovered_initiator = responder.remote_static().unwrap();
assert_eq!(discovered_initiator, &initiator_keypair.public_key());
// The discovered key can be used to look up peer config, verify against allow-list, etc.
}
// ===== ReplayWindow Tests =====
#[test]
fn test_replay_window_basic() {
let mut window = ReplayWindow::new();
// First packet is always acceptable
assert!(window.check(0));
window.accept(0);
assert_eq!(window.highest(), 0);
// Replay of 0 should fail
assert!(!window.check(0));
// New higher counter is acceptable
assert!(window.check(1));
window.accept(1);
assert_eq!(window.highest(), 1);
// Out-of-order within window is acceptable
// (after accepting 10, 2 is still in window)
window.accept(10);
assert!(window.check(5));
window.accept(5);
// Replay of 5 should now fail
assert!(!window.check(5));
}
#[test]
fn test_replay_window_large_jump() {
let mut window = ReplayWindow::new();
// Accept counter 0
window.accept(0);
// Jump to a large counter
window.accept(REPLAY_WINDOW_SIZE as u64 + 100);
// Old counter should be outside window
assert!(!window.check(0));
assert!(!window.check(50));
// Counters within window should work
assert!(window.check(REPLAY_WINDOW_SIZE as u64 + 99));
assert!(window.check(REPLAY_WINDOW_SIZE as u64 + 50));
}
#[test]
fn test_replay_window_boundary() {
let mut window = ReplayWindow::new();
// Accept at boundary
window.accept(REPLAY_WINDOW_SIZE as u64 - 1);
// Counter 0 should be exactly at the edge of the window
assert!(window.check(0));
window.accept(0);
// Move window forward by 1
window.accept(REPLAY_WINDOW_SIZE as u64);
// Counter 0 is now outside the window
assert!(!window.check(0));
// Counter 1 is still in the window
assert!(window.check(1));
}
#[test]
fn test_replay_window_sequential() {
let mut window = ReplayWindow::new();
// Accept counters 0-999 in order
for i in 0..1000 {
assert!(window.check(i), "Counter {} should be acceptable", i);
window.accept(i);
}
// All should be marked as seen
for i in 0..1000 {
assert!(!window.check(i), "Counter {} should be rejected as replay", i);
}
assert_eq!(window.highest(), 999);
}
#[test]
fn test_replay_window_reset() {
let mut window = ReplayWindow::new();
window.accept(100);
assert_eq!(window.highest(), 100);
assert!(!window.check(100));
window.reset();
assert_eq!(window.highest(), 0);
assert!(window.check(100));
}
#[test]
fn test_session_replay_protection() {
let keypair1 = generate_keypair();
let keypair2 = generate_keypair();
let mut init = HandshakeState::new_initiator(keypair1, keypair2.public_key());
let mut resp = HandshakeState::new_responder(keypair2);
let msg1 = init.write_message_1().unwrap();
resp.read_message_1(&msg1).unwrap();
let msg2 = resp.write_message_2().unwrap();
init.read_message_2(&msg2).unwrap();
let mut sender = init.into_session().unwrap();
let mut receiver = resp.into_session().unwrap();
// Encrypt a message
let counter = sender.current_send_counter();
let ciphertext = sender.encrypt(b"test message").unwrap();
// First decryption should succeed
let plaintext = receiver
.decrypt_with_replay_check(&ciphertext, counter)
.unwrap();
assert_eq!(plaintext, b"test message");
// Replay should fail
let result = receiver.decrypt_with_replay_check(&ciphertext, counter);
assert!(matches!(result, Err(NoiseError::ReplayDetected(_))));
// Check method alone also detects replay
assert!(receiver.check_replay(counter).is_err());
}
#[test]
fn test_handshake_with_odd_parity_responder() {
// Node B's secret key produces an odd-parity public key (0x03 prefix).
// When the initiator only has the npub (x-only), PeerIdentity::pubkey_full()
// returns even parity (0x02). The pre-message mix_hash must normalize
// parity so both sides produce matching hash chains.
let secp = secp256k1::Secp256k1::new();
// Node B (responder) - odd parity key
let sk_b = secp256k1::SecretKey::from_slice(
&hex::decode("b102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1fb0")
.unwrap(),
)
.unwrap();
let kp_b = secp256k1::Keypair::from_secret_key(&secp, &sk_b);
let (xonly_b, parity_b) = kp_b.public_key().x_only_public_key();
assert_eq!(parity_b, Parity::Odd, "Test requires odd-parity responder key");
// Node A (initiator) - even parity key
let sk_a = secp256k1::SecretKey::from_slice(
&hex::decode("0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20")
.unwrap(),
)
.unwrap();
let kp_a = secp256k1::Keypair::from_secret_key(&secp, &sk_a);
// Simulate the production path: initiator gets responder's key via npub
// (x-only -> assumed even parity)
let assumed_even_b = xonly_b.public_key(Parity::Even);
assert_ne!(
assumed_even_b, kp_b.public_key(),
"Even assumption should differ from actual odd key"
);
// Handshake using assumed-even key (as production code does)
let mut initiator = HandshakeState::new_initiator(kp_a, assumed_even_b);
let mut responder = HandshakeState::new_responder(kp_b);
let msg1 = initiator.write_message_1().unwrap();
responder.read_message_1(&msg1).unwrap();
let msg2 = responder.write_message_2().unwrap();
initiator.read_message_2(&msg2).unwrap();
assert!(initiator.is_complete());
assert!(responder.is_complete());
// Verify sessions can communicate
let mut sender = initiator.into_session().unwrap();
let mut receiver = responder.into_session().unwrap();
let counter = sender.current_send_counter();
let ciphertext = sender.encrypt(b"parity test").unwrap();
let plaintext = receiver
.decrypt_with_replay_check(&ciphertext, counter)
.unwrap();
assert_eq!(plaintext, b"parity test");
}
src/peer/active.rs
+1 -1
View File
@@ -581,7 +581,7 @@ mod tests {
#[test]
fn test_active_peer_creation() {
let identity = make_peer_identity();
let peer = ActivePeer::new(identity.clone(), LinkId::new(1), 1000);
let peer = ActivePeer::new(identity, LinkId::new(1), 1000);
assert_eq!(peer.identity().node_addr(), identity.node_addr());
assert_eq!(peer.link_id(), LinkId::new(1));
src/peer/connection.rs
+5 -6
View File
@@ -384,10 +384,9 @@ impl PeerConnection {
hs.read_message_1(message)?;
// Extract the discovered identity
let remote_static = hs
let remote_static = *hs
.remote_static()
.expect("remote static available after msg1")
.clone();
.expect("remote static available after msg1");
self.expected_identity = Some(PeerIdentity::from_pubkey_full(remote_static));
// Generate message 2
@@ -519,7 +518,7 @@ mod tests {
#[test]
fn test_outbound_connection() {
let identity = make_peer_identity();
let conn = PeerConnection::outbound(LinkId::new(1), identity.clone(), 1000);
let conn = PeerConnection::outbound(LinkId::new(1), identity, 1000);
assert!(conn.is_outbound());
assert!(!conn.is_inbound());
@@ -617,9 +616,9 @@ mod tests {
let keypair = make_keypair();
// Outbound can't receive_handshake_init
let mut outbound = PeerConnection::outbound(LinkId::new(1), identity.clone(), 1000);
let mut outbound = PeerConnection::outbound(LinkId::new(1), identity, 1000);
assert!(outbound
.receive_handshake_init(keypair.clone(), &[0u8; 82], 1100)
.receive_handshake_init(keypair, &[0u8; 82], 1100)
.is_err());
// Inbound can't start_handshake
src/peer/mod.rs
+10 -10
View File
@@ -153,25 +153,25 @@ pub fn cross_connection_winner(
#[derive(Debug)]
pub enum PeerSlot {
/// Connection in handshake phase.
Connecting(PeerConnection),
Connecting(Box<PeerConnection>),
/// Authenticated peer.
Active(ActivePeer),
Active(Box<ActivePeer>),
}
impl PeerSlot {
/// Create a new connecting slot (outbound).
pub fn outbound(conn: PeerConnection) -> Self {
PeerSlot::Connecting(conn)
PeerSlot::Connecting(Box::new(conn))
}
/// Create a new connecting slot (inbound).
pub fn inbound(conn: PeerConnection) -> Self {
PeerSlot::Connecting(conn)
PeerSlot::Connecting(Box::new(conn))
}
/// Create a new active slot.
pub fn active(peer: ActivePeer) -> Self {
PeerSlot::Active(peer)
PeerSlot::Active(Box::new(peer))
}
/// Check if this is a connecting slot.
@@ -307,7 +307,7 @@ mod tests {
fn test_peer_slot_connecting() {
let identity = make_peer_identity();
let conn = PeerConnection::outbound(LinkId::new(1), identity, 1000);
let slot = PeerSlot::Connecting(conn);
let slot = PeerSlot::Connecting(Box::new(conn));
assert!(slot.is_connecting());
assert!(!slot.is_active());
@@ -320,7 +320,7 @@ mod tests {
fn test_peer_slot_active() {
let identity = make_peer_identity();
let peer = ActivePeer::new(identity, LinkId::new(2), 2000);
let slot = PeerSlot::Active(peer);
let slot = PeerSlot::Active(Box::new(peer));
assert!(!slot.is_connecting());
assert!(slot.is_active());
@@ -373,19 +373,19 @@ mod tests {
let identity = make_peer_identity();
let expected_node_addr = *identity.node_addr();
let conn = PeerConnection::outbound(LinkId::new(1), identity, 1000);
let slot = PeerSlot::Connecting(conn);
let slot = PeerSlot::Connecting(Box::new(conn));
assert_eq!(slot.node_addr(), Some(&expected_node_addr));
// Inbound connection doesn't know identity yet
let conn_inbound = PeerConnection::inbound(LinkId::new(2), 2000);
let slot_inbound = PeerSlot::Connecting(conn_inbound);
let slot_inbound = PeerSlot::Connecting(Box::new(conn_inbound));
assert!(slot_inbound.node_addr().is_none());
// Active peer always knows identity
let identity2 = make_peer_identity();
let active_node_addr = *identity2.node_addr();
let peer = ActivePeer::new(identity2, LinkId::new(3), 3000);
let slot_active = PeerSlot::Active(peer);
let slot_active = PeerSlot::Active(Box::new(peer));
assert_eq!(slot_active.node_addr(), Some(&active_node_addr));
}
}
src/tree.rs
-1479
View File
File diff suppressed because it is too large Load Diff
src/tree/coordinate.rs
+230
View File
@@ -0,0 +1,230 @@
//! Tree coordinates and distance calculations.
use std::fmt;
use super::TreeError;
use crate::NodeAddr;
/// Metadata for a single node in a tree coordinate path.
///
/// Carries the node address and its declaration metadata (sequence number
/// and timestamp). Used in TreeCoordinate entries and TreeAnnounce wire
/// format.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CoordEntry {
/// The node's routing address.
pub node_addr: NodeAddr,
/// The node's declaration sequence number.
pub sequence: u64,
/// The node's declaration timestamp (Unix seconds).
pub timestamp: u64,
}
impl CoordEntry {
/// Wire size of a serialized entry: node_addr(16) + sequence(8) + timestamp(8).
pub const WIRE_SIZE: usize = 32;
/// Create a new coordinate entry.
pub fn new(node_addr: NodeAddr, sequence: u64, timestamp: u64) -> Self {
Self {
node_addr,
sequence,
timestamp,
}
}
/// Create an entry with default metadata (sequence=0, timestamp=0).
///
/// Useful for constructing coordinates when only routing (not wire format)
/// is needed, e.g., in tests or distance calculations.
pub fn addr_only(node_addr: NodeAddr) -> Self {
Self {
node_addr,
sequence: 0,
timestamp: 0,
}
}
}
/// A node's coordinates in the spanning tree.
///
/// Coordinates are the path from the node to the root:
/// `[self, parent, grandparent, ..., root]`
///
/// Each entry carries the node address plus declaration metadata (sequence
/// and timestamp) for the wire protocol. Routing operations (distance,
/// LCA) use only the node addresses.
///
/// The coordinate enables greedy routing via tree distance calculation.
/// Two nodes can compute the hops between them by finding their lowest
/// common ancestor (LCA) in the tree.
#[derive(Clone, PartialEq, Eq)]
pub struct TreeCoordinate(Vec<CoordEntry>);
impl TreeCoordinate {
/// Create a coordinate from a path of entries (self to root).
///
/// The path must be non-empty and ordered from the node to the root.
pub fn new(path: Vec<CoordEntry>) -> Result<Self, TreeError> {
if path.is_empty() {
return Err(TreeError::EmptyCoordinate);
}
Ok(Self(path))
}
/// Create a coordinate from node addresses only (no metadata).
///
/// Convenience constructor for cases where only routing is needed.
/// Each entry gets sequence=0, timestamp=0.
pub fn from_addrs(addrs: Vec<NodeAddr>) -> Result<Self, TreeError> {
if addrs.is_empty() {
return Err(TreeError::EmptyCoordinate);
}
Ok(Self(
addrs
.into_iter()
.map(CoordEntry::addr_only)
.collect(),
))
}
/// Create a coordinate for a root node.
pub fn root(node_addr: NodeAddr) -> Self {
Self(vec![CoordEntry::addr_only(node_addr)])
}
/// Create a root coordinate with metadata.
pub fn root_with_meta(node_addr: NodeAddr, sequence: u64, timestamp: u64) -> Self {
Self(vec![CoordEntry::new(node_addr, sequence, timestamp)])
}
/// The node this coordinate belongs to (first element).
pub fn node_addr(&self) -> &NodeAddr {
&self.0[0].node_addr
}
/// The root of the tree (last element).
pub fn root_id(&self) -> &NodeAddr {
&self.0.last().expect("coordinate never empty").node_addr
}
/// The immediate parent (second element, or self if root).
pub fn parent_id(&self) -> &NodeAddr {
self.0
.get(1)
.map(|e| &e.node_addr)
.unwrap_or(&self.0[0].node_addr)
}
/// Depth in the tree (0 = root).
pub fn depth(&self) -> usize {
self.0.len() - 1
}
/// The full path of entries with metadata.
pub fn entries(&self) -> &[CoordEntry] {
&self.0
}
/// Iterator over node addresses in the path (self to root).
///
/// Use this for routing operations (distance, LCA, ancestor checks)
/// that only need the address path.
pub fn node_addrs(&self) -> impl DoubleEndedIterator<Item = &NodeAddr> {
self.0.iter().map(|e| &e.node_addr)
}
/// Check if this coordinate is a root (length 1).
pub fn is_root(&self) -> bool {
self.0.len() == 1
}
/// Calculate tree distance to another coordinate.
///
/// Distance is hops through the lowest common ancestor (LCA).
/// If the coordinates have different roots, returns usize::MAX.
pub fn distance_to(&self, other: &TreeCoordinate) -> usize {
// Different trees have infinite distance
if self.root_id() != other.root_id() {
return usize::MAX;
}
let lca_depth = self.lca_depth(other);
let self_to_lca = self.depth() - lca_depth;
let other_to_lca = other.depth() - lca_depth;
self_to_lca + other_to_lca
}
/// Find the depth of the lowest common ancestor.
///
/// Since coordinates are self-to-root, common ancestry is a suffix match.
/// Returns the depth (from root) of the LCA.
pub fn lca_depth(&self, other: &TreeCoordinate) -> usize {
let mut common: usize = 0;
let self_rev = self.node_addrs().rev();
let other_rev = other.node_addrs().rev();
for (a, b) in self_rev.zip(other_rev) {
if a == b {
common += 1;
} else {
break;
}
}
// LCA depth is counted from root (depth 0)
common.saturating_sub(1)
}
/// Get the lowest common ancestor node ID.
pub fn lca(&self, other: &TreeCoordinate) -> Option<&NodeAddr> {
let self_rev: Vec<_> = self.node_addrs().rev().collect();
let other_rev: Vec<_> = other.node_addrs().rev().collect();
let mut lca = None;
for (a, b) in self_rev.iter().zip(other_rev.iter()) {
if a == b {
lca = Some(*a);
} else {
break;
}
}
lca
}
/// Check if `other` is an ancestor (appears in our path after self).
pub fn has_ancestor(&self, other: &NodeAddr) -> bool {
self.node_addrs().skip(1).any(|id| id == other)
}
/// Check if `other` is in our ancestry (including self).
pub fn contains(&self, other: &NodeAddr) -> bool {
self.node_addrs().any(|id| id == other)
}
/// Get the ancestor at a specific depth from self.
///
/// `ancestor_at(0)` returns self, `ancestor_at(1)` returns parent, etc.
pub fn ancestor_at(&self, depth: usize) -> Option<&NodeAddr> {
self.0.get(depth).map(|e| &e.node_addr)
}
}
impl fmt::Debug for TreeCoordinate {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "TreeCoordinate(depth={}, path=[", self.depth())?;
for (i, entry) in self.0.iter().enumerate() {
if i > 0 {
write!(f, " → ")?;
}
// Show first 4 bytes of each node ID
write!(
f,
"{:02x}{:02x}",
entry.node_addr.as_bytes()[0],
entry.node_addr.as_bytes()[1]
)?;
}
write!(f, "])")
}
}
src/tree/declaration.rs
+182
View File
@@ -0,0 +1,182 @@
//! Parent declarations for the spanning tree.
use secp256k1::schnorr::Signature;
use secp256k1::XOnlyPublicKey;
use std::fmt;
use super::TreeError;
use crate::{Identity, NodeAddr};
/// A node's declaration of its parent in the spanning tree.
///
/// Each node periodically announces its parent selection. The declaration
/// includes a monotonic sequence number for freshness and a signature
/// for authenticity. When `parent_id == node_addr`, the node declares itself
/// as a root candidate.
#[derive(Clone)]
pub struct ParentDeclaration {
/// The node making this declaration.
node_addr: NodeAddr,
/// The selected parent (equals node_addr if self-declaring as root).
parent_id: NodeAddr,
/// Monotonically increasing sequence number.
sequence: u64,
/// Timestamp when this declaration was created (Unix seconds).
timestamp: u64,
/// Schnorr signature over the declaration fields.
signature: Option<Signature>,
}
impl ParentDeclaration {
/// Create a new unsigned parent declaration.
///
/// The declaration must be signed before transmission using `set_signature()`.
pub fn new(node_addr: NodeAddr, parent_id: NodeAddr, sequence: u64, timestamp: u64) -> Self {
Self {
node_addr,
parent_id,
sequence,
timestamp,
signature: None,
}
}
/// Create a self-declaration (node is root candidate).
pub fn self_root(node_addr: NodeAddr, sequence: u64, timestamp: u64) -> Self {
Self::new(node_addr, node_addr, sequence, timestamp)
}
/// Create a declaration with a pre-computed signature.
pub fn with_signature(
node_addr: NodeAddr,
parent_id: NodeAddr,
sequence: u64,
timestamp: u64,
signature: Signature,
) -> Self {
Self {
node_addr,
parent_id,
sequence,
timestamp,
signature: Some(signature),
}
}
/// Get the declaring node's ID.
pub fn node_addr(&self) -> &NodeAddr {
&self.node_addr
}
/// Get the parent node's ID.
pub fn parent_id(&self) -> &NodeAddr {
&self.parent_id
}
/// Get the sequence number.
pub fn sequence(&self) -> u64 {
self.sequence
}
/// Get the timestamp.
pub fn timestamp(&self) -> u64 {
self.timestamp
}
/// Get the signature, if set.
pub fn signature(&self) -> Option<&Signature> {
self.signature.as_ref()
}
/// Set the signature after signing.
pub fn set_signature(&mut self, signature: Signature) {
self.signature = Some(signature);
}
/// Sign this declaration with the given identity.
///
/// The identity's node_addr must match this declaration's node_addr.
/// Returns an error if the node_addrs don't match.
pub fn sign(&mut self, identity: &Identity) -> Result<(), TreeError> {
if identity.node_addr() != &self.node_addr {
return Err(TreeError::InvalidSignature(self.node_addr));
}
let signature = identity.sign(&self.signing_bytes());
self.signature = Some(signature);
Ok(())
}
/// Check if this is a root declaration (parent == self).
pub fn is_root(&self) -> bool {
self.node_addr == self.parent_id
}
/// Check if this declaration is signed.
pub fn is_signed(&self) -> bool {
self.signature.is_some()
}
/// Get the bytes that should be signed.
///
/// Format: node_addr (16) || parent_id (16) || sequence (8) || timestamp (8)
pub fn signing_bytes(&self) -> Vec<u8> {
let mut bytes = Vec::with_capacity(48);
bytes.extend_from_slice(self.node_addr.as_bytes());
bytes.extend_from_slice(self.parent_id.as_bytes());
bytes.extend_from_slice(&self.sequence.to_le_bytes());
bytes.extend_from_slice(&self.timestamp.to_le_bytes());
bytes
}
/// Verify the signature on this declaration.
///
/// Returns Ok(()) if the signature is valid, or an error otherwise.
pub fn verify(&self, pubkey: &XOnlyPublicKey) -> Result<(), TreeError> {
let signature = self
.signature
.as_ref()
.ok_or(TreeError::InvalidSignature(self.node_addr))?;
let secp = secp256k1::Secp256k1::verification_only();
let hash = self.signing_hash();
secp.verify_schnorr(signature, &hash, pubkey)
.map_err(|_| TreeError::InvalidSignature(self.node_addr))
}
/// Compute the SHA-256 hash of the signing bytes.
fn signing_hash(&self) -> [u8; 32] {
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
hasher.update(self.signing_bytes());
hasher.finalize().into()
}
/// Check if this declaration is fresher than another.
pub fn is_fresher_than(&self, other: &ParentDeclaration) -> bool {
self.sequence > other.sequence
}
}
impl fmt::Debug for ParentDeclaration {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ParentDeclaration")
.field("node_addr", &self.node_addr)
.field("parent_id", &self.parent_id)
.field("sequence", &self.sequence)
.field("is_root", &self.is_root())
.field("signed", &self.is_signed())
.finish()
}
}
impl PartialEq for ParentDeclaration {
fn eq(&self, other: &Self) -> bool {
self.node_addr == other.node_addr
&& self.parent_id == other.parent_id
&& self.sequence == other.sequence
&& self.timestamp == other.timestamp
}
}
impl Eq for ParentDeclaration {}
src/tree/mod.rs
+42
View File
@@ -0,0 +1,42 @@
//! Spanning Tree Protocol Entities
//!
//! Tree coordinates and parent declarations for the FIPS spanning tree.
//! The spanning tree provides a routing topology where each node maintains
//! a path to a common root, enabling greedy distance-based routing.
mod coordinate;
mod declaration;
mod state;
use thiserror::Error;
use crate::{IdentityError, NodeAddr};
pub use coordinate::{CoordEntry, TreeCoordinate};
pub use declaration::ParentDeclaration;
pub use state::TreeState;
/// Errors related to spanning tree operations.
#[derive(Debug, Error)]
pub enum TreeError {
#[error("invalid tree coordinate: empty path")]
EmptyCoordinate,
#[error("invalid ancestry: does not reach claimed root")]
AncestryNotToRoot,
#[error("signature verification failed for node {0:?}")]
InvalidSignature(NodeAddr),
#[error("sequence number regression: got {got}, expected > {expected}")]
SequenceRegression { got: u64, expected: u64 },
#[error("parent not in peers: {0:?}")]
ParentNotPeer(NodeAddr),
#[error("identity error: {0}")]
Identity(#[from] IdentityError),
}
#[cfg(test)]
mod tests;
src/tree/state.rs
+357
View File
@@ -0,0 +1,357 @@
//! Local spanning tree state for a node.
use std::collections::HashMap;
use std::fmt;
use super::{CoordEntry, ParentDeclaration, TreeCoordinate, TreeError};
use crate::{Identity, NodeAddr};
/// Local spanning tree state for a node.
///
/// Contains this node's declaration, coordinates, and view of peers'
/// tree positions. State is bounded by O(P × D) where P is peer count
/// and D is tree depth.
pub struct TreeState {
/// This node's NodeAddr.
my_node_addr: NodeAddr,
/// This node's current parent declaration.
my_declaration: ParentDeclaration,
/// This node's current coordinates (computed from declaration chain).
pub(super) my_coords: TreeCoordinate,
/// The current elected root (smallest reachable node_addr).
pub(super) root: NodeAddr,
/// Each peer's most recent parent declaration.
peer_declarations: HashMap<NodeAddr, ParentDeclaration>,
/// Each peer's full ancestry to root.
peer_ancestry: HashMap<NodeAddr, TreeCoordinate>,
/// Minimum depth improvement required to switch parents (same root).
parent_switch_threshold: usize,
}
impl TreeState {
/// Create initial tree state for a node (as root candidate).
///
/// The node starts as its own root until it learns of a smaller node_addr.
/// Initial sequence is 1 per protocol spec; timestamp is current Unix time.
pub fn new(my_node_addr: NodeAddr) -> Self {
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
let my_declaration = ParentDeclaration::self_root(my_node_addr, 1, timestamp);
let my_coords = TreeCoordinate::root_with_meta(my_node_addr, 1, timestamp);
Self {
my_node_addr,
my_declaration,
my_coords,
root: my_node_addr,
peer_declarations: HashMap::new(),
peer_ancestry: HashMap::new(),
parent_switch_threshold: 1,
}
}
/// Get this node's NodeAddr.
pub fn my_node_addr(&self) -> &NodeAddr {
&self.my_node_addr
}
/// Get this node's current declaration.
pub fn my_declaration(&self) -> &ParentDeclaration {
&self.my_declaration
}
/// Get this node's current coordinates.
pub fn my_coords(&self) -> &TreeCoordinate {
&self.my_coords
}
/// Get the current root.
pub fn root(&self) -> &NodeAddr {
&self.root
}
/// Check if this node is currently the root.
pub fn is_root(&self) -> bool {
self.root == self.my_node_addr
}
/// Get coordinates for a peer, if known.
pub fn peer_coords(&self, peer_id: &NodeAddr) -> Option<&TreeCoordinate> {
self.peer_ancestry.get(peer_id)
}
/// Get declaration for a peer, if known.
pub fn peer_declaration(&self, peer_id: &NodeAddr) -> Option<&ParentDeclaration> {
self.peer_declarations.get(peer_id)
}
/// Number of known peers.
pub fn peer_count(&self) -> usize {
self.peer_declarations.len()
}
/// Iterate over all peer node IDs.
pub fn peer_ids(&self) -> impl Iterator<Item = &NodeAddr> {
self.peer_declarations.keys()
}
/// Add or update a peer's tree state.
///
/// Returns true if the state was updated (new or fresher declaration).
pub fn update_peer(
&mut self,
declaration: ParentDeclaration,
ancestry: TreeCoordinate,
) -> bool {
let peer_id = *declaration.node_addr();
// Check if this is a fresh update
if let Some(existing) = self.peer_declarations.get(&peer_id)
&& !declaration.is_fresher_than(existing)
{
return false;
}
self.peer_declarations.insert(peer_id, declaration);
self.peer_ancestry.insert(peer_id, ancestry);
true
}
/// Remove a peer from the tree state.
pub fn remove_peer(&mut self, peer_id: &NodeAddr) {
self.peer_declarations.remove(peer_id);
self.peer_ancestry.remove(peer_id);
}
/// Update this node's parent selection.
///
/// Call this when switching parents. Updates the declaration and coordinates.
pub fn set_parent(&mut self, parent_id: NodeAddr, sequence: u64, timestamp: u64) {
self.my_declaration = ParentDeclaration::new(self.my_node_addr, parent_id, sequence, timestamp);
// Coordinates will be recomputed when ancestry is available
}
/// Update this node's coordinates based on current parent's ancestry.
pub fn recompute_coords(&mut self) {
if self.my_declaration.is_root() {
self.my_coords = TreeCoordinate::root_with_meta(
self.my_node_addr,
self.my_declaration.sequence(),
self.my_declaration.timestamp(),
);
self.root = self.my_node_addr;
return;
}
let parent_id = self.my_declaration.parent_id();
if let Some(parent_coords) = self.peer_ancestry.get(parent_id) {
// Our coords = [self_entry] ++ parent_coords entries
let self_entry = CoordEntry::new(
self.my_node_addr,
self.my_declaration.sequence(),
self.my_declaration.timestamp(),
);
let mut entries = vec![self_entry];
entries.extend_from_slice(parent_coords.entries());
self.my_coords = TreeCoordinate::new(entries).expect("non-empty path");
self.root = *self.my_coords.root_id();
}
}
/// Calculate tree distance to a peer.
pub fn distance_to_peer(&self, peer_id: &NodeAddr) -> Option<usize> {
self.peer_ancestry
.get(peer_id)
.map(|coords| self.my_coords.distance_to(coords))
}
/// Find the best next hop toward a destination using greedy tree routing.
///
/// Returns the peer that minimizes tree distance to the destination,
/// but only if that peer is strictly closer than we are (prevents
/// routing loops at local minima). Tie-breaks equal distance by
/// smallest node_addr.
///
/// Returns `None` if:
/// - No peers have coordinates
/// - Destination is in a different tree (different root)
/// - No peer is closer to the destination than we are
pub fn find_next_hop(&self, dest_coords: &TreeCoordinate) -> Option<NodeAddr> {
if self.my_coords.root_id() != dest_coords.root_id() {
return None;
}
let my_distance = self.my_coords.distance_to(dest_coords);
let mut best: Option<(NodeAddr, usize)> = None;
for (peer_id, peer_coords) in &self.peer_ancestry {
let distance = peer_coords.distance_to(dest_coords);
let dominated = match &best {
None => true,
Some((best_id, best_dist)) => {
distance < *best_dist
|| (distance == *best_dist && peer_id < best_id)
}
};
if dominated {
best = Some((*peer_id, distance));
}
}
match best {
Some((peer_id, distance)) if distance < my_distance => Some(peer_id),
_ => None,
}
}
/// Set the parent switch threshold.
pub fn set_parent_switch_threshold(&mut self, threshold: usize) {
self.parent_switch_threshold = threshold;
}
/// Evaluate whether to switch parents based on current peer tree state.
///
/// v1 algorithm: depth-based, no latency/loss metrics.
///
/// Returns `Some(peer_node_addr)` if a parent switch is recommended,
/// or `None` if the current parent is adequate.
pub fn evaluate_parent(&self) -> Option<NodeAddr> {
if self.peer_ancestry.is_empty() {
return None;
}
// Find the smallest root visible across all peers
let mut smallest_root: Option<NodeAddr> = None;
for coords in self.peer_ancestry.values() {
let peer_root = coords.root_id();
smallest_root = Some(match smallest_root {
None => *peer_root,
Some(current) => {
if *peer_root < current {
*peer_root
} else {
current
}
}
});
}
let smallest_root = smallest_root?;
// If we are the smallest node in the network, stay root
if self.my_node_addr <= smallest_root && self.is_root() {
return None;
}
// Among peers that reach the smallest root, find the shallowest
let mut best_peer: Option<(NodeAddr, usize)> = None; // (peer_addr, depth)
for (peer_id, coords) in &self.peer_ancestry {
if *coords.root_id() != smallest_root {
continue;
}
let depth = coords.depth();
match &best_peer {
None => best_peer = Some((*peer_id, depth)),
Some((_, best_depth)) => {
if depth < *best_depth {
best_peer = Some((*peer_id, depth));
}
}
}
}
let (best_peer_id, best_depth) = best_peer?;
// If already using this peer as parent, no switch needed
if *self.my_declaration.parent_id() == best_peer_id && !self.is_root() {
return None;
}
// If our current parent is gone from peer_ancestry, our path is broken — always switch
if !self.is_root() && !self.peer_ancestry.contains_key(self.my_declaration.parent_id()) {
return Some(best_peer_id);
}
// Switching roots (smaller root found) → always switch
if smallest_root < self.root || (self.is_root() && smallest_root < self.my_node_addr) {
return Some(best_peer_id);
}
// Same root: require depth improvement ≥ threshold
if self.is_root() {
// We're root but shouldn't be (peers have a smaller root) — always switch
return Some(best_peer_id);
}
// Compare depth: our current depth vs what we'd get through best_peer
// Our new depth would be best_depth + 1
let current_depth = self.my_coords.depth();
let proposed_depth = best_depth + 1;
if current_depth >= proposed_depth + self.parent_switch_threshold {
return Some(best_peer_id);
}
None
}
/// Handle loss of current parent.
///
/// Tries to find an alternative parent among remaining peers.
/// If none available, becomes its own root (increments sequence).
///
/// Returns `true` if the tree state changed (caller should re-announce).
pub fn handle_parent_lost(&mut self) -> bool {
// Try to find an alternative parent
if let Some(new_parent) = self.evaluate_parent() {
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
let new_seq = self.my_declaration.sequence() + 1;
self.set_parent(new_parent, new_seq, timestamp);
self.recompute_coords();
return true;
}
// No alternative: become own root
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
let new_seq = self.my_declaration.sequence() + 1;
self.my_declaration =
ParentDeclaration::self_root(self.my_node_addr, new_seq, timestamp);
self.recompute_coords();
true
}
/// Sign this node's declaration with the given identity.
///
/// The identity's node_addr must match this TreeState's node_addr.
pub fn sign_declaration(&mut self, identity: &Identity) -> Result<(), TreeError> {
self.my_declaration.sign(identity)
}
/// Check if this node's declaration is signed.
pub fn is_declaration_signed(&self) -> bool {
self.my_declaration.is_signed()
}
}
impl fmt::Debug for TreeState {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("TreeState")
.field("my_node_addr", &self.my_node_addr)
.field("root", &self.root)
.field("is_root", &self.is_root())
.field("depth", &self.my_coords.depth())
.field("peers", &self.peer_count())
.finish()
}
}
src/tree/tests.rs
+687
View File
@@ -0,0 +1,687 @@
use super::*;
fn make_node_addr(val: u8) -> NodeAddr {
let mut bytes = [0u8; 16];
bytes[0] = val;
NodeAddr::from_bytes(bytes)
}
fn make_coords(ids: &[u8]) -> TreeCoordinate {
TreeCoordinate::from_addrs(ids.iter().map(|&v| make_node_addr(v)).collect()).unwrap()
}
// ===== TreeCoordinate Tests =====
#[test]
fn test_tree_coordinate_root() {
let root_id = make_node_addr(1);
let coord = TreeCoordinate::root(root_id);
assert!(coord.is_root());
assert_eq!(coord.depth(), 0);
assert_eq!(coord.node_addr(), &root_id);
assert_eq!(coord.root_id(), &root_id);
assert_eq!(coord.parent_id(), &root_id);
}
#[test]
fn test_tree_coordinate_path() {
let node = make_node_addr(1);
let parent = make_node_addr(2);
let root = make_node_addr(3);
let coord = make_coords(&[1, 2, 3]);
assert!(!coord.is_root());
assert_eq!(coord.depth(), 2);
assert_eq!(coord.node_addr(), &node);
assert_eq!(coord.parent_id(), &parent);
assert_eq!(coord.root_id(), &root);
}
#[test]
fn test_tree_coordinate_empty_fails() {
let result = TreeCoordinate::from_addrs(vec![]);
assert!(matches!(result, Err(TreeError::EmptyCoordinate)));
}
#[test]
fn test_tree_coordinate_entries_metadata() {
let node = make_node_addr(1);
let root = make_node_addr(0);
let coord = TreeCoordinate::new(vec![
CoordEntry::new(node, 5, 1000),
CoordEntry::new(root, 1, 500),
])
.unwrap();
assert_eq!(coord.entries()[0].sequence, 5);
assert_eq!(coord.entries()[0].timestamp, 1000);
assert_eq!(coord.entries()[1].sequence, 1);
assert_eq!(coord.entries()[1].timestamp, 500);
}
#[test]
fn test_tree_distance_same_node() {
let node = make_node_addr(1);
let coord = TreeCoordinate::root(node);
assert_eq!(coord.distance_to(&coord), 0);
}
#[test]
fn test_tree_distance_siblings() {
let coord_a = make_coords(&[1, 0]);
let coord_b = make_coords(&[2, 0]);
// a -> root -> b = 2 hops
assert_eq!(coord_a.distance_to(&coord_b), 2);
}
#[test]
fn test_tree_distance_ancestor() {
let coord_parent = make_coords(&[1, 0]);
let coord_child = make_coords(&[2, 1, 0]);
// child -> parent = 1 hop
assert_eq!(coord_child.distance_to(&coord_parent), 1);
}
#[test]
fn test_tree_distance_cousins() {
// Tree structure:
// root(0)
// / \
// a(1) b(2)
// / \
// c(3) d(4)
let coord_c = make_coords(&[3, 1, 0]);
let coord_d = make_coords(&[4, 2, 0]);
// c -> a -> root -> b -> d = 4 hops
assert_eq!(coord_c.distance_to(&coord_d), 4);
}
#[test]
fn test_tree_distance_different_roots() {
let coord1 = TreeCoordinate::root(make_node_addr(1));
let coord2 = TreeCoordinate::root(make_node_addr(2));
assert_eq!(coord1.distance_to(&coord2), usize::MAX);
}
#[test]
fn test_has_ancestor() {
let root = make_node_addr(0);
let parent = make_node_addr(1);
let child = make_node_addr(2);
let coord = make_coords(&[2, 1, 0]);
assert!(coord.has_ancestor(&parent));
assert!(coord.has_ancestor(&root));
assert!(!coord.has_ancestor(&child)); // self is not an ancestor
}
#[test]
fn test_contains() {
let root = make_node_addr(0);
let parent = make_node_addr(1);
let child = make_node_addr(2);
let other = make_node_addr(99);
let coord = make_coords(&[2, 1, 0]);
assert!(coord.contains(&child));
assert!(coord.contains(&parent));
assert!(coord.contains(&root));
assert!(!coord.contains(&other));
}
#[test]
fn test_ancestor_at() {
let root = make_node_addr(0);
let parent = make_node_addr(1);
let child = make_node_addr(2);
let coord = make_coords(&[2, 1, 0]);
assert_eq!(coord.ancestor_at(0), Some(&child));
assert_eq!(coord.ancestor_at(1), Some(&parent));
assert_eq!(coord.ancestor_at(2), Some(&root));
assert_eq!(coord.ancestor_at(3), None);
}
#[test]
fn test_lca() {
let root = make_node_addr(0);
let a = make_node_addr(1);
// c under a, d under b, both under root
let coord_c = make_coords(&[3, 1, 0]);
let coord_d = make_coords(&[4, 2, 0]);
assert_eq!(coord_c.lca(&coord_d), Some(&root));
// c and a share ancestry through a and root
let coord_a = make_coords(&[1, 0]);
assert_eq!(coord_c.lca(&coord_a), Some(&a));
}
// ===== ParentDeclaration Tests =====
#[test]
fn test_parent_declaration_new() {
let node = make_node_addr(1);
let parent = make_node_addr(2);
let decl = ParentDeclaration::new(node, parent, 1, 1000);
assert_eq!(decl.node_addr(), &node);
assert_eq!(decl.parent_id(), &parent);
assert_eq!(decl.sequence(), 1);
assert_eq!(decl.timestamp(), 1000);
assert!(!decl.is_root());
assert!(!decl.is_signed());
}
#[test]
fn test_parent_declaration_self_root() {
let node = make_node_addr(1);
let decl = ParentDeclaration::self_root(node, 5, 2000);
assert!(decl.is_root());
assert_eq!(decl.node_addr(), decl.parent_id());
}
#[test]
fn test_parent_declaration_freshness() {
let node = make_node_addr(1);
let parent = make_node_addr(2);
let old_decl = ParentDeclaration::new(node, parent, 1, 1000);
let new_decl = ParentDeclaration::new(node, parent, 2, 2000);
assert!(new_decl.is_fresher_than(&old_decl));
assert!(!old_decl.is_fresher_than(&new_decl));
assert!(!old_decl.is_fresher_than(&old_decl));
}
#[test]
fn test_parent_declaration_signing_bytes() {
let node = make_node_addr(1);
let parent = make_node_addr(2);
let decl = ParentDeclaration::new(node, parent, 100, 1234567890);
let bytes = decl.signing_bytes();
// Should be 48 bytes: 16 + 16 + 8 + 8
assert_eq!(bytes.len(), 48);
// Verify structure
assert_eq!(&bytes[0..16], node.as_bytes());
assert_eq!(&bytes[16..32], parent.as_bytes());
assert_eq!(&bytes[32..40], &100u64.to_le_bytes());
assert_eq!(&bytes[40..48], &1234567890u64.to_le_bytes());
}
#[test]
fn test_parent_declaration_equality() {
let node = make_node_addr(1);
let parent = make_node_addr(2);
let decl1 = ParentDeclaration::new(node, parent, 1, 1000);
let decl2 = ParentDeclaration::new(node, parent, 1, 1000);
let decl3 = ParentDeclaration::new(node, parent, 2, 1000);
assert_eq!(decl1, decl2);
assert_ne!(decl1, decl3);
}
// ===== TreeState Tests =====
#[test]
fn test_tree_state_new() {
let node = make_node_addr(1);
let state = TreeState::new(node);
assert_eq!(state.my_node_addr(), &node);
assert!(state.is_root());
assert_eq!(state.root(), &node);
assert_eq!(state.my_coords().depth(), 0);
assert_eq!(state.peer_count(), 0);
}
#[test]
fn test_tree_state_update_peer() {
let my_node = make_node_addr(0);
let mut state = TreeState::new(my_node);
let peer = make_node_addr(1);
let root = make_node_addr(2);
let decl = ParentDeclaration::new(peer, root, 1, 1000);
let coords = make_coords(&[1, 2]);
assert!(state.update_peer(decl.clone(), coords.clone()));
assert_eq!(state.peer_count(), 1);
assert!(state.peer_coords(&peer).is_some());
assert!(state.peer_declaration(&peer).is_some());
// Same sequence should not update
let decl2 = ParentDeclaration::new(peer, root, 1, 1000);
assert!(!state.update_peer(decl2, coords.clone()));
// Higher sequence should update
let decl3 = ParentDeclaration::new(peer, root, 2, 2000);
assert!(state.update_peer(decl3, coords));
}
#[test]
fn test_tree_state_remove_peer() {
let my_node = make_node_addr(0);
let mut state = TreeState::new(my_node);
let peer = make_node_addr(1);
let root = make_node_addr(2);
let decl = ParentDeclaration::new(peer, root, 1, 1000);
let coords = make_coords(&[1, 2]);
state.update_peer(decl, coords);
assert_eq!(state.peer_count(), 1);
state.remove_peer(&peer);
assert_eq!(state.peer_count(), 0);
assert!(state.peer_coords(&peer).is_none());
}
#[test]
fn test_tree_state_distance_to_peer() {
let my_node = make_node_addr(0);
let mut state = TreeState::new(my_node);
let peer = make_node_addr(1);
// Both are roots in their own trees initially - different roots
let peer_coords = TreeCoordinate::root(peer);
let decl = ParentDeclaration::self_root(peer, 1, 1000);
state.update_peer(decl, peer_coords);
// Different roots = MAX distance
assert_eq!(state.distance_to_peer(&peer), Some(usize::MAX));
// If they share a root, distance should be finite
let shared_root = make_node_addr(99);
// Update my state to have shared root
state.set_parent(shared_root, 1, 1000);
let my_new_coords = make_coords(&[0, 99]);
// Manually set coords for test (normally done by recompute_coords)
state.my_coords = my_new_coords;
state.root = shared_root;
// Update peer to have same root
let peer_coords = make_coords(&[1, 99]);
let decl = ParentDeclaration::new(peer, shared_root, 2, 2000);
state.update_peer(decl, peer_coords);
// Now distance should be 2 (me -> root -> peer)
assert_eq!(state.distance_to_peer(&peer), Some(2));
}
#[test]
fn test_tree_state_peer_ids() {
let my_node = make_node_addr(0);
let mut state = TreeState::new(my_node);
let peer1 = make_node_addr(1);
let peer2 = make_node_addr(2);
state.update_peer(
ParentDeclaration::self_root(peer1, 1, 1000),
TreeCoordinate::root(peer1),
);
state.update_peer(
ParentDeclaration::self_root(peer2, 1, 1000),
TreeCoordinate::root(peer2),
);
let ids: Vec<_> = state.peer_ids().collect();
assert_eq!(ids.len(), 2);
assert!(ids.contains(&&peer1));
assert!(ids.contains(&&peer2));
}
// ===== Parent Selection Tests =====
#[test]
fn test_evaluate_parent_picks_smallest_root() {
// Node 5 starts as root. Peers 3 and 7 each claim different roots.
// Peer 3's path: [3, 1] (root=1)
// Peer 7's path: [7, 2] (root=2)
// Should pick peer 3 because root 1 < root 2.
let my_node = make_node_addr(5);
let mut state = TreeState::new(my_node);
let peer3 = make_node_addr(3);
let peer7 = make_node_addr(7);
state.update_peer(
ParentDeclaration::new(peer3, make_node_addr(1), 1, 1000),
make_coords(&[3, 1]),
);
state.update_peer(
ParentDeclaration::new(peer7, make_node_addr(2), 1, 1000),
make_coords(&[7, 2]),
);
let result = state.evaluate_parent();
assert_eq!(result, Some(peer3));
}
#[test]
fn test_evaluate_parent_prefers_shallowest_depth() {
// Node 5, root=0 (shared). Peer 1 at depth 1, peer 2 at depth 3.
// Both reach root 0. Should pick peer 1 (shallowest).
let my_node = make_node_addr(5);
let mut state = TreeState::new(my_node);
let peer1 = make_node_addr(1);
let peer2 = make_node_addr(2);
let root = make_node_addr(0);
// Peer 1: depth 1 (path = [1, 0])
state.update_peer(
ParentDeclaration::new(peer1, root, 1, 1000),
make_coords(&[1, 0]),
);
// Peer 2: depth 3 (path = [2, 3, 4, 0])
state.update_peer(
ParentDeclaration::new(peer2, make_node_addr(3), 1, 1000),
make_coords(&[2, 3, 4, 0]),
);
let result = state.evaluate_parent();
assert_eq!(result, Some(peer1));
}
#[test]
fn test_evaluate_parent_stays_root_when_smallest() {
// Node 0 (smallest possible) should stay root even if peers exist.
let my_node = make_node_addr(0);
let mut state = TreeState::new(my_node);
let peer1 = make_node_addr(1);
// Peer 1 has root 0 (us) — shouldn't trigger switch
state.update_peer(
ParentDeclaration::new(peer1, my_node, 1, 1000),
make_coords(&[1, 0]),
);
assert_eq!(state.evaluate_parent(), None);
}
#[test]
fn test_evaluate_parent_no_switch_when_already_best() {
// Node 5, already using peer 1 as parent. No better option.
let my_node = make_node_addr(5);
let mut state = TreeState::new(my_node);
let peer1 = make_node_addr(1);
let root = make_node_addr(0);
state.update_peer(
ParentDeclaration::new(peer1, root, 1, 1000),
make_coords(&[1, 0]),
);
// Switch to peer1 as parent first
state.set_parent(peer1, 1, 1000);
state.recompute_coords();
// Now evaluate — should return None since peer1 is already our parent
assert_eq!(state.evaluate_parent(), None);
}
#[test]
fn test_evaluate_parent_no_peers() {
let my_node = make_node_addr(5);
let state = TreeState::new(my_node);
assert_eq!(state.evaluate_parent(), None);
}
#[test]
fn test_evaluate_parent_depth_threshold() {
// Node 5, currently at depth 4 through peer 2.
// Peer 1 offers depth 3 (improvement of 1, which equals threshold).
// Peer 3 offers depth 1 (improvement of 3, exceeds threshold).
// Should switch to peer 3.
let my_node = make_node_addr(5);
let mut state = TreeState::new(my_node);
let peer2 = make_node_addr(2);
let peer3 = make_node_addr(3);
let root = make_node_addr(0);
// Peer 2: depth 3 (we'd be depth 4 through them)
state.update_peer(
ParentDeclaration::new(peer2, make_node_addr(6), 1, 1000),
make_coords(&[2, 6, 7, 0]),
);
// Set peer2 as our parent, making us depth 4
state.set_parent(peer2, 1, 1000);
state.recompute_coords();
assert_eq!(state.my_coords().depth(), 4);
// Peer 3: depth 1 (we'd be depth 2 through them) — improvement of 2
state.update_peer(
ParentDeclaration::new(peer3, root, 1, 1000),
make_coords(&[3, 0]),
);
let result = state.evaluate_parent();
assert_eq!(result, Some(peer3));
}
#[test]
fn test_handle_parent_lost_finds_alternative() {
let my_node = make_node_addr(5);
let mut state = TreeState::new(my_node);
let peer1 = make_node_addr(1);
let peer2 = make_node_addr(2);
let root = make_node_addr(0);
state.update_peer(
ParentDeclaration::new(peer1, root, 1, 1000),
make_coords(&[1, 0]),
);
state.update_peer(
ParentDeclaration::new(peer2, root, 1, 1000),
make_coords(&[2, 0]),
);
// Set peer1 as parent
state.set_parent(peer1, 1, 1000);
state.recompute_coords();
// Remove peer1 (parent lost)
state.remove_peer(&peer1);
let changed = state.handle_parent_lost();
assert!(changed);
// Should have switched to peer2
assert_eq!(state.my_declaration().parent_id(), &peer2);
assert!(!state.is_root());
}
#[test]
fn test_handle_parent_lost_becomes_root() {
let my_node = make_node_addr(5);
let mut state = TreeState::new(my_node);
let peer1 = make_node_addr(1);
let root = make_node_addr(0);
state.update_peer(
ParentDeclaration::new(peer1, root, 1, 1000),
make_coords(&[1, 0]),
);
// Set peer1 as parent
state.set_parent(peer1, 1, 1000);
state.recompute_coords();
let seq_before = state.my_declaration().sequence();
// Remove peer1 (only parent)
state.remove_peer(&peer1);
let changed = state.handle_parent_lost();
assert!(changed);
assert!(state.is_root());
assert!(state.my_declaration().sequence() > seq_before);
assert_eq!(state.root(), &my_node);
}
// === find_next_hop tests ===
/// Build a TreeState with our own coordinates set.
fn make_tree_state(my_addr: u8, coord_path: &[u8]) -> TreeState {
let my_node = make_node_addr(my_addr);
let mut state = TreeState::new(my_node);
let coords = make_coords(coord_path);
state.root = *coords.root_id();
state.my_coords = coords;
state
}
/// Add a peer with given coordinates to the tree state.
fn add_peer(state: &mut TreeState, peer_addr: u8, coord_path: &[u8]) {
let peer = make_node_addr(peer_addr);
let parent = make_node_addr(coord_path[1]);
state.update_peer(
ParentDeclaration::new(peer, parent, 1, 1000),
make_coords(coord_path),
);
}
#[test]
fn test_find_next_hop_chain() {
// Chain: 0 (root) <- 5 (us) <- 1 <- 2
// Both peers 1 and 2 are in our peer_ancestry. Peer 2 IS the
// destination (distance 0), so it's the best next hop.
let mut state = make_tree_state(5, &[5, 0]);
add_peer(&mut state, 1, &[1, 5, 0]);
add_peer(&mut state, 2, &[2, 1, 5, 0]);
let dest = make_coords(&[2, 1, 5, 0]);
assert_eq!(state.find_next_hop(&dest), Some(make_node_addr(2)));
}
#[test]
fn test_find_next_hop_chain_indirect() {
// Chain: 0 (root) <- 5 (us) <- 1
// Dest is node 2 at [2, 1, 5, 0] but peer 2 is NOT in our peer
// list — only peer 1 is. So we route via peer 1 (distance 1).
let mut state = make_tree_state(5, &[5, 0]);
add_peer(&mut state, 1, &[1, 5, 0]);
let dest = make_coords(&[2, 1, 5, 0]);
assert_eq!(state.find_next_hop(&dest), Some(make_node_addr(1)));
}
#[test]
fn test_find_next_hop_toward_root() {
// Tree: 0 (root) <- 1 <- 5 (us)
// Routing toward root should pick node 1 (our parent).
let mut state = make_tree_state(5, &[5, 1, 0]);
add_peer(&mut state, 1, &[1, 0]);
let dest = make_coords(&[0]);
assert_eq!(state.find_next_hop(&dest), Some(make_node_addr(1)));
}
#[test]
fn test_find_next_hop_sibling() {
// Tree: 0 (root) <- 5 (us), 0 <- 3
// Routing to sibling 3: should go through parent 0... but 0 is
// the root and not in our peer list. Our only peer is 3 itself.
// But 3 is not a "closer" peer in tree distance — distance from
// us to 3 is 2 (up to root, down to 3), and distance from 3 to
// 3 is 0, so 3 IS closer. Should pick 3.
let mut state = make_tree_state(5, &[5, 0]);
add_peer(&mut state, 3, &[3, 0]);
let dest = make_coords(&[3, 0]);
assert_eq!(state.find_next_hop(&dest), Some(make_node_addr(3)));
}
#[test]
fn test_find_next_hop_tie_breaking() {
// Tree: 0 (root) <- 5 (us), 0 <- 3, 0 <- 2
// Both peers are siblings at depth 1, equidistant to a dest
// at [4, 0]. Should pick node 2 (smaller node_addr).
let mut state = make_tree_state(5, &[5, 0]);
add_peer(&mut state, 3, &[3, 0]);
add_peer(&mut state, 2, &[2, 0]);
let dest = make_coords(&[4, 0]);
// Our distance: 2 (up to root, down to 4)
// Peer 3 distance: 2 (up to root, down to 4)
// Peer 2 distance: 2 (up to root, down to 4)
// All equal to our distance — no peer is strictly closer.
assert_eq!(state.find_next_hop(&dest), None);
}
#[test]
fn test_find_next_hop_different_root() {
let mut state = make_tree_state(5, &[5, 0]);
add_peer(&mut state, 1, &[1, 0]);
// Destination in a different tree (root = 9)
let dest = make_coords(&[3, 9]);
assert_eq!(state.find_next_hop(&dest), None);
}
#[test]
fn test_find_next_hop_no_peers() {
let state = make_tree_state(5, &[5, 0]);
let dest = make_coords(&[3, 0]);
assert_eq!(state.find_next_hop(&dest), None);
}
#[test]
fn test_find_next_hop_local_minimum() {
// Tree: 0 (root) <- 5 (us), 5 <- 8
// Routing to node 3 at [3, 0]. Our distance = 2.
// Peer 8's distance = 4 (8→5→0→3 but via coords: [8,5,0] to [3,0] = 3).
// Actually: lca of [8,5,0] and [3,0] is root 0 at depth 0.
// dist = (2-0) + (1-0) = 3. Our dist = (1-0) + (1-0) = 2.
// Peer is farther, so no hop.
let mut state = make_tree_state(5, &[5, 0]);
add_peer(&mut state, 8, &[8, 5, 0]);
let dest = make_coords(&[3, 0]);
assert_eq!(state.find_next_hop(&dest), None);
}
#[test]
fn test_find_next_hop_best_of_multiple() {
// Tree: 0 (root) <- 1 <- 5 (us), 1 <- 3 <- 7
// Dest is node 7 at [7, 3, 1, 0].
// Peer 1 coords [1, 0]: dist to dest = 0 + 2 = 2
// Peer 3 coords [3, 1, 0]: dist to dest = 0 + 1 = 1
// Our coords [5, 1, 0]: dist to dest = 1 + 2 = 3
// Peer 3 is closest. Should pick 3.
let mut state = make_tree_state(5, &[5, 1, 0]);
add_peer(&mut state, 1, &[1, 0]);
add_peer(&mut state, 3, &[3, 1, 0]);
let dest = make_coords(&[7, 3, 1, 0]);
assert_eq!(state.find_next_hop(&dest), Some(make_node_addr(3)));
}