use std::collections::HashMap;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum SessionCapability {
BlockExchange,
VectorSearch,
TensorLogic,
GossipRelay,
Bootstrap,
}
fn fnv1a(data: &[u8]) -> u64 {
const OFFSET_BASIS: u64 = 14_695_981_039_346_656_037;
const PRIME: u64 = 1_099_511_628_211;
let mut hash = OFFSET_BASIS;
for &byte in data {
hash ^= u64::from(byte);
hash = hash.wrapping_mul(PRIME);
}
hash
}
fn compute_token_id(peer_id: &str, nonce: u64) -> u64 {
let combined = format!("{}{}", peer_id, nonce);
fnv1a(combined.as_bytes())
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SessionToken {
pub token_id: u64,
pub peer_id: String,
pub nonce: u64,
pub issued_at_secs: u64,
pub expires_at_secs: u64,
}
impl SessionToken {
pub fn is_expired(&self, now_secs: u64) -> bool {
now_secs >= self.expires_at_secs
}
pub fn remaining_secs(&self, now_secs: u64) -> u64 {
self.expires_at_secs.saturating_sub(now_secs)
}
}
#[derive(Clone, Debug)]
pub struct PeerSession {
pub token: SessionToken,
pub capabilities: Vec<SessionCapability>,
pub bytes_sent: u64,
pub bytes_received: u64,
pub message_count: u64,
}
impl PeerSession {
pub fn has_capability(&self, cap: SessionCapability) -> bool {
self.capabilities.contains(&cap)
}
}
#[derive(Clone, Debug)]
pub struct SessionManagerConfig {
pub session_ttl_secs: u64,
pub max_sessions: usize,
pub nonce_seed: u64,
}
impl Default for SessionManagerConfig {
fn default() -> Self {
Self {
session_ttl_secs: 3600,
max_sessions: 1000,
nonce_seed: 0,
}
}
}
pub struct PeerSessionManager {
sessions: HashMap<u64, PeerSession>,
config: SessionManagerConfig,
next_nonce: u64,
}
impl PeerSessionManager {
pub fn new(config: SessionManagerConfig) -> Self {
let next_nonce = config.nonce_seed;
Self {
sessions: HashMap::new(),
config,
next_nonce,
}
}
pub fn open_session(
&mut self,
peer_id: String,
capabilities: Vec<SessionCapability>,
now_secs: u64,
) -> Result<u64, String> {
if self.sessions.len() >= self.config.max_sessions {
return Err("session limit reached".to_string());
}
let nonce = self.next_nonce;
self.next_nonce = self.next_nonce.wrapping_add(1);
let token_id = compute_token_id(&peer_id, nonce);
let expires_at_secs = now_secs.saturating_add(self.config.session_ttl_secs);
let token = SessionToken {
token_id,
peer_id,
nonce,
issued_at_secs: now_secs,
expires_at_secs,
};
let session = PeerSession {
token,
capabilities,
bytes_sent: 0,
bytes_received: 0,
message_count: 0,
};
self.sessions.insert(token_id, session);
Ok(token_id)
}
pub fn get_session(&self, token_id: u64) -> Option<&PeerSession> {
self.sessions.get(&token_id)
}
pub fn record_traffic(&mut self, token_id: u64, sent: u64, received: u64) {
if let Some(session) = self.sessions.get_mut(&token_id) {
session.bytes_sent = session.bytes_sent.saturating_add(sent);
session.bytes_received = session.bytes_received.saturating_add(received);
session.message_count = session.message_count.saturating_add(1);
}
}
pub fn close_session(&mut self, token_id: u64) -> bool {
self.sessions.remove(&token_id).is_some()
}
pub fn evict_expired(&mut self, now_secs: u64) -> usize {
let before = self.sessions.len();
self.sessions
.retain(|_, session| !session.token.is_expired(now_secs));
before - self.sessions.len()
}
pub fn active_sessions(&self, now_secs: u64) -> Vec<&PeerSession> {
self.sessions
.values()
.filter(|s| !s.token.is_expired(now_secs))
.collect()
}
pub fn sessions_with_capability(
&self,
cap: SessionCapability,
now_secs: u64,
) -> Vec<&PeerSession> {
self.sessions
.values()
.filter(|s| !s.token.is_expired(now_secs) && s.has_capability(cap))
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn default_manager() -> PeerSessionManager {
PeerSessionManager::new(SessionManagerConfig::default())
}
fn now() -> u64 {
1_700_000_000_u64
}
#[test]
fn open_session_returns_token_id() {
let mut mgr = default_manager();
let result = mgr.open_session(
"peer-alpha".to_string(),
vec![SessionCapability::BlockExchange],
now(),
);
assert!(result.is_ok());
}
#[test]
fn token_id_is_correct_fnv1a_hash() {
let mut mgr = PeerSessionManager::new(SessionManagerConfig {
nonce_seed: 42,
..Default::default()
});
let token_id = mgr
.open_session("peer-x".to_string(), vec![], now())
.expect("open_session failed");
let expected = compute_token_id("peer-x", 42);
assert_eq!(token_id, expected);
}
#[test]
fn get_session_returns_some_for_existing() {
let mut mgr = default_manager();
let tid = mgr
.open_session("peer-a".to_string(), vec![], now())
.expect("open_session failed");
assert!(mgr.get_session(tid).is_some());
}
#[test]
fn get_session_returns_none_for_unknown() {
let mgr = default_manager();
assert!(mgr.get_session(9999).is_none());
}
#[test]
fn open_session_errors_when_at_limit() {
let mut mgr = PeerSessionManager::new(SessionManagerConfig {
max_sessions: 2,
..Default::default()
});
mgr.open_session("p1".to_string(), vec![], now())
.expect("first session");
mgr.open_session("p2".to_string(), vec![], now())
.expect("second session");
let err = mgr.open_session("p3".to_string(), vec![], now());
assert_eq!(err, Err("session limit reached".to_string()));
}
#[test]
fn has_capability_returns_true_when_present() {
let mut mgr = default_manager();
let tid = mgr
.open_session(
"peer-b".to_string(),
vec![
SessionCapability::VectorSearch,
SessionCapability::Bootstrap,
],
now(),
)
.expect("open_session failed");
let session = mgr.get_session(tid).expect("session not found");
assert!(session.has_capability(SessionCapability::VectorSearch));
assert!(session.has_capability(SessionCapability::Bootstrap));
}
#[test]
fn has_capability_returns_false_when_absent() {
let mut mgr = default_manager();
let tid = mgr
.open_session(
"peer-c".to_string(),
vec![SessionCapability::BlockExchange],
now(),
)
.expect("open_session failed");
let session = mgr.get_session(tid).expect("session not found");
assert!(!session.has_capability(SessionCapability::TensorLogic));
}
#[test]
fn record_traffic_updates_bytes_and_count() {
let mut mgr = default_manager();
let tid = mgr
.open_session("peer-d".to_string(), vec![], now())
.expect("open_session failed");
mgr.record_traffic(tid, 100, 200);
mgr.record_traffic(tid, 50, 75);
let session = mgr.get_session(tid).expect("session not found");
assert_eq!(session.bytes_sent, 150);
assert_eq!(session.bytes_received, 275);
assert_eq!(session.message_count, 2);
}
#[test]
fn record_traffic_unknown_id_is_noop() {
let mut mgr = default_manager();
mgr.record_traffic(0xdeadbeef, 1, 1);
}
#[test]
fn close_session_removes_session() {
let mut mgr = default_manager();
let tid = mgr
.open_session("peer-e".to_string(), vec![], now())
.expect("open_session failed");
assert!(mgr.close_session(tid));
assert!(mgr.get_session(tid).is_none());
}
#[test]
fn close_session_returns_false_for_unknown() {
let mut mgr = default_manager();
assert!(!mgr.close_session(0xdeadbeef));
}
#[test]
fn evict_expired_removes_only_expired() {
let mut mgr = PeerSessionManager::new(SessionManagerConfig {
session_ttl_secs: 100,
..Default::default()
});
let t0 = 1_000_u64;
let tid_a = mgr
.open_session("peer-f".to_string(), vec![], t0)
.expect("open a");
let tid_b = mgr
.open_session("peer-g".to_string(), vec![], t0)
.expect("open b");
let t1 = t0 + 200;
let tid_c = mgr
.open_session("peer-h".to_string(), vec![], t1)
.expect("open c");
let removed = mgr.evict_expired(t1);
assert_eq!(removed, 2);
assert!(mgr.get_session(tid_a).is_none());
assert!(mgr.get_session(tid_b).is_none());
assert!(mgr.get_session(tid_c).is_some());
}
#[test]
fn evict_expired_returns_zero_when_nothing_expired() {
let mut mgr = default_manager();
mgr.open_session("peer-i".to_string(), vec![], now())
.expect("open");
let removed = mgr.evict_expired(now());
assert_eq!(removed, 0);
}
#[test]
fn active_sessions_counts_correctly() {
let mut mgr = PeerSessionManager::new(SessionManagerConfig {
session_ttl_secs: 50,
..Default::default()
});
let t0 = 1_000_u64;
mgr.open_session("p1".to_string(), vec![], t0).expect("p1");
mgr.open_session("p2".to_string(), vec![], t0).expect("p2");
assert_eq!(mgr.active_sessions(t0 + 30).len(), 2);
assert_eq!(mgr.active_sessions(t0 + 60).len(), 0);
}
#[test]
fn sessions_with_capability_filters_correctly() {
let mut mgr = default_manager();
let t = now();
mgr.open_session(
"p1".to_string(),
vec![SessionCapability::GossipRelay, SessionCapability::Bootstrap],
t,
)
.expect("p1");
mgr.open_session("p2".to_string(), vec![SessionCapability::BlockExchange], t)
.expect("p2");
mgr.open_session("p3".to_string(), vec![SessionCapability::GossipRelay], t)
.expect("p3");
let gossip_peers = mgr.sessions_with_capability(SessionCapability::GossipRelay, t);
assert_eq!(gossip_peers.len(), 2);
let block_peers = mgr.sessions_with_capability(SessionCapability::BlockExchange, t);
assert_eq!(block_peers.len(), 1);
let tensor_peers = mgr.sessions_with_capability(SessionCapability::TensorLogic, t);
assert_eq!(tensor_peers.len(), 0);
}
#[test]
fn token_is_expired_and_remaining_secs() {
let token = SessionToken {
token_id: 1,
peer_id: "p".to_string(),
nonce: 0,
issued_at_secs: 1000,
expires_at_secs: 1100,
};
assert!(!token.is_expired(1099));
assert!(token.is_expired(1100));
assert!(token.is_expired(1200));
assert_eq!(token.remaining_secs(1000), 100);
assert_eq!(token.remaining_secs(1099), 1);
assert_eq!(token.remaining_secs(1100), 0);
assert_eq!(token.remaining_secs(1200), 0);
}
#[test]
fn sessions_with_capability_excludes_expired() {
let mut mgr = PeerSessionManager::new(SessionManagerConfig {
session_ttl_secs: 10,
..Default::default()
});
let t0 = 500_u64;
mgr.open_session(
"old-peer".to_string(),
vec![SessionCapability::VectorSearch],
t0,
)
.expect("old peer");
let t1 = t0 + 20;
mgr.open_session(
"new-peer".to_string(),
vec![SessionCapability::VectorSearch],
t1,
)
.expect("new peer");
let results = mgr.sessions_with_capability(SessionCapability::VectorSearch, t1);
assert_eq!(results.len(), 1);
assert_eq!(results[0].token.peer_id, "new-peer");
}
#[test]
fn nonce_increments_per_session() {
let mut mgr = PeerSessionManager::new(SessionManagerConfig {
nonce_seed: 10,
..Default::default()
});
let t = now();
let tid1 = mgr.open_session("pa".to_string(), vec![], t).expect("s1");
let tid2 = mgr.open_session("pa".to_string(), vec![], t).expect("s2");
assert_ne!(tid1, tid2);
let s1 = mgr.get_session(tid1).expect("s1 not found");
let s2 = mgr.get_session(tid2).expect("s2 not found");
assert_eq!(s1.token.nonce, 10);
assert_eq!(s2.token.nonce, 11);
}
#[test]
fn token_timestamps_are_set_correctly() {
let mut mgr = PeerSessionManager::new(SessionManagerConfig {
session_ttl_secs: 7200,
..Default::default()
});
let t = now();
let tid = mgr
.open_session("peer-ts".to_string(), vec![], t)
.expect("open");
let session = mgr.get_session(tid).expect("get");
assert_eq!(session.token.issued_at_secs, t);
assert_eq!(session.token.expires_at_secs, t + 7200);
}
#[test]
fn close_session_frees_slot_for_reuse() {
let mut mgr = PeerSessionManager::new(SessionManagerConfig {
max_sessions: 1,
..Default::default()
});
let t = now();
let tid = mgr.open_session("p1".to_string(), vec![], t).expect("p1");
assert!(mgr.open_session("p2".to_string(), vec![], t).is_err());
mgr.close_session(tid);
assert!(mgr.open_session("p2".to_string(), vec![], t).is_ok());
}
}