use libp2p::gossipsub::{IdentTopic, TopicHash};
use std::collections::{HashMap, VecDeque};
use std::time::Instant;
pub struct GossipTopics {
topics: HashMap<String, IdentTopic>,
}
impl GossipTopics {
pub fn new() -> Self {
let mut topics = HashMap::new();
topics.insert("blocks".to_string(), IdentTopic::new("tenzro/blocks"));
topics.insert("transactions".to_string(), IdentTopic::new("tenzro/transactions"));
topics.insert("consensus".to_string(), IdentTopic::new("tenzro/consensus"));
topics.insert("attestations".to_string(), IdentTopic::new("tenzro/attestations"));
topics.insert("models".to_string(), IdentTopic::new("tenzro/models"));
topics.insert("inference".to_string(), IdentTopic::new("tenzro/inference"));
topics.insert("status".to_string(), IdentTopic::new("tenzro/status"));
topics.insert("agents".to_string(), IdentTopic::new("tenzro/agents"));
topics.insert("providers".to_string(), IdentTopic::new("tenzro/providers"));
topics.insert("cortex".to_string(), IdentTopic::new("tenzro/cortex"));
topics.insert("training".to_string(), IdentTopic::new("tenzro/training"));
topics.insert("training_syncer".to_string(), IdentTopic::new("tenzro/training/syncer"));
topics.insert("seed_agents".to_string(), IdentTopic::new("tenzro/seed-agents"));
Self { topics }
}
pub fn get(&self, name: &str) -> Option<&IdentTopic> {
self.topics.get(name)
}
pub fn all_hashes(&self) -> Vec<TopicHash> {
self.topics.values().map(|t| t.hash()).collect()
}
pub fn all(&self) -> Vec<&IdentTopic> {
self.topics.values().collect()
}
pub fn add_custom(&mut self, name: String, topic: IdentTopic) {
self.topics.insert(name, topic);
}
pub fn blocks(&self) -> &IdentTopic {
self.topics.get("blocks").expect("blocks topic must exist")
}
pub fn transactions(&self) -> &IdentTopic {
self.topics.get("transactions").expect("transactions topic must exist")
}
pub fn consensus(&self) -> &IdentTopic {
self.topics.get("consensus").expect("consensus topic must exist")
}
pub fn attestations(&self) -> &IdentTopic {
self.topics.get("attestations").expect("attestations topic must exist")
}
pub fn models(&self) -> &IdentTopic {
self.topics.get("models").expect("models topic must exist")
}
pub fn inference(&self) -> &IdentTopic {
self.topics.get("inference").expect("inference topic must exist")
}
pub fn status(&self) -> &IdentTopic {
self.topics.get("status").expect("status topic must exist")
}
pub fn agents(&self) -> &IdentTopic {
self.topics.get("agents").expect("agents topic must exist")
}
pub fn providers(&self) -> &IdentTopic {
self.topics.get("providers").expect("providers topic must exist")
}
pub fn cortex(&self) -> &IdentTopic {
self.topics.get("cortex").expect("cortex topic must exist")
}
pub fn seed_agents(&self) -> &IdentTopic {
self.topics.get("seed_agents").expect("seed_agents topic must exist")
}
}
impl Default for GossipTopics {
fn default() -> Self {
Self::new()
}
}
pub struct TopicSubscriptions {
subscribed: HashMap<TopicHash, bool>,
}
impl TopicSubscriptions {
pub fn new() -> Self {
Self {
subscribed: HashMap::new(),
}
}
pub fn subscribe(&mut self, topic: &TopicHash) {
self.subscribed.insert(topic.clone(), true);
}
pub fn unsubscribe(&mut self, topic: &TopicHash) {
self.subscribed.remove(topic);
}
pub fn is_subscribed(&self, topic: &TopicHash) -> bool {
self.subscribed.get(topic).copied().unwrap_or(false)
}
pub fn subscribed_topics(&self) -> Vec<&TopicHash> {
self.subscribed.keys().collect()
}
pub fn count(&self) -> usize {
self.subscribed.len()
}
}
impl Default for TopicSubscriptions {
fn default() -> Self {
Self::new()
}
}
pub struct MessageDeduplicator {
seen: VecDeque<([u8; 32], Instant)>,
seen_set: HashMap<[u8; 32], ()>,
max_entries: usize,
expiry_secs: u64,
}
impl MessageDeduplicator {
pub fn new(max_entries: usize, expiry_secs: u64) -> Self {
Self {
seen: VecDeque::with_capacity(max_entries),
seen_set: HashMap::with_capacity(max_entries),
max_entries,
expiry_secs,
}
}
pub fn is_duplicate(&mut self, data: &[u8]) -> bool {
use sha2::{Digest, Sha256};
let hash: [u8; 32] = Sha256::digest(data).into();
let now = Instant::now();
while let Some(&(_, ts)) = self.seen.front() {
if now.duration_since(ts).as_secs() > self.expiry_secs {
if let Some((old_hash, _)) = self.seen.pop_front() {
self.seen_set.remove(&old_hash);
}
} else {
break;
}
}
if self.seen.len() >= self.max_entries
&& let Some((old_hash, _)) = self.seen.pop_front()
{
self.seen_set.remove(&old_hash);
}
if self.seen_set.contains_key(&hash) {
return true;
}
self.seen_set.insert(hash, ());
self.seen.push_back((hash, now));
false
}
pub fn len(&self) -> usize {
self.seen.len()
}
pub fn is_empty(&self) -> bool {
self.seen.is_empty()
}
}
impl Default for MessageDeduplicator {
fn default() -> Self {
Self::new(100_000, 300)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MessageValidation {
Accept,
Reject,
Ignore,
}
pub fn validate_gossip_message(_topic: &TopicHash, data: &[u8]) -> MessageValidation {
if data.is_empty() {
return MessageValidation::Reject;
}
if data.len() > 1_048_576 {
return MessageValidation::Reject;
}
match crate::message::NetworkMessage::from_bytes(data) {
Ok(msg) => {
if let Err(e) = crate::message::validate_message(&msg) {
tracing::warn!("Message validation failed: {}", e);
return MessageValidation::Reject;
}
MessageValidation::Accept
}
Err(e) => {
tracing::warn!("Failed to parse network message: {}", e);
MessageValidation::Reject
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_gossip_topics() {
let topics = GossipTopics::new();
assert!(topics.get("blocks").is_some());
assert!(topics.get("transactions").is_some());
assert!(topics.get("consensus").is_some());
assert!(topics.get("nonexistent").is_none());
}
#[test]
fn test_topic_subscriptions() {
let mut subs = TopicSubscriptions::new();
let topics = GossipTopics::new();
let blocks_hash = topics.blocks().hash();
assert!(!subs.is_subscribed(&blocks_hash));
subs.subscribe(&blocks_hash);
assert!(subs.is_subscribed(&blocks_hash));
assert_eq!(subs.count(), 1);
subs.unsubscribe(&blocks_hash);
assert!(!subs.is_subscribed(&blocks_hash));
assert_eq!(subs.count(), 0);
}
#[test]
fn test_message_deduplication() {
let mut dedup = MessageDeduplicator::new(100, 60);
let msg1 = b"message one";
let msg2 = b"message two";
assert!(!dedup.is_duplicate(msg1));
assert!(dedup.is_duplicate(msg1));
assert!(!dedup.is_duplicate(msg2));
assert_eq!(dedup.len(), 2);
}
#[test]
fn test_dedup_capacity_eviction() {
let mut dedup = MessageDeduplicator::new(3, 300);
assert!(!dedup.is_duplicate(b"a"));
assert!(!dedup.is_duplicate(b"b"));
assert!(!dedup.is_duplicate(b"c"));
assert!(!dedup.is_duplicate(b"d"));
assert_eq!(dedup.len(), 3);
assert!(!dedup.is_duplicate(b"a"));
}
#[test]
fn test_message_validation() {
assert_eq!(
validate_gossip_message(&TopicHash::from_raw("test"), &[]),
MessageValidation::Reject
);
let msg = crate::message::NetworkMessage::new(crate::message::MessagePayload::Ping);
let bytes = msg.to_bytes().unwrap();
let topics = GossipTopics::new();
let result = validate_gossip_message(&topics.status().hash(), &bytes);
assert_eq!(result, MessageValidation::Accept);
}
}