lanspread/crates/lanspread-peer/src/content_quarantine.rs

//! Runtime-local quarantine for peers that served invalid catalog content.

use std::{
    collections::HashSet,
    sync::{Arc, RwLock},
};

use lanspread_db::content_manifest::ContentId;
use lanspread_proto::{PeerEndpoint, PeerId};

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
struct QuarantineKey {
    peer_id: PeerId,
    content_id: ContentId,
}

/// In-memory quarantine shared by every transfer in one peer runtime.
///
/// Clones share the same set. Constructing a new value starts empty; quarantine
/// is intentionally not durable trust state.
#[derive(Clone, Debug, Default)]
pub(crate) struct ContentQuarantine {
    quarantined: Arc<RwLock<HashSet<QuarantineKey>>>,
}

impl ContentQuarantine {
    /// Records a typed integrity failure for this peer and exact catalog
    /// content. Callers must not use this for transport or local I/O failures.
    pub(crate) fn record_integrity_failure(
        &self,
        source: &PeerEndpoint,
        content_id: ContentId,
    ) -> bool {
        self.quarantined
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .insert(QuarantineKey {
                peer_id: source.peer_id,
                content_id,
            })
    }

    #[must_use]
    pub(crate) fn is_quarantined(&self, source: &PeerEndpoint, content_id: ContentId) -> bool {
        self.quarantined
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .contains(&QuarantineKey {
                peer_id: source.peer_id,
                content_id,
            })
    }
}

#[cfg(test)]
mod tests {
    use std::net::SocketAddr;

    use super::*;

    fn source(peer_id: &str, port: u16) -> PeerEndpoint {
        PeerEndpoint::new(
            PeerId::from_bytes(*blake3::hash(peer_id.as_bytes()).as_bytes()),
            SocketAddr::from(([127, 0, 0, 1], port)),
        )
    }

    fn content(seed: u8) -> ContentId {
        ContentId::from_bytes([seed; 32])
    }

    #[test]
    fn bad_source_is_quarantined_without_blocking_good_source() {
        let quarantine = ContentQuarantine::default();
        let bad = source("bad", 12000);
        let good = source("good", 12001);
        let content_id = content(1);

        assert!(quarantine.record_integrity_failure(&bad, content_id));
        assert!(quarantine.is_quarantined(&bad, content_id));
        assert!(!quarantine.is_quarantined(&good, content_id));
    }

    #[test]
    fn address_rotation_does_not_escape_peer_content_quarantine() {
        let quarantine = ContentQuarantine::default();
        let original = source("peer", 12000);
        let rotated = source("peer", 22000);
        let content_id = content(2);

        quarantine.record_integrity_failure(&original, content_id);

        assert!(quarantine.is_quarantined(&rotated, content_id));
    }

    #[test]
    fn quarantine_for_one_content_id_does_not_block_another() {
        let quarantine = ContentQuarantine::default();
        let source = source("peer", 12000);

        quarantine.record_integrity_failure(&source, content(3));

        assert!(quarantine.is_quarantined(&source, content(3)));
        assert!(!quarantine.is_quarantined(&source, content(4)));
    }

    #[test]
    fn clones_share_runtime_state_but_a_new_runtime_starts_empty() {
        let runtime = ContentQuarantine::default();
        let runtime_clone = runtime.clone();
        let source = source("peer", 12000);
        let content_id = content(5);

        runtime.record_integrity_failure(&source, content_id);

        assert!(runtime_clone.is_quarantined(&source, content_id));
        assert!(!ContentQuarantine::default().is_quarantined(&source, content_id));
    }
}