varde/varde-daemon/src/daemon.rs

//! Request dispatch and daemon state.

use std::collections::{BTreeMap, BTreeSet};
use std::path::PathBuf;
use std::str::FromStr;
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Duration, Instant};

use anyhow::Context;
use iroh_blobs::{BlobFormat, Hash, HashAndFormat};
use tokio::sync::broadcast;
use tracing::{debug, info, warn};
use varde_proto::{
    Envelope, ErrorCode, Event, GlobalStatus, HashStatus, PeerInfo, PinInfo, Request, ResponseData,
};

use crate::config::Config;
use crate::meta::{Meta, StoredFormat};
use crate::store::{BlobStore, OpError};
use crate::transfer::Transfer;

/// A LAN peer is considered present if discovery has seen it within this
/// window (mdns re-announces well below this).
const PEER_PRESENCE_WINDOW: Duration = Duration::from_secs(120);

/// Suppress duplicate PeerJoined events within this window.
const PEER_REJOIN_WINDOW: Duration = Duration::from_secs(60);

/// Shared daemon state.
pub struct Daemon {
    config: Config,
    store: BlobStore,
    meta: Arc<Meta>,
    transfer: Transfer,
    events: broadcast::Sender<Event>,
    /// Hashes servable to untrusted peers (open_lan pins, exported
    /// tickets, and their hashseq children). Read by the provider filter
    /// on every untrusted request.
    open_set: Arc<RwLock<BTreeSet<Hash>>>,
    /// LAN peers seen via discovery: node id -> last sighting.
    lan_peers: Mutex<BTreeMap<String, Instant>>,
}

impl Daemon {
    /// Open the blob store and metadata under the configured store dir,
    /// bring up the iroh endpoint, and start the discovery/sync loops.
    pub async fn open(config: Config) -> anyhow::Result<Arc<Daemon>> {
        std::fs::create_dir_all(&config.store_dir)
            .with_context(|| format!("creating store directory {}", config.store_dir.display()))?;
        let store = BlobStore::open(&config.store_dir).await?;
        let meta = Arc::new(Meta::load(&config.store_dir)?);
        // Capacity bounds memory if a subscriber stalls; laggards get a
        // Lagged error, not unbounded buffering.
        let (events, _) = broadcast::channel(1024);

        let open_set: Arc<RwLock<BTreeSet<Hash>>> = Arc::default();
        let filter_set = open_set.clone();
        let open_filter: crate::shaped::ServeFilter =
            Arc::new(move |hash| filter_set.read().expect("open set lock").contains(hash));

        let metered = crate::metered::start();
        let transfer = Transfer::start(
            &config,
            &store,
            meta.clone(),
            open_filter,
            store.pool_handle(),
            events.clone(),
            metered,
        )
        .await?;

        let daemon = Arc::new(Daemon {
            config,
            store,
            meta,
            transfer,
            events,
            open_set,
            lan_peers: Mutex::new(BTreeMap::new()),
        });
        daemon.recompute_open_set().await;

        // Completed fetches can turn hashseq roots expandable; refresh
        // the open set whenever a pin finishes.
        let weak = Arc::downgrade(&daemon);
        let mut pin_events = daemon.events.subscribe();
        tokio::spawn(async move {
            loop {
                match pin_events.recv().await {
                    Ok(Event::PinComplete { .. }) => {
                        let Some(daemon) = weak.upgrade() else { break };
                        daemon.recompute_open_set().await;
                    }
                    Ok(_) => {}
                    Err(broadcast::error::RecvError::Lagged(_)) => {}
                    Err(broadcast::error::RecvError::Closed) => break,
                }
            }
        });

        // Discovery loop: track LAN peer sightings, emit join events for
        // trusted peers and sync any incomplete pins from them.
        if let Some(mut discovered) = daemon.transfer.take_discovery() {
            let weak = Arc::downgrade(&daemon);
            tokio::spawn(async move {
                while let Some(node_id) = discovered.recv().await {
                    let Some(daemon) = weak.upgrade() else { break };
                    daemon.on_peer_seen(node_id).await;
                }
            });
        }

        Ok(daemon)
    }

    /// Flush the store and close the endpoint.
    pub async fn shutdown(&self) {
        self.transfer.shutdown().await;
        self.store.shutdown().await;
    }

    pub fn config(&self) -> &Config {
        &self.config
    }

    pub fn store(&self) -> &BlobStore {
        &self.store
    }

    /// New receiver for the daemon event stream.
    pub fn subscribe_events(&self) -> broadcast::Receiver<Event> {
        self.events.subscribe()
    }

    /// Rebuild the set of hashes servable to untrusted peers.
    async fn recompute_open_set(&self) {
        let mut open = BTreeSet::new();
        let mut roots: Vec<String> = self
            .meta
            .pins()
            .into_iter()
            .filter(|(_, record)| record.policy.open_lan)
            .map(|(hash, _)| hash)
            .collect();
        roots.extend(self.meta.exported());

        for hex in roots {
            let Ok(hash) = Hash::from_str(&hex) else {
                continue;
            };
            open.insert(hash);
            if self.meta.format_of(&hex) == StoredFormat::HashSeq {
                match self.store.hashseq_children(hash).await {
                    Ok(children) => open.extend(children),
                    Err(e) => warn!(hash = %hex, error = %e, "expanding open hashseq"),
                }
            }
        }
        debug!(hashes = open.len(), "open set recomputed");
        *self.open_set.write().expect("open set lock") = open;
    }

    /// Handle a discovery sighting of `node_id`.
    async fn on_peer_seen(&self, node_id: iroh::NodeId) {
        let id = node_id.to_string();
        let rejoined = {
            let mut peers = self.lan_peers.lock().expect("lan peers lock");
            let now = Instant::now();
            let fresh = match peers.get(&id) {
                Some(last) => now.duration_since(*last) > PEER_REJOIN_WINDOW,
                None => true,
            };
            peers.insert(id.clone(), now);
            fresh
        };
        if !self.meta.is_trusted(&id) {
            return;
        }
        if rejoined {
            info!(peer = %id, "trusted peer present on LAN");
            let _ = self.events.send(Event::PeerJoined { node_id: id });
        }
        self.sync_pins_from(node_id).await;
    }

    /// Queue fetches for every incomplete pin from a trusted LAN peer.
    /// The downloader dedupes by content, so repeated sightings are
    /// cheap; discovery resolves the peer's addresses.
    async fn sync_pins_from(&self, node_id: iroh::NodeId) {
        for (hex, _record) in self.meta.pins() {
            let Ok(hash) = Hash::from_str(&hex) else {
                continue;
            };
            match self.store.presence(hash).await {
                Ok((_, _, true)) => {} // complete, nothing to do
                Ok(_) => {
                    debug!(hash = %hex, peer = %node_id, "syncing incomplete pin");
                    self.transfer.spawn_fetch(
                        HashAndFormat {
                            hash,
                            format: blob_format(self.meta.format_of(&hex)),
                        },
                        vec![iroh::NodeAddr::new(node_id)],
                    );
                }
                Err(e) => warn!(hash = %hex, error = %e, "checking pin presence"),
            }
        }
    }

    /// Trusted peers seen on the LAN recently.
    fn present_trusted_peers(&self) -> Vec<PeerInfo> {
        let peers = self.lan_peers.lock().expect("lan peers lock");
        let now = Instant::now();
        self.meta
            .trusted_peers()
            .into_iter()
            .map(|node_id| {
                let connected = peers
                    .get(&node_id)
                    .is_some_and(|last| now.duration_since(*last) < PEER_PRESENCE_WINDOW);
                PeerInfo { node_id, connected }
            })
            .collect()
    }

    /// Handle one API request. Infallible at this level: all failures are
    /// mapped to structured error envelopes.
    pub async fn handle(&self, request: Request) -> Envelope {
        match self.dispatch(request).await {
            Ok(data) => Envelope::ok(data),
            Err(e) => Envelope::err(error_code(&e), e.to_string()),
        }
    }

    async fn dispatch(&self, request: Request) -> Result<ResponseData, OpError> {
        match request {
            Request::Add { path, recursive } => {
                let (hash, bytes, format) =
                    self.store.add_path(&PathBuf::from(path), recursive).await?;
                self.meta.record_format(&hash.to_hex(), format)?;
                Ok(ResponseData::Added {
                    hash: hash.to_hex().to_string(),
                    bytes,
                })
            }
            Request::Pin { hash, policy } => {
                let parsed = parse_hash(&hash)?;
                self.meta.pin(&parsed.to_hex(), policy)?;
                self.recompute_open_set().await;
                // Fetch if absent: try every trusted peer currently
                // present on the LAN.
                if let Ok((_, _, false)) = self.store.presence(parsed).await {
                    let present: Vec<iroh::NodeAddr> = self
                        .present_trusted_peers()
                        .into_iter()
                        .filter(|p| p.connected)
                        .filter_map(|p| p.node_id.parse().ok().map(iroh::NodeAddr::new))
                        .collect();
                    if !present.is_empty() {
                        self.transfer.spawn_fetch(
                            HashAndFormat {
                                hash: parsed,
                                format: blob_format(self.meta.format_of(&parsed.to_hex())),
                            },
                            present,
                        );
                    }
                }
                Ok(ResponseData::Pinned {
                    hash: parsed.to_hex().to_string(),
                })
            }
            Request::Unpin { hash } => {
                let parsed = parse_hash(&hash)?;
                if !self.meta.unpin(&parsed.to_hex())? {
                    return Err(OpError::NotFound(format!("no pin for {}", parsed.to_hex())));
                }
                self.recompute_open_set().await;
                Ok(ResponseData::Done {})
            }
            Request::Materialize { hash, dest, mode } => {
                let parsed = parse_hash(&hash)?;
                let format = self.meta.format_of(&parsed.to_hex());
                let (bytes, reflinked) = self
                    .store
                    .materialize(parsed, format, &PathBuf::from(dest), mode)
                    .await?;
                Ok(ResponseData::Materialized { bytes, reflinked })
            }
            Request::Status { hash: None } => {
                let connected = self
                    .present_trusted_peers()
                    .iter()
                    .filter(|p| p.connected)
                    .count() as u64;
                Ok(ResponseData::Status(GlobalStatus {
                    version: env!("CARGO_PKG_VERSION").to_string(),
                    protocol: varde_proto::PROTOCOL_VERSION,
                    node_id: Some(self.transfer.node_id()),
                    store_path: self.config.store_dir.display().to_string(),
                    pins: self.meta.pins().len() as u64,
                    connected_peers: connected,
                }))
            }
            Request::Status { hash: Some(hash) } => {
                let parsed = parse_hash(&hash)?;
                let (have_bytes, total_bytes, complete) = self.store.presence(parsed).await?;
                Ok(ResponseData::HashStatus(HashStatus {
                    hash: parsed.to_hex().to_string(),
                    have_bytes,
                    total_bytes,
                    complete,
                    pinned: self.meta.is_pinned(&parsed.to_hex()),
                    active_peers: Vec::new(),
                }))
            }
            Request::List {} => {
                let mut pins = Vec::new();
                for (hash, record) in self.meta.pins() {
                    let parsed = parse_hash(&hash)?;
                    let (have_bytes, total_bytes, complete) = self.store.presence(parsed).await?;
                    pins.push(PinInfo {
                        hash,
                        policy: record.policy,
                        complete,
                        have_bytes,
                        total_bytes,
                    });
                }
                Ok(ResponseData::Pins { pins })
            }
            Request::Gc {} => {
                let roots = self.pin_roots()?;
                let blobs_removed = self.store.gc(roots).await?;
                Ok(ResponseData::GcDone { blobs_removed })
            }
            Request::TicketExport { hash } => {
                let parsed = parse_hash(&hash)?;
                let (_have, _total, complete) = self.store.presence(parsed).await?;
                if !complete {
                    return Err(OpError::NotFound(format!(
                        "{} is not fully present; cannot export a ticket",
                        parsed.to_hex()
                    )));
                }
                let content = HashAndFormat {
                    hash: parsed,
                    format: blob_format(self.meta.format_of(&parsed.to_hex())),
                };
                let ticket = self.transfer.export_ticket(content).await?;
                // Exporting a ticket is deliberate publication: the
                // holder must be able to fetch without being trusted.
                self.meta.record_exported(&parsed.to_hex())?;
                self.recompute_open_set().await;
                Ok(ResponseData::Ticket { ticket })
            }
            Request::TicketImport { ticket, pin_policy } => {
                let (hash, format, provider) =
                    crate::transfer::parse_ticket(&ticket).map_err(OpError::InvalidArgument)?;
                let stored = match format {
                    BlobFormat::Raw => StoredFormat::Raw,
                    BlobFormat::HashSeq => StoredFormat::HashSeq,
                };
                self.meta.record_format(&hash.to_hex(), stored)?;
                // Pin first: the pin is the GC root that keeps the fetch's
                // data alive, and it is the standing intent even if the
                // provider is unreachable right now.
                self.meta.pin(&hash.to_hex(), pin_policy)?;
                self.recompute_open_set().await;
                self.transfer
                    .spawn_fetch(HashAndFormat { hash, format }, vec![provider]);
                Ok(ResponseData::Pinned {
                    hash: hash.to_hex().to_string(),
                })
            }
            Request::PeerTrust { node_id } => {
                let parsed: iroh::NodeId = node_id
                    .parse()
                    .map_err(|e| OpError::InvalidArgument(format!("invalid node id: {e}")))?;
                self.meta.trust_peer(&parsed.to_string())?;
                info!(peer = %parsed, "peer trusted");
                // If the peer is already visible on the LAN, sync now.
                let present = {
                    let peers = self.lan_peers.lock().expect("lan peers lock");
                    peers.get(&parsed.to_string()).is_some_and(|last| {
                        Instant::now().duration_since(*last) < PEER_PRESENCE_WINDOW
                    })
                };
                if present {
                    self.sync_pins_from(parsed).await;
                }
                Ok(ResponseData::Done {})
            }
            Request::PeerUntrust { node_id } => {
                let parsed: iroh::NodeId = node_id
                    .parse()
                    .map_err(|e| OpError::InvalidArgument(format!("invalid node id: {e}")))?;
                if !self.meta.untrust_peer(&parsed.to_string())? {
                    return Err(OpError::NotFound(format!("{parsed} was not trusted")));
                }
                Ok(ResponseData::Done {})
            }
            Request::PeerList {} => Ok(ResponseData::Peers {
                peers: self.present_trusted_peers(),
            }),
            Request::Subscribe {} => Ok(ResponseData::Done {}),
        }
    }

    /// The GC roots: every pinned hash with its recorded format.
    fn pin_roots(&self) -> Result<Vec<HashAndFormat>, OpError> {
        let mut roots = Vec::new();
        for (hash, _record) in self.meta.pins() {
            let parsed = parse_hash(&hash)?;
            roots.push(HashAndFormat {
                hash: parsed,
                format: blob_format(self.meta.format_of(&hash)),
            });
        }
        Ok(roots)
    }
}

fn blob_format(format: StoredFormat) -> BlobFormat {
    match format {
        StoredFormat::Raw => BlobFormat::Raw,
        StoredFormat::HashSeq => BlobFormat::HashSeq,
    }
}

fn parse_hash(s: &str) -> Result<Hash, OpError> {
    Hash::from_str(s).map_err(|e| OpError::InvalidArgument(format!("invalid hash {s:?}: {e}")))
}

fn error_code(e: &OpError) -> ErrorCode {
    match e {
        OpError::NotFound(_) => ErrorCode::NotFound,
        OpError::InvalidArgument(_) => ErrorCode::InvalidArgument,
        OpError::Io(_) => ErrorCode::Io,
        OpError::Internal(_) => ErrorCode::Internal,
        OpError::Unimplemented(_) => ErrorCode::Unimplemented,
    }
}