xlb 0.8.45

Chunked, Bao-verified blob distribution with multi-source concurrent fetch (LAN + peer + edge), app-namespaced
Documentation
//! In-process mock swarm for testing the xlb fetch chain.
//!
//! `MockSwarm` builds an `AssetClass` wired to a set of virtual peers — no
//! real network, no iroh, no QUIC. All blobs live in `HashMap`s.
//!
//! Covered by the test suite below:
//! - Fetch-chain priority (LAN > Swarm > Seed > CDN).
//! - Cache hits on second fetch.
//! - Poisoned-source rejection (hash mismatch → skip source, try next).
//! - `FetchFailed` when no source has the blob.

use std::{collections::HashMap, sync::Arc};

use async_trait::async_trait;
use bytes::Bytes;

use crate::{
    source::{BlobSource, FetchTier},
    AssetClass, AssetClassConfig, BlakeHash, Result,
};

// ─── MockPeer ─────────────────────────────────────────────────────────────────

/// A single virtual peer holding a set of blobs (in-memory).
pub struct MockPeer {
    blobs: HashMap<BlakeHash, Vec<u8>>,
    /// Stand-in `NodeId` reported through [`BlobSource::peer_id`] (R423-T7).
    /// `None` mimics a source with no peer identity, like cache or CDN.
    id: Option<String>,
}

impl Default for MockPeer {
    fn default() -> Self {
        Self::new()
    }
}

impl MockPeer {
    pub fn new() -> Self {
        Self {
            blobs: HashMap::new(),
            id: None,
        }
    }

    /// Give this peer an identity, so a [`crate::FetchReport`] produced by
    /// fetching from it carries a `peer_id`. Without this, per-peer attribution
    /// is untestable — which is how it would ship broken.
    pub fn with_id(mut self, id: impl Into<String>) -> Self {
        self.id = Some(id.into());
        self
    }

    /// Register `data` keyed by its BLAKE3 hash.
    pub fn with_blob(mut self, data: impl Into<Vec<u8>>) -> Self {
        let data = data.into();
        let hash = BlakeHash::hash(&data);
        self.blobs.insert(hash, data);
        self
    }

    /// Register `data` keyed by an explicit `hash` (allows injecting poisoned data).
    pub fn with_blob_at(mut self, hash: BlakeHash, data: impl Into<Vec<u8>>) -> Self {
        self.blobs.insert(hash, data.into());
        self
    }
}

// ─── MockBlobSource (internal) ───────────────────────────────────────────────

struct MockBlobSource {
    tier: FetchTier,
    blobs: Arc<HashMap<BlakeHash, Vec<u8>>>,
    peer_id: Option<String>,
}

#[async_trait]
impl BlobSource for MockBlobSource {
    fn tier(&self) -> FetchTier {
        self.tier
    }

    async fn fetch_raw(&self, hash: &BlakeHash) -> Option<Bytes> {
        self.blobs.get(hash).map(|v| Bytes::from(v.clone()))
    }

    fn peer_id(&self) -> Option<String> {
        self.peer_id.clone()
    }
}

// ─── MockSwarm ────────────────────────────────────────────────────────────────

/// Builder for a test `AssetClass` backed by in-memory peers.
///
/// ```ignore
/// use xlb::{testing::{MockPeer, MockSwarm}, AssetClassConfig, BlakeHash};
///
/// let data = b"hello world";
/// let hash = BlakeHash::hash(data);
///
/// let class = MockSwarm::new()
///     .with_lan_peer(MockPeer::new().with_blob(data.as_ref()))
///     .build_class(AssetClassConfig::default())
///     .await
///     .unwrap();
///
/// let fetched = class.asset(hash).fetch().await.unwrap();
/// assert_eq!(&fetched[..], data);
/// ```
pub struct MockSwarm {
    lan_peers: Vec<MockPeer>,
    swarm_peers: Vec<MockPeer>,
    seeds: Vec<MockPeer>,
    cdn: Option<MockPeer>,
}

impl Default for MockSwarm {
    fn default() -> Self {
        Self::new()
    }
}

impl MockSwarm {
    pub fn new() -> Self {
        Self {
            lan_peers: vec![],
            swarm_peers: vec![],
            seeds: vec![],
            cdn: None,
        }
    }

    pub fn with_lan_peer(mut self, peer: MockPeer) -> Self {
        self.lan_peers.push(peer);
        self
    }

    pub fn with_swarm_peer(mut self, peer: MockPeer) -> Self {
        self.swarm_peers.push(peer);
        self
    }

    pub fn with_seed(mut self, peer: MockPeer) -> Self {
        self.seeds.push(peer);
        self
    }

    pub fn with_cdn(mut self, cdn: MockPeer) -> Self {
        self.cdn = Some(cdn);
        self
    }

    /// Build an `AssetClass` backed by this swarm's virtual peers.
    pub async fn build_class(&self, config: AssetClassConfig) -> Result<AssetClass> {
        let class = AssetClass::register(config).await?;
        self.attach(&class);
        Ok(class)
    }

    /// Attach mock peers to an already-registered class.
    pub fn attach(&self, class: &AssetClass) {
        // Merge all peers at each tier into one source per tier. The merged
        // source reports the first declared id at that tier — enough to prove
        // attribution flows, without pretending one source is several peers.
        let merge = |peers: &[MockPeer]| -> (HashMap<BlakeHash, Vec<u8>>, Option<String>) {
            let mut blobs = HashMap::new();
            let mut id = None;
            for p in peers {
                blobs.extend(p.blobs.clone());
                if id.is_none() {
                    id = p.id.clone();
                }
            }
            (blobs, id)
        };

        for (tier, peers) in [
            (FetchTier::Lan, &self.lan_peers),
            (FetchTier::Swarm, &self.swarm_peers),
            (FetchTier::Seed, &self.seeds),
        ] {
            let (blobs, peer_id) = merge(peers);
            if !blobs.is_empty() {
                class.add_source(Arc::new(MockBlobSource {
                    tier,
                    blobs: Arc::new(blobs),
                    peer_id,
                }));
            }
        }

        if let Some(cdn) = &self.cdn {
            class.add_source(Arc::new(MockBlobSource {
                tier: FetchTier::Cdn,
                blobs: Arc::new(cdn.blobs.clone()),
                // Deliberately not `cdn.id`: the CDN tier is HTTP and has no
                // NodeId in production, so the mock must not be able to prove
                // an attribution the real system cannot produce.
                peer_id: None,
            }));
        }
    }
}

// ─── Tests ────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{AssetClassConfig, BlakeHash, Error};

    async fn lan_class(data: &[u8]) -> (AssetClass, BlakeHash) {
        let hash = BlakeHash::hash(data);
        let class = MockSwarm::new()
            .with_lan_peer(MockPeer::new().with_blob(data))
            .build_class(AssetClassConfig::default())
            .await
            .unwrap();
        (class, hash)
    }

    #[tokio::test]
    async fn fetch_from_lan_peer() {
        let data = b"hello from lan";
        let (class, hash) = lan_class(data).await;
        let fetched = class.asset(hash).fetch().await.unwrap();
        assert_eq!(&fetched[..], data);
    }

    #[tokio::test]
    async fn cache_hit_on_second_fetch() {
        let data = b"once and cache";
        let (class, hash) = lan_class(data).await;
        let asset = class.asset(hash);

        assert!(
            !asset.is_cached().await,
            "should not be cached before first fetch"
        );
        let _ = asset.fetch().await.unwrap();
        assert!(
            asset.is_cached().await,
            "should be cached after first fetch"
        );

        let fetched = asset.fetch().await.unwrap();
        assert_eq!(&fetched[..], data);
    }

    #[tokio::test]
    async fn fetch_failed_when_no_source_has_blob() {
        let hash = BlakeHash::hash(b"orphan blob");
        let class = MockSwarm::new()
            .build_class(AssetClassConfig::default())
            .await
            .unwrap();
        let err = class.asset(hash).fetch().await.unwrap_err();
        assert!(
            matches!(err, Error::FetchFailed(_)),
            "expected FetchFailed, got {err:?}"
        );
    }

    #[tokio::test]
    async fn falls_through_to_cdn_when_lan_missing() {
        let data = b"cdn fallback data";
        let hash = BlakeHash::hash(data);

        let class = MockSwarm::new()
            .with_lan_peer(MockPeer::new().with_blob(b"unrelated lan blob"))
            .with_cdn(MockPeer::new().with_blob_at(hash, data.to_vec()))
            .build_class(AssetClassConfig::default())
            .await
            .unwrap();

        let fetched = class.asset(hash).fetch().await.unwrap();
        assert_eq!(&fetched[..], data);
    }

    #[tokio::test]
    async fn lan_wins_over_cdn_when_both_have_data() {
        let lan_data = b"lan data";
        let cdn_data = b"cdn data";
        let lan_hash = BlakeHash::hash(lan_data);
        let cdn_hash = BlakeHash::hash(cdn_data);

        let class = MockSwarm::new()
            .with_lan_peer(MockPeer::new().with_blob(lan_data.as_ref()))
            .with_cdn(MockPeer::new().with_blob(cdn_data.as_ref()))
            .build_class(AssetClassConfig::default())
            .await
            .unwrap();

        // LAN wins when it has the blob.
        let fetched = class.asset(lan_hash).fetch().await.unwrap();
        assert_eq!(&fetched[..], lan_data);

        // CDN serves the blob LAN doesn't have.
        let fetched = class.asset(cdn_hash).fetch().await.unwrap();
        assert_eq!(&fetched[..], cdn_data);
    }

    #[tokio::test]
    async fn priority_lan_then_swarm_then_seed_then_cdn() {
        // Put the correct blob only on LAN; put wrong bytes on lower tiers.
        // Wrong bytes have a different BLAKE3, so they'll be rejected.
        let real = b"authoritative blob";
        let hash = BlakeHash::hash(real);
        let bad = b"poisoned (wrong hash)";

        let class = MockSwarm::new()
            .with_swarm_peer(MockPeer::new().with_blob_at(hash, bad.to_vec()))
            .with_seed(MockPeer::new().with_blob_at(hash, bad.to_vec()))
            .with_cdn(MockPeer::new().with_blob_at(hash, bad.to_vec()))
            .with_lan_peer(MockPeer::new().with_blob_at(hash, real.to_vec()))
            .build_class(AssetClassConfig::default())
            .await
            .unwrap();

        let fetched = class.asset(hash).fetch().await.unwrap();
        assert_eq!(&fetched[..], real, "LAN peer should win");
    }

    #[tokio::test]
    async fn poisoned_source_rejected_falls_through_to_cdn() {
        // LAN serves wrong bytes → hash mismatch → skip → CDN wins.
        let real = b"real payload";
        let hash = BlakeHash::hash(real);
        let poison = b"bad bytes from malicious lan peer";

        let class = MockSwarm::new()
            .with_lan_peer(MockPeer::new().with_blob_at(hash, poison.to_vec()))
            .with_cdn(MockPeer::new().with_blob_at(hash, real.to_vec()))
            .build_class(AssetClassConfig::default())
            .await
            .unwrap();

        let fetched = class.asset(hash).fetch().await.unwrap();
        assert_eq!(&fetched[..], real, "CDN serves after LAN is rejected");
    }

    #[tokio::test]
    async fn blake3_hash_roundtrip() {
        let data = b"roundtrip test data";
        let hash = BlakeHash::hash(data);
        assert!(hash.verify(data));
        assert!(!hash.verify(b"wrong data"));
    }

    #[tokio::test]
    async fn blake3_hex_roundtrip() {
        let hash = BlakeHash::hash(b"hex roundtrip");
        let reparsed = BlakeHash::from_hex(&hash.to_hex()).unwrap();
        assert_eq!(hash, reparsed);
    }

    #[tokio::test]
    async fn multiple_assets_in_one_class() {
        let blobs: &[&[u8]] = &[b"alpha", b"beta", b"gamma"];
        let mut peer = MockPeer::new();
        for b in blobs {
            peer = peer.with_blob(*b);
        }
        let class = MockSwarm::new()
            .with_swarm_peer(peer)
            .build_class(AssetClassConfig::default())
            .await
            .unwrap();

        for data in blobs {
            let hash = BlakeHash::hash(data);
            let fetched = class.asset(hash).fetch().await.unwrap();
            assert_eq!(&fetched[..], *data);
        }
    }
}