mirror of
https://github.com/jmcorgan/fips.git
synced 2026-07-22 07:48:26 +00:00
Discovery protocol: LookupRequest/LookupResponse handlers
Implement the coordinate discovery protocol with flood-based lookup and reverse-path response routing. Wire format: LookupRequest (0x30) encode/decode with TTL, visited bloom filter, origin coords. LookupResponse (0x31) encode/decode with target coords and Schnorr proof signature. Handler logic: request dedup by request_id, visited filter loop prevention, TTL enforcement, lazy purge of expired entries (10s). Response routing: originator caches route in route_cache, transit nodes reverse-path forward via recent_requests. Node state: RecentRequest struct, route_cache (RouteCache), and recent_requests map for dedup + reverse-path forwarding. 13 handler tests (9 unit + 4 integration) plus 4 protocol tests. 392 tests pass, clean build.
This commit is contained in:
299
src/node/handlers/discovery.rs
Normal file
299
src/node/handlers/discovery.rs
Normal file
@@ -0,0 +1,299 @@
|
||||
//! LookupRequest/LookupResponse discovery protocol handlers.
|
||||
//!
|
||||
//! Handles coordinate discovery requests: flood-based lookup with TTL,
|
||||
//! visited filter for loop prevention, and reverse-path forwarding for
|
||||
//! responses.
|
||||
|
||||
use crate::node::{Node, RecentRequest};
|
||||
use crate::protocol::{LookupRequest, LookupResponse};
|
||||
use crate::NodeAddr;
|
||||
use tracing::{debug, trace};
|
||||
|
||||
impl Node {
|
||||
/// Handle an incoming LookupRequest from a peer.
|
||||
///
|
||||
/// Processing steps:
|
||||
/// 1. Decode and validate
|
||||
/// 2. Check request_id for duplicates (dedup)
|
||||
/// 3. Record request for reverse-path forwarding
|
||||
/// 4. Lazy purge expired entries
|
||||
/// 5. Check visited filter (loop prevention)
|
||||
/// 6. If we're the target, generate and send response
|
||||
/// 7. If TTL > 0, forward to peers not in visited filter
|
||||
pub(in crate::node) async fn handle_lookup_request(
|
||||
&mut self,
|
||||
from: &NodeAddr,
|
||||
payload: &[u8],
|
||||
) {
|
||||
let request = match LookupRequest::decode(payload) {
|
||||
Ok(req) => req,
|
||||
Err(e) => {
|
||||
debug!(from = %from, error = %e, "Malformed LookupRequest");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let now_ms = Self::now_ms();
|
||||
|
||||
// Dedup: drop if we've already seen this request_id
|
||||
if self.recent_requests.contains_key(&request.request_id) {
|
||||
trace!(
|
||||
request_id = request.request_id,
|
||||
from = %from,
|
||||
"Duplicate LookupRequest, dropping"
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
// Record for reverse-path forwarding and dedup
|
||||
self.recent_requests.insert(
|
||||
request.request_id,
|
||||
RecentRequest::new(*from, now_ms),
|
||||
);
|
||||
|
||||
// Lazy purge expired entries
|
||||
self.purge_expired_requests(now_ms);
|
||||
|
||||
// Loop prevention: drop if we've already been visited
|
||||
if request.was_visited(self.node_addr()) {
|
||||
trace!(
|
||||
request_id = request.request_id,
|
||||
target = %request.target,
|
||||
"Already visited, dropping LookupRequest"
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
// Are we the target?
|
||||
if request.target == *self.node_addr() {
|
||||
debug!(
|
||||
request_id = request.request_id,
|
||||
origin = %request.origin,
|
||||
"We are the lookup target, generating response"
|
||||
);
|
||||
self.send_lookup_response(&request).await;
|
||||
return;
|
||||
}
|
||||
|
||||
// Forward if TTL permits
|
||||
if request.can_forward() {
|
||||
self.forward_lookup_request(request).await;
|
||||
} else {
|
||||
trace!(
|
||||
request_id = request.request_id,
|
||||
target = %request.target,
|
||||
"LookupRequest TTL exhausted, not forwarding"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle an incoming LookupResponse from a peer.
|
||||
///
|
||||
/// Processing steps:
|
||||
/// 1. Decode and validate
|
||||
/// 2. Check recent_requests to determine if we originated or are forwarding
|
||||
/// 3. If originator: cache target_coords in route_cache
|
||||
/// 4. If transit: reverse-path forward to from_peer
|
||||
pub(in crate::node) async fn handle_lookup_response(
|
||||
&mut self,
|
||||
from: &NodeAddr,
|
||||
payload: &[u8],
|
||||
) {
|
||||
let response = match LookupResponse::decode(payload) {
|
||||
Ok(resp) => resp,
|
||||
Err(e) => {
|
||||
debug!(from = %from, error = %e, "Malformed LookupResponse");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let now_ms = Self::now_ms();
|
||||
|
||||
// Check if we forwarded this request (transit node) or originated it
|
||||
if let Some(recent) = self.recent_requests.get(&response.request_id) {
|
||||
// Transit node: reverse-path forward
|
||||
let from_peer = recent.from_peer;
|
||||
|
||||
debug!(
|
||||
request_id = response.request_id,
|
||||
target = %response.target,
|
||||
next_hop = %from_peer,
|
||||
"Reverse-path forwarding LookupResponse"
|
||||
);
|
||||
|
||||
let encoded = response.encode();
|
||||
if let Err(e) = self.send_encrypted_link_message(&from_peer, &encoded).await {
|
||||
debug!(
|
||||
next_hop = %from_peer,
|
||||
error = %e,
|
||||
"Failed to forward LookupResponse"
|
||||
);
|
||||
}
|
||||
} else {
|
||||
// We originated this request — cache the discovered coordinates
|
||||
debug!(
|
||||
request_id = response.request_id,
|
||||
target = %response.target,
|
||||
depth = response.target_coords.depth(),
|
||||
"Received LookupResponse, caching route"
|
||||
);
|
||||
|
||||
self.route_cache.insert(
|
||||
response.target,
|
||||
response.target_coords,
|
||||
now_ms,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate and send a LookupResponse when we are the target.
|
||||
///
|
||||
/// Signs a proof using our identity and routes the response toward
|
||||
/// the origin. The first hop uses find_next_hop; subsequent hops use
|
||||
/// reverse-path forwarding via recent_requests.
|
||||
async fn send_lookup_response(&mut self, request: &LookupRequest) {
|
||||
let our_coords = self.tree_state().my_coords().clone();
|
||||
|
||||
// Sign proof: Identity::sign hashes with SHA-256 internally
|
||||
let proof_data = LookupResponse::proof_bytes(request.request_id, &request.target);
|
||||
let proof = self.identity().sign(&proof_data);
|
||||
|
||||
let response = LookupResponse::new(
|
||||
request.request_id,
|
||||
request.target,
|
||||
our_coords,
|
||||
proof,
|
||||
);
|
||||
|
||||
// Route toward origin
|
||||
let next_hop_addr = match self.find_next_hop(&request.origin) {
|
||||
Some(peer) => *peer.node_addr(),
|
||||
None => {
|
||||
// Origin might be our direct peer that sent us the request
|
||||
// Check if origin == the peer we received from
|
||||
if let Some(recent) = self.recent_requests.get(&request.request_id) {
|
||||
recent.from_peer
|
||||
} else {
|
||||
debug!(
|
||||
origin = %request.origin,
|
||||
"Cannot route LookupResponse: no path to origin"
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
debug!(
|
||||
request_id = request.request_id,
|
||||
origin = %request.origin,
|
||||
next_hop = %next_hop_addr,
|
||||
"Sending LookupResponse"
|
||||
);
|
||||
|
||||
let encoded = response.encode();
|
||||
if let Err(e) = self.send_encrypted_link_message(&next_hop_addr, &encoded).await {
|
||||
debug!(
|
||||
next_hop = %next_hop_addr,
|
||||
error = %e,
|
||||
"Failed to send LookupResponse"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward a LookupRequest to peers not in the visited filter.
|
||||
///
|
||||
/// Decrements TTL, adds self to visited, and sends to all eligible peers.
|
||||
async fn forward_lookup_request(&mut self, mut request: LookupRequest) {
|
||||
if !request.forward(self.node_addr()) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Collect peers not in visited filter
|
||||
let forward_to: Vec<NodeAddr> = self
|
||||
.peers
|
||||
.keys()
|
||||
.filter(|addr| !request.was_visited(addr))
|
||||
.copied()
|
||||
.collect();
|
||||
|
||||
if forward_to.is_empty() {
|
||||
trace!(
|
||||
request_id = request.request_id,
|
||||
"No eligible peers to forward LookupRequest"
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
debug!(
|
||||
request_id = request.request_id,
|
||||
target = %request.target,
|
||||
ttl = request.ttl,
|
||||
peer_count = forward_to.len(),
|
||||
"Forwarding LookupRequest"
|
||||
);
|
||||
|
||||
let encoded = request.encode();
|
||||
|
||||
for peer_addr in forward_to {
|
||||
if let Err(e) = self.send_encrypted_link_message(&peer_addr, &encoded).await {
|
||||
debug!(
|
||||
peer = %peer_addr,
|
||||
error = %e,
|
||||
"Failed to forward LookupRequest to peer"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Initiate a discovery lookup for a target node.
|
||||
///
|
||||
/// Creates a LookupRequest and floods it to all peers. The originator
|
||||
/// does NOT record the request_id in recent_requests, so when the
|
||||
/// response arrives, it's recognized as "our request" and the
|
||||
/// target's coordinates are cached in route_cache.
|
||||
#[allow(dead_code)] // Called from integration tests; will be used from event loop
|
||||
pub(in crate::node) async fn initiate_lookup(&mut self, target: &NodeAddr, ttl: u8) {
|
||||
let origin = *self.node_addr();
|
||||
let origin_coords = self.tree_state().my_coords().clone();
|
||||
let mut request = LookupRequest::generate(*target, origin, origin_coords, ttl);
|
||||
|
||||
// Add ourselves to the visited filter so forwarding nodes
|
||||
// won't send the request back to us
|
||||
request.visited.insert(&origin);
|
||||
|
||||
debug!(
|
||||
request_id = request.request_id,
|
||||
target = %target,
|
||||
ttl = ttl,
|
||||
"Initiating LookupRequest"
|
||||
);
|
||||
|
||||
// Send to all peers (flood)
|
||||
let peer_addrs: Vec<NodeAddr> = self.peers.keys().copied().collect();
|
||||
let encoded = request.encode();
|
||||
|
||||
for peer_addr in peer_addrs {
|
||||
if let Err(e) = self.send_encrypted_link_message(&peer_addr, &encoded).await {
|
||||
debug!(
|
||||
peer = %peer_addr,
|
||||
error = %e,
|
||||
"Failed to send LookupRequest to peer"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove expired entries from the recent_requests cache.
|
||||
fn purge_expired_requests(&mut self, current_time_ms: u64) {
|
||||
self.recent_requests
|
||||
.retain(|_, entry| !entry.is_expired(current_time_ms));
|
||||
}
|
||||
|
||||
/// Get current time in milliseconds since Unix epoch.
|
||||
fn now_ms() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as u64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
}
|
||||
@@ -16,7 +16,6 @@ impl Node {
|
||||
let msg_type = plaintext[0];
|
||||
let payload = &plaintext[1..];
|
||||
|
||||
// TODO: Implement remaining link message handlers
|
||||
match msg_type {
|
||||
0x10 => {
|
||||
// TreeAnnounce
|
||||
@@ -28,11 +27,11 @@ impl Node {
|
||||
}
|
||||
0x30 => {
|
||||
// LookupRequest
|
||||
debug!("Received LookupRequest (not yet implemented)");
|
||||
self.handle_lookup_request(from, payload).await;
|
||||
}
|
||||
0x31 => {
|
||||
// LookupResponse
|
||||
debug!("Received LookupResponse (not yet implemented)");
|
||||
self.handle_lookup_response(from, payload).await;
|
||||
}
|
||||
0x40 => {
|
||||
// SessionDatagram
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
//! RX event loop and message handlers.
|
||||
|
||||
mod discovery;
|
||||
mod dispatch;
|
||||
mod encrypted;
|
||||
mod forwarding;
|
||||
|
||||
@@ -13,7 +13,7 @@ mod tree;
|
||||
mod tests;
|
||||
|
||||
use crate::bloom::BloomState;
|
||||
use crate::cache::CoordCache;
|
||||
use crate::cache::{CoordCache, RouteCache};
|
||||
use crate::index::IndexAllocator;
|
||||
use crate::peer::{ActivePeer, PeerConnection};
|
||||
use crate::rate_limit::HandshakeRateLimiter;
|
||||
@@ -142,6 +142,33 @@ impl fmt::Display for NodeState {
|
||||
}
|
||||
}
|
||||
|
||||
/// Recent request tracking for dedup and reverse-path forwarding.
|
||||
///
|
||||
/// When a LookupRequest is forwarded through a node, the node stores the
|
||||
/// request_id and which peer sent it. When the corresponding LookupResponse
|
||||
/// arrives, it's forwarded back to that peer (reverse-path forwarding).
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct RecentRequest {
|
||||
/// The peer who sent this request to us.
|
||||
pub(crate) from_peer: NodeAddr,
|
||||
/// When we received this request (Unix milliseconds).
|
||||
pub(crate) timestamp_ms: u64,
|
||||
}
|
||||
|
||||
impl RecentRequest {
|
||||
pub(crate) fn new(from_peer: NodeAddr, timestamp_ms: u64) -> Self {
|
||||
Self {
|
||||
from_peer,
|
||||
timestamp_ms,
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if this entry has expired (older than 10 seconds).
|
||||
pub(crate) fn is_expired(&self, current_time_ms: u64) -> bool {
|
||||
current_time_ms.saturating_sub(self.timestamp_ms) > 10_000
|
||||
}
|
||||
}
|
||||
|
||||
/// Key for addr_to_link reverse lookup.
|
||||
type AddrKey = (TransportId, TransportAddr);
|
||||
|
||||
@@ -182,8 +209,13 @@ pub struct Node {
|
||||
bloom_state: BloomState,
|
||||
|
||||
// === Routing ===
|
||||
/// Address -> coordinates cache.
|
||||
/// Address -> coordinates cache (from session setup).
|
||||
coord_cache: CoordCache,
|
||||
/// Discovered routes (from discovery protocol).
|
||||
route_cache: RouteCache,
|
||||
/// Recent discovery requests (dedup + reverse-path forwarding).
|
||||
/// Maps request_id → RecentRequest.
|
||||
recent_requests: HashMap<u64, RecentRequest>,
|
||||
|
||||
// === Transports & Links ===
|
||||
/// Active transports (owned by Node).
|
||||
@@ -294,6 +326,8 @@ impl Node {
|
||||
tree_state,
|
||||
bloom_state,
|
||||
coord_cache: CoordCache::with_defaults(),
|
||||
route_cache: RouteCache::with_defaults(),
|
||||
recent_requests: HashMap::new(),
|
||||
transports: HashMap::new(),
|
||||
links: HashMap::new(),
|
||||
addr_to_link: HashMap::new(),
|
||||
@@ -343,6 +377,8 @@ impl Node {
|
||||
tree_state,
|
||||
bloom_state: BloomState::new(node_addr),
|
||||
coord_cache: CoordCache::with_defaults(),
|
||||
route_cache: RouteCache::with_defaults(),
|
||||
recent_requests: HashMap::new(),
|
||||
transports: HashMap::new(),
|
||||
links: HashMap::new(),
|
||||
addr_to_link: HashMap::new(),
|
||||
@@ -492,6 +528,18 @@ impl Node {
|
||||
&mut self.coord_cache
|
||||
}
|
||||
|
||||
// === Route Cache ===
|
||||
|
||||
/// Get the route cache (discovery protocol).
|
||||
pub fn route_cache(&self) -> &RouteCache {
|
||||
&self.route_cache
|
||||
}
|
||||
|
||||
/// Get mutable route cache.
|
||||
pub fn route_cache_mut(&mut self) -> &mut RouteCache {
|
||||
&mut self.route_cache
|
||||
}
|
||||
|
||||
// === TUN Interface ===
|
||||
|
||||
/// Get the TUN state.
|
||||
|
||||
361
src/node/tests/discovery.rs
Normal file
361
src/node/tests/discovery.rs
Normal file
@@ -0,0 +1,361 @@
|
||||
//! Discovery protocol tests: LookupRequest and LookupResponse.
|
||||
//!
|
||||
//! Unit tests for handler logic (dedup, visited filter, TTL, response
|
||||
//! caching) and integration tests for multi-node forwarding and
|
||||
//! reverse-path response routing.
|
||||
|
||||
use super::*;
|
||||
use crate::node::RecentRequest;
|
||||
use crate::protocol::{LookupRequest, LookupResponse};
|
||||
use crate::tree::TreeCoordinate;
|
||||
use spanning_tree::{cleanup_nodes, process_available_packets, run_tree_test};
|
||||
|
||||
// ============================================================================
|
||||
// Unit Tests — LookupRequest Handler
|
||||
// ============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_request_decode_error() {
|
||||
let mut node = make_node();
|
||||
let from = make_node_addr(0xAA);
|
||||
// Too-short payload: should log error and return without panic
|
||||
node.handle_lookup_request(&from, &[0x00; 5]).await;
|
||||
assert!(node.recent_requests.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_request_dedup() {
|
||||
let mut node = make_node();
|
||||
let from = make_node_addr(0xAA);
|
||||
let target = make_node_addr(0xBB);
|
||||
let origin = make_node_addr(0xCC);
|
||||
let coords = TreeCoordinate::from_addrs(vec![origin, make_node_addr(0)]).unwrap();
|
||||
|
||||
let request = LookupRequest::new(999, target, origin, coords, 5);
|
||||
let payload = &request.encode()[1..]; // skip msg_type byte
|
||||
|
||||
// First request: accepted
|
||||
node.handle_lookup_request(&from, payload).await;
|
||||
assert_eq!(node.recent_requests.len(), 1);
|
||||
|
||||
// Duplicate request: dropped
|
||||
node.handle_lookup_request(&from, payload).await;
|
||||
assert_eq!(node.recent_requests.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_request_visited_filter_self() {
|
||||
let mut node = make_node();
|
||||
let from = make_node_addr(0xAA);
|
||||
let target = make_node_addr(0xBB);
|
||||
let origin = make_node_addr(0xCC);
|
||||
let coords = TreeCoordinate::from_addrs(vec![origin, make_node_addr(0)]).unwrap();
|
||||
|
||||
let mut request = LookupRequest::new(888, target, origin, coords, 5);
|
||||
// Mark ourselves as already visited
|
||||
request.visited.insert(node.node_addr());
|
||||
|
||||
let payload = &request.encode()[1..];
|
||||
node.handle_lookup_request(&from, payload).await;
|
||||
|
||||
// Request was recorded (dedup happens before visited check)
|
||||
// but the handler should have stopped after detecting self in visited filter
|
||||
assert!(node.recent_requests.contains_key(&888));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_request_target_is_self() {
|
||||
let mut node = make_node();
|
||||
let from = make_node_addr(0xAA);
|
||||
let origin = make_node_addr(0xCC);
|
||||
let my_addr = *node.node_addr();
|
||||
let coords = TreeCoordinate::from_addrs(vec![origin, make_node_addr(0)]).unwrap();
|
||||
|
||||
// Request targeting us
|
||||
let request = LookupRequest::new(777, my_addr, origin, coords, 5);
|
||||
let payload = &request.encode()[1..];
|
||||
|
||||
// Should succeed without panic (response send will fail silently
|
||||
// since we have no peers to route toward origin)
|
||||
node.handle_lookup_request(&from, payload).await;
|
||||
assert!(node.recent_requests.contains_key(&777));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_request_ttl_zero_not_forwarded() {
|
||||
let mut node = make_node();
|
||||
let from = make_node_addr(0xAA);
|
||||
let target = make_node_addr(0xBB);
|
||||
let origin = make_node_addr(0xCC);
|
||||
let coords = TreeCoordinate::from_addrs(vec![origin, make_node_addr(0)]).unwrap();
|
||||
|
||||
let request = LookupRequest::new(666, target, origin, coords, 0);
|
||||
let payload = &request.encode()[1..];
|
||||
|
||||
node.handle_lookup_request(&from, payload).await;
|
||||
// Request recorded, but not forwarded (TTL=0, and no peers anyway)
|
||||
assert!(node.recent_requests.contains_key(&666));
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Unit Tests — LookupResponse Handler
|
||||
// ============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_response_decode_error() {
|
||||
let mut node = make_node();
|
||||
let from = make_node_addr(0xAA);
|
||||
node.handle_lookup_response(&from, &[0x00; 10]).await;
|
||||
// No panic, no route cached
|
||||
assert!(node.route_cache.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_response_originator_caches_route() {
|
||||
let mut node = make_node();
|
||||
let from = make_node_addr(0xAA);
|
||||
let target = make_node_addr(0xBB);
|
||||
let root = make_node_addr(0xF0);
|
||||
let coords = TreeCoordinate::from_addrs(vec![target, root]).unwrap();
|
||||
|
||||
// Create a valid response with a real proof signature
|
||||
let proof_data = LookupResponse::proof_bytes(555, &target);
|
||||
let target_identity = Identity::generate();
|
||||
let proof = target_identity.sign(&proof_data);
|
||||
|
||||
let response = LookupResponse::new(555, target, coords.clone(), proof);
|
||||
let payload = &response.encode()[1..]; // skip msg_type
|
||||
|
||||
// No entry in recent_requests for 555 → we're the originator
|
||||
assert!(!node.recent_requests.contains_key(&555));
|
||||
|
||||
node.handle_lookup_response(&from, payload).await;
|
||||
|
||||
// Route should be cached
|
||||
assert!(node.route_cache.contains(&target));
|
||||
let cached = node.route_cache.get(&target).unwrap();
|
||||
assert_eq!(cached.coords(), &coords);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_response_transit_needs_recent_request() {
|
||||
let mut node = make_node();
|
||||
let from = make_node_addr(0xAA);
|
||||
let target = make_node_addr(0xBB);
|
||||
let root = make_node_addr(0xF0);
|
||||
let coords = TreeCoordinate::from_addrs(vec![target, root]).unwrap();
|
||||
|
||||
let proof_data = LookupResponse::proof_bytes(444, &target);
|
||||
let target_identity = Identity::generate();
|
||||
let proof = target_identity.sign(&proof_data);
|
||||
|
||||
let response = LookupResponse::new(444, target, coords, proof);
|
||||
let payload = &response.encode()[1..];
|
||||
|
||||
// Simulate being a transit node: record a recent_request for this ID
|
||||
let now_ms = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_millis() as u64;
|
||||
node.recent_requests.insert(
|
||||
444,
|
||||
RecentRequest::new(make_node_addr(0xDD), now_ms),
|
||||
);
|
||||
|
||||
// Handle response — should try to reverse-path forward to 0xDD
|
||||
// (will fail silently since 0xDD is not an actual peer)
|
||||
node.handle_lookup_response(&from, payload).await;
|
||||
|
||||
// Should NOT cache in route_cache (we're transit, not originator)
|
||||
assert!(!node.route_cache.contains(&target));
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Unit Tests — RecentRequest Expiry
|
||||
// ============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_recent_request_expiry() {
|
||||
let mut node = make_node();
|
||||
|
||||
let now_ms = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_millis() as u64;
|
||||
|
||||
// Insert an old request (11 seconds ago)
|
||||
node.recent_requests.insert(
|
||||
123,
|
||||
RecentRequest::new(make_node_addr(1), now_ms - 11_000),
|
||||
);
|
||||
|
||||
// Insert a recent request
|
||||
node.recent_requests.insert(
|
||||
456,
|
||||
RecentRequest::new(make_node_addr(2), now_ms),
|
||||
);
|
||||
|
||||
assert_eq!(node.recent_requests.len(), 2);
|
||||
|
||||
// Trigger purge via a new lookup request
|
||||
let target = make_node_addr(0xBB);
|
||||
let origin = make_node_addr(0xCC);
|
||||
let coords = TreeCoordinate::from_addrs(vec![origin, make_node_addr(0)]).unwrap();
|
||||
let request = LookupRequest::new(789, target, origin, coords, 3);
|
||||
let payload = &request.encode()[1..];
|
||||
node.handle_lookup_request(&make_node_addr(0xAA), payload).await;
|
||||
|
||||
// Old entry (123) should be purged, recent entry (456) and new entry (789) kept
|
||||
assert!(!node.recent_requests.contains_key(&123));
|
||||
assert!(node.recent_requests.contains_key(&456));
|
||||
assert!(node.recent_requests.contains_key(&789));
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Integration Tests — Multi-Node Forwarding
|
||||
// ============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_request_forwarding_two_node() {
|
||||
// Set up a two-node topology: node0 — node1
|
||||
// Send a LookupRequest from node0 targeting some unknown node.
|
||||
// Node1 should receive the forwarded request.
|
||||
let edges = vec![(0, 1)];
|
||||
let mut nodes = run_tree_test(2, &edges, false).await;
|
||||
|
||||
let node0_addr = *nodes[0].node.node_addr();
|
||||
let target = make_node_addr(0xEE); // unknown node
|
||||
let root = make_node_addr(0);
|
||||
|
||||
let coords = TreeCoordinate::from_addrs(vec![node0_addr, root]).unwrap();
|
||||
let request = LookupRequest::new(42, target, node0_addr, coords, 5);
|
||||
let payload = &request.encode()[1..];
|
||||
|
||||
// Handle on node0 as if we received it from outside
|
||||
nodes[0]
|
||||
.node
|
||||
.handle_lookup_request(&node0_addr, payload)
|
||||
.await;
|
||||
|
||||
// Process packets — node1 should receive the forwarded request
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
let count = process_available_packets(&mut nodes).await;
|
||||
assert!(count > 0, "Expected forwarded LookupRequest to arrive at node 1");
|
||||
|
||||
// Node1 should have recorded the request
|
||||
assert!(
|
||||
nodes[1].node.recent_requests.contains_key(&42),
|
||||
"Node 1 should have recorded the forwarded request"
|
||||
);
|
||||
|
||||
cleanup_nodes(&mut nodes).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_request_target_found_generates_response() {
|
||||
// Set up a two-node topology: node0 — node1
|
||||
// Node0 initiates a lookup targeting node1.
|
||||
// Node1 receives, detects it's the target, generates a LookupResponse.
|
||||
// Response routes back to node0 which caches the coordinates.
|
||||
let edges = vec![(0, 1)];
|
||||
let mut nodes = run_tree_test(2, &edges, false).await;
|
||||
|
||||
let node1_addr = *nodes[1].node.node_addr();
|
||||
|
||||
// Node0 initiates lookup (doesn't record in recent_requests)
|
||||
nodes[0].node.initiate_lookup(&node1_addr, 5).await;
|
||||
|
||||
// Process packets in rounds to allow request + response
|
||||
for _ in 0..4 {
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
process_available_packets(&mut nodes).await;
|
||||
}
|
||||
|
||||
// Node0 should have cached node1's route (it originated the request)
|
||||
assert!(
|
||||
nodes[0].node.route_cache.contains(&node1_addr),
|
||||
"Node 0 should have cached node 1's route from LookupResponse"
|
||||
);
|
||||
|
||||
cleanup_nodes(&mut nodes).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_request_three_node_chain() {
|
||||
// Topology: node0 — node1 — node2
|
||||
// Node0 initiates a lookup targeting node2.
|
||||
// Request should propagate: node0 → node1 → node2.
|
||||
// Node2 generates response, reverse-path: node2 → node1 → node0.
|
||||
let edges = vec![(0, 1), (1, 2)];
|
||||
let mut nodes = run_tree_test(3, &edges, false).await;
|
||||
|
||||
let node2_addr = *nodes[2].node.node_addr();
|
||||
|
||||
// Node0 initiates lookup (doesn't record in recent_requests)
|
||||
nodes[0].node.initiate_lookup(&node2_addr, 8).await;
|
||||
|
||||
// Process packets in rounds to allow multi-hop propagation + response
|
||||
// Chain: node0→node1→node2 (request), node2→node1→node0 (response)
|
||||
for _ in 0..10 {
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
process_available_packets(&mut nodes).await;
|
||||
}
|
||||
|
||||
// Node1 should have been a transit node (has the request_id in recent_requests)
|
||||
assert!(
|
||||
!nodes[1].node.recent_requests.is_empty(),
|
||||
"Node 1 should have recorded the forwarded request"
|
||||
);
|
||||
|
||||
// Node2 should have received the request (it's the target)
|
||||
assert!(
|
||||
!nodes[2].node.recent_requests.is_empty(),
|
||||
"Node 2 should have received the request"
|
||||
);
|
||||
|
||||
// Node0 should have cached node2's route
|
||||
assert!(
|
||||
nodes[0].node.route_cache.contains(&node2_addr),
|
||||
"Node 0 should have cached node 2's route through 3-node chain"
|
||||
);
|
||||
|
||||
cleanup_nodes(&mut nodes).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_request_dedup_convergent_paths() {
|
||||
// Topology: triangle (node0 — node1, node0 — node2, node1 — node2)
|
||||
// A request from node0 reaches node2 via two paths: 0→1→2 and 0→2.
|
||||
// The second arrival at node2 should be deduped.
|
||||
let edges = vec![(0, 1), (0, 2), (1, 2)];
|
||||
let mut nodes = run_tree_test(3, &edges, false).await;
|
||||
|
||||
let node0_addr = *nodes[0].node.node_addr();
|
||||
let target = make_node_addr(0xEE);
|
||||
let root = make_node_addr(0);
|
||||
|
||||
let coords = TreeCoordinate::from_addrs(vec![node0_addr, root]).unwrap();
|
||||
let request = LookupRequest::new(300, target, node0_addr, coords, 5);
|
||||
let payload = &request.encode()[1..];
|
||||
|
||||
// Node0 handles the request (forwards to both node1 and node2)
|
||||
nodes[0]
|
||||
.node
|
||||
.handle_lookup_request(&node0_addr, payload)
|
||||
.await;
|
||||
|
||||
// Process several rounds
|
||||
for _ in 0..5 {
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
process_available_packets(&mut nodes).await;
|
||||
}
|
||||
|
||||
// Both node1 and node2 should have recorded the request
|
||||
assert!(nodes[1].node.recent_requests.contains_key(&300));
|
||||
assert!(nodes[2].node.recent_requests.contains_key(&300));
|
||||
|
||||
// The request should appear exactly once in each node's recent_requests
|
||||
// (dedup prevents duplicate processing via convergent paths)
|
||||
|
||||
cleanup_nodes(&mut nodes).await;
|
||||
}
|
||||
@@ -6,6 +6,7 @@ use std::time::Duration;
|
||||
|
||||
mod bloom;
|
||||
mod disconnect;
|
||||
mod discovery;
|
||||
mod forwarding;
|
||||
mod handshake;
|
||||
mod routing;
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
//! Discovery messages: LookupRequest and LookupResponse.
|
||||
|
||||
use crate::bloom::BloomFilter;
|
||||
use crate::protocol::error::ProtocolError;
|
||||
use crate::protocol::session::{decode_coords, encode_coords};
|
||||
use crate::tree::TreeCoordinate;
|
||||
use crate::NodeAddr;
|
||||
use secp256k1::schnorr::Signature;
|
||||
@@ -79,6 +81,90 @@ impl LookupRequest {
|
||||
pub fn was_visited(&self, node_addr: &NodeAddr) -> bool {
|
||||
self.visited.contains(node_addr)
|
||||
}
|
||||
|
||||
/// Encode as wire format (includes msg_type byte).
|
||||
///
|
||||
/// Format: `[0x30][request_id:8][target:16][origin:16][ttl:1]`
|
||||
/// `[origin_coords_cnt:2][origin_coords:16×n]`
|
||||
/// `[visited_hash_cnt:1][visited_bits:256]`
|
||||
pub fn encode(&self) -> Vec<u8> {
|
||||
let visited_bytes = self.visited.as_bytes();
|
||||
let mut buf = Vec::with_capacity(44 + self.origin_coords.depth() * 16 + 1 + visited_bytes.len());
|
||||
|
||||
buf.push(0x30); // msg_type
|
||||
buf.extend_from_slice(&self.request_id.to_le_bytes());
|
||||
buf.extend_from_slice(self.target.as_bytes());
|
||||
buf.extend_from_slice(self.origin.as_bytes());
|
||||
buf.push(self.ttl);
|
||||
encode_coords(&self.origin_coords, &mut buf);
|
||||
buf.push(self.visited.hash_count());
|
||||
buf.extend_from_slice(visited_bytes);
|
||||
|
||||
buf
|
||||
}
|
||||
|
||||
/// Decode from wire format (after msg_type byte has been consumed).
|
||||
pub fn decode(payload: &[u8]) -> Result<Self, ProtocolError> {
|
||||
// Minimum: request_id(8) + target(16) + origin(16) + ttl(1)
|
||||
// + coords_count(2) + hash_count(1) = 44 bytes
|
||||
if payload.len() < 44 {
|
||||
return Err(ProtocolError::MessageTooShort {
|
||||
expected: 44,
|
||||
got: payload.len(),
|
||||
});
|
||||
}
|
||||
|
||||
let mut pos = 0;
|
||||
|
||||
let request_id = u64::from_le_bytes(
|
||||
payload[pos..pos + 8]
|
||||
.try_into()
|
||||
.map_err(|_| ProtocolError::Malformed("bad request_id".into()))?,
|
||||
);
|
||||
pos += 8;
|
||||
|
||||
let mut target_bytes = [0u8; 16];
|
||||
target_bytes.copy_from_slice(&payload[pos..pos + 16]);
|
||||
let target = NodeAddr::from_bytes(target_bytes);
|
||||
pos += 16;
|
||||
|
||||
let mut origin_bytes = [0u8; 16];
|
||||
origin_bytes.copy_from_slice(&payload[pos..pos + 16]);
|
||||
let origin = NodeAddr::from_bytes(origin_bytes);
|
||||
pos += 16;
|
||||
|
||||
let ttl = payload[pos];
|
||||
pos += 1;
|
||||
|
||||
let (origin_coords, consumed) = decode_coords(&payload[pos..])?;
|
||||
pos += consumed;
|
||||
|
||||
if payload.len() < pos + 1 {
|
||||
return Err(ProtocolError::MessageTooShort {
|
||||
expected: pos + 1,
|
||||
got: payload.len(),
|
||||
});
|
||||
}
|
||||
let hash_count = payload[pos];
|
||||
pos += 1;
|
||||
|
||||
let filter_bytes = &payload[pos..];
|
||||
if filter_bytes.is_empty() {
|
||||
return Err(ProtocolError::Malformed("visited filter missing".into()));
|
||||
}
|
||||
|
||||
let visited = BloomFilter::from_slice(filter_bytes, hash_count)
|
||||
.map_err(|e| ProtocolError::Malformed(format!("bad visited filter: {e}")))?;
|
||||
|
||||
Ok(Self {
|
||||
request_id,
|
||||
target,
|
||||
origin,
|
||||
origin_coords,
|
||||
ttl,
|
||||
visited,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Response to a lookup request with target's coordinates.
|
||||
@@ -121,6 +207,65 @@ impl LookupResponse {
|
||||
bytes.extend_from_slice(target.as_bytes());
|
||||
bytes
|
||||
}
|
||||
|
||||
/// Encode as wire format (includes msg_type byte).
|
||||
///
|
||||
/// Format: `[0x31][request_id:8][target:16][target_coords_cnt:2][target_coords:16×n][proof:64]`
|
||||
pub fn encode(&self) -> Vec<u8> {
|
||||
let mut buf = Vec::with_capacity(91 + self.target_coords.depth() * 16);
|
||||
|
||||
buf.push(0x31); // msg_type
|
||||
buf.extend_from_slice(&self.request_id.to_le_bytes());
|
||||
buf.extend_from_slice(self.target.as_bytes());
|
||||
encode_coords(&self.target_coords, &mut buf);
|
||||
buf.extend_from_slice(self.proof.as_ref());
|
||||
|
||||
buf
|
||||
}
|
||||
|
||||
/// Decode from wire format (after msg_type byte has been consumed).
|
||||
pub fn decode(payload: &[u8]) -> Result<Self, ProtocolError> {
|
||||
// Minimum: request_id(8) + target(16) + coords_count(2) + proof(64) = 90
|
||||
if payload.len() < 90 {
|
||||
return Err(ProtocolError::MessageTooShort {
|
||||
expected: 90,
|
||||
got: payload.len(),
|
||||
});
|
||||
}
|
||||
|
||||
let mut pos = 0;
|
||||
|
||||
let request_id = u64::from_le_bytes(
|
||||
payload[pos..pos + 8]
|
||||
.try_into()
|
||||
.map_err(|_| ProtocolError::Malformed("bad request_id".into()))?,
|
||||
);
|
||||
pos += 8;
|
||||
|
||||
let mut target_bytes = [0u8; 16];
|
||||
target_bytes.copy_from_slice(&payload[pos..pos + 16]);
|
||||
let target = NodeAddr::from_bytes(target_bytes);
|
||||
pos += 16;
|
||||
|
||||
let (target_coords, consumed) = decode_coords(&payload[pos..])?;
|
||||
pos += consumed;
|
||||
|
||||
if payload.len() < pos + 64 {
|
||||
return Err(ProtocolError::MessageTooShort {
|
||||
expected: pos + 64,
|
||||
got: payload.len(),
|
||||
});
|
||||
}
|
||||
let proof = Signature::from_slice(&payload[pos..pos + 64])
|
||||
.map_err(|_| ProtocolError::Malformed("bad proof signature".into()))?;
|
||||
|
||||
Ok(Self {
|
||||
request_id,
|
||||
target,
|
||||
target_coords,
|
||||
proof,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -190,4 +335,62 @@ mod tests {
|
||||
assert_eq!(&bytes[0..8], &12345u64.to_le_bytes());
|
||||
assert_eq!(&bytes[8..24], target.as_bytes());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lookup_request_encode_decode_roundtrip() {
|
||||
let target = make_node_addr(10);
|
||||
let origin = make_node_addr(20);
|
||||
let coords = make_coords(&[20, 0]);
|
||||
|
||||
let mut request = LookupRequest::new(12345, target, origin, coords.clone(), 8);
|
||||
request.forward(&make_node_addr(30));
|
||||
|
||||
let encoded = request.encode();
|
||||
assert_eq!(encoded[0], 0x30);
|
||||
|
||||
let decoded = LookupRequest::decode(&encoded[1..]).unwrap();
|
||||
assert_eq!(decoded.request_id, 12345);
|
||||
assert_eq!(decoded.target, target);
|
||||
assert_eq!(decoded.origin, origin);
|
||||
assert_eq!(decoded.ttl, 7); // decremented by forward()
|
||||
assert!(decoded.was_visited(&make_node_addr(30)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lookup_request_decode_too_short() {
|
||||
assert!(LookupRequest::decode(&[]).is_err());
|
||||
assert!(LookupRequest::decode(&[0u8; 40]).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lookup_response_encode_decode_roundtrip() {
|
||||
use secp256k1::Secp256k1;
|
||||
|
||||
let target = make_node_addr(42);
|
||||
let coords = make_coords(&[42, 1, 0]);
|
||||
|
||||
// Create a dummy signature for testing
|
||||
let secp = Secp256k1::new();
|
||||
let keypair = secp256k1::Keypair::new(&secp, &mut rand::thread_rng());
|
||||
let proof_data = LookupResponse::proof_bytes(999, &target);
|
||||
use sha2::Digest;
|
||||
let digest: [u8; 32] = sha2::Sha256::digest(&proof_data).into();
|
||||
let sig = secp.sign_schnorr(&digest, &keypair);
|
||||
|
||||
let response = LookupResponse::new(999, target, coords.clone(), sig);
|
||||
|
||||
let encoded = response.encode();
|
||||
assert_eq!(encoded[0], 0x31);
|
||||
|
||||
let decoded = LookupResponse::decode(&encoded[1..]).unwrap();
|
||||
assert_eq!(decoded.request_id, 999);
|
||||
assert_eq!(decoded.target, target);
|
||||
assert_eq!(decoded.proof, sig);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lookup_response_decode_too_short() {
|
||||
assert!(LookupResponse::decode(&[]).is_err());
|
||||
assert!(LookupResponse::decode(&[0u8; 50]).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -76,7 +76,7 @@ impl fmt::Display for SessionMessageType {
|
||||
///
|
||||
/// Session-layer messages serialize coordinates as NodeAddr arrays (16 bytes each),
|
||||
/// without the sequence/timestamp metadata used by the tree gossip protocol.
|
||||
fn encode_coords(coords: &TreeCoordinate, buf: &mut Vec<u8>) {
|
||||
pub(crate) fn encode_coords(coords: &TreeCoordinate, buf: &mut Vec<u8>) {
|
||||
let addrs: Vec<&NodeAddr> = coords.node_addrs().collect();
|
||||
let count = addrs.len() as u16;
|
||||
buf.extend_from_slice(&count.to_le_bytes());
|
||||
@@ -88,7 +88,7 @@ fn encode_coords(coords: &TreeCoordinate, buf: &mut Vec<u8>) {
|
||||
/// Decode a TreeCoordinate from address-only wire format.
|
||||
///
|
||||
/// Returns the decoded coordinate and the number of bytes consumed.
|
||||
fn decode_coords(data: &[u8]) -> Result<(TreeCoordinate, usize), ProtocolError> {
|
||||
pub(crate) fn decode_coords(data: &[u8]) -> Result<(TreeCoordinate, usize), ProtocolError> {
|
||||
if data.len() < 2 {
|
||||
return Err(ProtocolError::MessageTooShort {
|
||||
expected: 2,
|
||||
|
||||
Reference in New Issue
Block a user