//! Peer database and consensus validation for tracking remote peers and their games. use std::{ cmp::Reverse, collections::HashMap, net::SocketAddr, time::{Duration, Instant}, }; use lanspread_db::db::{Game, GameFileDescription}; /// Information about a discovered peer. #[derive(Clone, Debug)] pub struct PeerInfo { /// Network address of the peer. pub addr: SocketAddr, /// Last time we heard from this peer. pub last_seen: Instant, /// Games this peer has available, keyed by game ID. pub games: HashMap, /// File descriptions for each game, keyed by game ID. pub files: HashMap>, } /// Database tracking all discovered peers and their games. #[derive(Debug)] pub struct PeerGameDB { peers: HashMap, } impl Default for PeerGameDB { fn default() -> Self { Self::new() } } impl PeerGameDB { #[must_use] pub fn new() -> Self { Self { peers: HashMap::new(), } } /// Adds a new peer to the database. pub fn add_peer(&mut self, addr: SocketAddr) { let peer_info = PeerInfo { addr, last_seen: Instant::now(), games: HashMap::new(), files: HashMap::new(), }; self.peers.insert(addr, peer_info); log::info!("Added peer: {addr}"); } /// Removes a peer from the database. pub fn remove_peer(&mut self, addr: &SocketAddr) -> Option { self.peers.remove(addr) } /// Updates the games list for a peer. pub fn update_peer_games(&mut self, addr: SocketAddr, games: Vec) { if let Some(peer) = self.peers.get_mut(&addr) { let mut map = HashMap::with_capacity(games.len()); for game in games { map.insert(game.id.clone(), game); } peer.games = map; peer.last_seen = Instant::now(); log::info!("Updated games for peer: {addr}"); } } /// Updates the file descriptions for a specific game from a peer. pub fn update_peer_game_files( &mut self, addr: SocketAddr, game_id: &str, files: Vec, ) { if let Some(peer) = self.peers.get_mut(&addr) { peer.files.insert(game_id.to_string(), files); peer.last_seen = Instant::now(); } } /// Updates the last seen timestamp for a peer. pub fn update_last_seen(&mut self, addr: &SocketAddr) { if let Some(peer) = self.peers.get_mut(addr) { peer.last_seen = Instant::now(); } } /// Returns all games aggregated from all peers. #[must_use] pub fn get_all_games(&self) -> Vec { let mut aggregated: HashMap = HashMap::new(); let mut peer_counts: HashMap = HashMap::new(); // Count peers per game for peer in self.peers.values() { for game_id in peer.games.keys() { *peer_counts.entry(game_id.clone()).or_insert(0) += 1; } } // Aggregate games with peer counts for peer in self.peers.values() { for game in peer.games.values() { aggregated .entry(game.id.clone()) .and_modify(|existing| { if let (Some(new_version), Some(current)) = (&game.eti_game_version, &existing.eti_game_version) { if new_version > current { existing.eti_game_version = Some(new_version.clone()); } } else if existing.eti_game_version.is_none() { existing.eti_game_version.clone_from(&game.eti_game_version); } // Update peer count existing.peer_count = peer_counts[&game.id]; }) .or_insert_with(|| { let mut game_clone = game.clone(); game_clone.peer_count = peer_counts[&game.id]; game_clone }); } } let mut games: Vec = aggregated.into_values().collect(); games.sort_by(|a, b| a.name.cmp(&b.name)); games } /// Returns the latest version of a game across all peers. #[must_use] pub fn get_latest_version_for_game(&self, game_id: &str) -> Option { let mut latest_version: Option = None; for peer in self.peers.values() { if let Some(game) = peer.games.get(game_id) && let Some(ref version) = game.eti_game_version { match &latest_version { None => latest_version = Some(version.clone()), Some(current_latest) => { if version > current_latest { latest_version = Some(version.clone()); } } } } } latest_version } /// Returns all peer addresses. #[must_use] pub fn get_peer_addresses(&self) -> Vec { self.peers.keys().copied().collect() } /// Checks if a peer is in the database. #[must_use] pub fn contains_peer(&self, addr: &SocketAddr) -> bool { self.peers.contains_key(addr) } /// Returns addresses of peers that have a specific game. #[must_use] pub fn peers_with_game(&self, game_id: &str) -> Vec { self.peers .iter() .filter(|(_, peer)| peer.games.contains_key(game_id)) .map(|(addr, _)| *addr) .collect() } /// Returns addresses of peers that have the latest version of a game. #[must_use] pub fn peers_with_latest_version(&self, game_id: &str) -> Vec { let latest_version = self.get_latest_version_for_game(game_id); if let Some(ref latest) = latest_version { self.peers .iter() .filter(|(_, peer)| { if let Some(game) = peer.games.get(game_id) { if let Some(ref version) = game.eti_game_version { version == latest } else { false } } else { false } }) .map(|(addr, _)| *addr) .collect() } else { // If no version info is available, fall back to all peers with the game self.peers_with_game(game_id) } } /// Returns file descriptions for a game from all peers. #[must_use] pub fn game_files_for(&self, game_id: &str) -> Vec<(SocketAddr, Vec)> { self.peers .iter() .filter_map(|(addr, peer)| peer.files.get(game_id).cloned().map(|files| (*addr, files))) .collect() } /// Returns aggregated file descriptions for a game across all peers. #[must_use] pub fn aggregated_game_files(&self, game_id: &str) -> Vec { let mut seen: HashMap = HashMap::new(); for (_, files) in self.game_files_for(game_id) { for file in files { seen.entry(file.relative_path.clone()).or_insert(file); } } seen.into_values().collect() } /// Returns the majority-agreed size for a game. #[must_use] pub fn majority_game_size(&self, game_id: &str) -> Option { let mut size_counts: HashMap = HashMap::new(); for peer in self.peers.values() { if let Some(game) = peer.games.get(game_id) { if game.size == 0 { continue; } *size_counts.entry(game.size).or_insert(0) += 1; } } size_counts .into_iter() .max_by(|(size_a, count_a), (size_b, count_b)| { count_a.cmp(count_b).then_with(|| size_a.cmp(size_b)) }) .map(|(size, _)| size) } /// Validates file sizes across all peers and returns only the files with majority consensus. /// /// Returns a tuple of (`validated_files`, `peer_whitelist`, `file_peer_map`) where /// `peer_whitelist` contains peers that have at least one majority-approved file and /// `file_peer_map` lists which peers were validated for each file. pub fn validate_file_sizes_majority( &self, game_id: &str, ) -> eyre::Result { let game_files = self.game_files_for(game_id); if game_files.is_empty() { return Ok((Vec::new(), Vec::new(), HashMap::new())); } let (file_size_map, _peer_files) = collect_file_sizes(&game_files); let (validated_files, peer_scores, file_peer_map) = self.validate_each_file_consensus(game_id, file_size_map)?; let peer_whitelist = create_peer_whitelist(peer_scores); Ok((validated_files, peer_whitelist, file_peer_map)) } /// Validates consensus for each file and returns validated files with peer scores. fn validate_each_file_consensus( &self, game_id: &str, file_size_map: FileSizeMap, ) -> eyre::Result { let mut validated_files = Vec::new(); let mut peer_whitelist_scores: HashMap = HashMap::new(); let mut file_peer_map: HashMap> = HashMap::new(); for (relative_path, size_map) in file_size_map { let total_peers: usize = size_map.values().map(Vec::len).sum(); if total_peers == 0 { continue; // Skip files with no size information } let (consensus_size, consensus_peers) = self.determine_size_consensus(&size_map, total_peers, &relative_path)?; update_peer_scores(&consensus_peers, &mut peer_whitelist_scores); if let Some((size, peers)) = consensus_size && let Some(file_desc) = self.create_validated_file_description(game_id, &relative_path, size, &peers) { file_peer_map.insert(relative_path.clone(), peers.clone()); validated_files.push(file_desc); } } Ok((validated_files, peer_whitelist_scores, file_peer_map)) } /// Determines the consensus size for a file based on peer reports. /// /// # Panics /// /// Panics if `size_map.iter().next()` returns None when `total_peers` == 1 #[allow(clippy::unused_self)] fn determine_size_consensus( &self, size_map: &HashMap>, total_peers: usize, relative_path: &str, ) -> eyre::Result<(ConsensusResult, Vec)> { if total_peers == 1 { // Only one peer has this file - trust it let (&size, peers) = size_map .iter() .next() .expect("size_map should have at least one entry when total_peers == 1"); return Ok((Some((size, peers.clone())), peers.clone())); } let (majority_size, _majority_count) = find_majority_size(size_map); if let Some(size) = majority_size { let majority_peers = &size_map[&size]; let is_majority = majority_peers.len() > total_peers / 2; if is_majority { // We have a clear majority Ok((Some((size, majority_peers.clone())), majority_peers.clone())) } else if total_peers == 2 { // Two peers with different sizes - ambiguous, fail eyre::bail!( "File size ambiguity for '{}': two peers report different sizes, cannot determine majority", relative_path ); } // If no majority and more than 2 peers, we fall back to plurality (largest group) else { Ok((Some((size, majority_peers.clone())), majority_peers.clone())) } } else { // No clear majority and it's a tie between different sizes if total_peers == 2 { eyre::bail!( "File size ambiguity for '{}': two peers report different sizes, cannot determine majority", relative_path ); } // For more than 2 peers, we could fall back to plurality, but for now let's be strict eyre::bail!( "File size ambiguity for '{}': no clear majority among {} peers", relative_path, total_peers ); } } /// Creates a validated file description from consensus data. fn create_validated_file_description( &self, game_id: &str, relative_path: &str, size: u64, peers: &[SocketAddr], ) -> Option { if let Some(first_peer) = peers.first() && let Some(files) = self .peers .get(first_peer) .and_then(|p| p.files.get(game_id)) && let Some(file_desc) = files .iter() .find(|f| f.relative_path == relative_path && f.size == size) { return Some(file_desc.clone()); } None } /// Returns peers that haven't been seen within the timeout duration. #[must_use] pub fn get_stale_peers(&self, timeout: Duration) -> Vec { self.peers .iter() .filter(|(_, peer)| peer.last_seen.elapsed() > timeout) .map(|(addr, _)| *addr) .collect() } } // ============================================================================= // Type aliases for consensus validation // ============================================================================= /// Type alias for file size mapping: path -> size -> peers type FileSizeMap = HashMap>>; /// Type alias for peer file mapping: peer -> path -> size type PeerFileMap = HashMap>; /// Type alias for consensus result: (size, peers) or None type ConsensusResult = Option<(u64, Vec)>; /// Type alias for the aggregated majority validation result. pub type MajorityValidationResult = ( Vec, Vec, HashMap>, ); /// Type alias for per-file consensus aggregation results. type FileConsensusAggregation = ( Vec, HashMap, HashMap>, ); // ============================================================================= // Helper functions for consensus validation // ============================================================================= /// Collects file sizes from all peers and organizes them by path and size. fn collect_file_sizes( game_files: &[(SocketAddr, Vec)], ) -> (FileSizeMap, PeerFileMap) { let mut file_size_map: FileSizeMap = HashMap::new(); let mut peer_files: PeerFileMap = HashMap::new(); for (peer_addr, files) in game_files { let mut peer_file_sizes = HashMap::new(); for file in files { if !file.is_dir { let size = file.size; file_size_map .entry(file.relative_path.clone()) .or_default() .entry(size) .or_default() .push(*peer_addr); peer_file_sizes.insert(file.relative_path.clone(), size); } } peer_files.insert(*peer_addr, peer_file_sizes); } (file_size_map, peer_files) } /// Finds the majority size from a map of sizes to peer lists. fn find_majority_size(size_map: &HashMap>) -> (Option, usize) { let mut majority_size = None; let mut majority_count = 0; for (&size, peers) in size_map { let count = peers.len(); if count > majority_count { majority_count = count; majority_size = Some(size); } else if count == majority_count { // Tie between different sizes - ambiguous, fail majority_size = None; break; } } (majority_size, majority_count) } /// Updates peer scores based on consensus participation. fn update_peer_scores( peers: &[SocketAddr], peer_whitelist_scores: &mut HashMap, ) { for &peer in peers { *peer_whitelist_scores.entry(peer).or_insert(0) += 1; } } /// Creates a peer whitelist from scores, including peers with the highest scores. fn create_peer_whitelist(peer_scores: HashMap) -> Vec { if peer_scores.is_empty() { return Vec::new(); } let mut peers: Vec<_> = peer_scores .into_iter() .filter_map(|(peer, score)| (score > 0).then_some((peer, score))) .collect(); peers.sort_by_key(|(peer, score)| (Reverse(*score), *peer)); peers.into_iter().map(|(peer, _)| peer).collect() }