varde/varde-daemon/src/store.rs

//! Blob store operations on top of the iroh-blobs persistent fs store.

use std::collections::BTreeSet;
use std::path::{Path, PathBuf};

use anyhow::Context;
use iroh_blobs::format::collection::Collection;
use iroh_blobs::hashseq::HashSeq;
use iroh_blobs::store::{
    EntryStatus, ImportMode, Map, MapEntry, MapMut, ReadableStore, Store as _,
};
use iroh_blobs::util::local_pool::{LocalPool, LocalPoolHandle};
use iroh_blobs::util::progress::IgnoreProgressSender;
use iroh_blobs::{BlobFormat, Hash, HashAndFormat};
use iroh_io::{AsyncSliceReader, AsyncSliceReaderExt};
use tokio::io::AsyncWriteExt;
use tracing::{debug, info};
use varde_proto::MaterializeMode;

use crate::meta::StoredFormat;

/// Operation errors that map onto protocol error codes.
#[derive(Debug, thiserror::Error)]
pub enum OpError {
    #[error("{0}")]
    NotFound(String),
    #[error("{0}")]
    InvalidArgument(String),
    #[error(transparent)]
    Io(#[from] std::io::Error),
    #[error(transparent)]
    Internal(#[from] anyhow::Error),
    #[error("{0}")]
    Unimplemented(String),
}

/// The daemon's blob store: an iroh-blobs fs store rooted at
/// `<store_dir>/blobs`.
///
/// Reads from store entries yield non-`Send` futures, so those operations
/// run on a dedicated [`LocalPool`] — the same pattern iroh-blobs itself
/// uses for its provider and GC tasks.
#[derive(Debug, Clone)]
pub struct BlobStore {
    store: iroh_blobs::store::fs::Store,
    data_dir: PathBuf,
    pool: LocalPoolHandle,
    // Owns the pool threads; dropped when the last clone goes away.
    _pool: std::sync::Arc<LocalPool>,
}

/// Reading a blob entry for a streaming copy happens in chunks this size.
const COPY_CHUNK: u64 = 1024 * 1024;

impl BlobStore {
    pub async fn open(store_dir: &Path) -> anyhow::Result<BlobStore> {
        let blobs_dir = store_dir.join("blobs");
        let store = iroh_blobs::store::fs::Store::load(&blobs_dir)
            .await
            .with_context(|| format!("opening blob store at {}", blobs_dir.display()))?;
        let pool = LocalPool::default();
        Ok(BlobStore {
            store,
            data_dir: blobs_dir.join("data"),
            pool: pool.handle().clone(),
            _pool: std::sync::Arc::new(pool),
        })
    }

    /// Run a non-`Send` store operation on the local pool.
    async fn on_pool<T, F, Fut>(&self, f: F) -> Result<T, OpError>
    where
        F: FnOnce(BlobStore) -> Fut + Send + 'static,
        Fut: std::future::Future<Output = Result<T, OpError>> + 'static,
        T: Send + 'static,
    {
        let this = self.clone();
        self.pool
            .spawn(move || f(this))
            .await
            .map_err(|e| OpError::Internal(anyhow::anyhow!("local pool: {e}")))?
    }

    /// Access to the underlying iroh-blobs store (used by the transfer
    /// layer).
    pub fn inner(&self) -> &iroh_blobs::store::fs::Store {
        &self.store
    }

    /// Import a file or directory. Returns the root hash, total imported
    /// bytes, and the resulting format (`Raw` for files, `HashSeq` for
    /// directories).
    pub async fn add_path(
        &self,
        path: &Path,
        recursive: bool,
    ) -> Result<(Hash, u64, StoredFormat), OpError> {
        if !path.is_absolute() {
            return Err(OpError::InvalidArgument(format!(
                "path must be absolute: {}",
                path.display()
            )));
        }
        let meta = std::fs::metadata(path)?;
        if meta.is_file() {
            let (tag, size) = self.import_one(path.to_owned()).await?;
            info!(hash = %tag.hash().to_hex(), size, "imported file");
            return Ok((*tag.hash(), size, StoredFormat::Raw));
        }
        if !meta.is_dir() {
            return Err(OpError::InvalidArgument(format!(
                "not a file or directory: {}",
                path.display()
            )));
        }
        if !recursive {
            return Err(OpError::InvalidArgument(format!(
                "{} is a directory; pass recursive=true",
                path.display()
            )));
        }

        let mut files = Vec::new();
        collect_files(path, path, &mut files)?;
        if files.is_empty() {
            return Err(OpError::InvalidArgument(format!(
                "directory {} contains no files",
                path.display()
            )));
        }

        // Temp tags keep the children alive until the collection root is
        // stored (and this method's caller records/pins the root).
        let mut children = Vec::new();
        let mut total = 0u64;
        for (name, file_path) in files {
            let (tag, size) = self.import_one(file_path).await?;
            total += size;
            children.push((name, tag));
        }
        let collection: Collection = children
            .iter()
            .map(|(name, tag)| (name.clone(), *tag.hash()))
            .collect();
        let root_tag = collection
            .store(&self.store)
            .await
            .context("storing collection")?;
        let root = *root_tag.hash();
        info!(hash = %root.to_hex(), files = children.len(), total, "imported directory");
        Ok((root, total, StoredFormat::HashSeq))
    }

    async fn import_one(&self, path: PathBuf) -> Result<(iroh_blobs::TempTag, u64), OpError> {
        let (tag, size) = self
            .store
            .import_file(
                path,
                ImportMode::Copy,
                BlobFormat::Raw,
                IgnoreProgressSender::default(),
            )
            .await?;
        Ok((tag, size))
    }

    /// Export `hash` to `dest`. Returns bytes written and whether at least
    /// one file was reflinked.
    pub async fn materialize(
        &self,
        hash: Hash,
        format: StoredFormat,
        dest: &Path,
        mode: MaterializeMode,
    ) -> Result<(u64, bool), OpError> {
        if !dest.is_absolute() {
            return Err(OpError::InvalidArgument(format!(
                "dest must be absolute: {}",
                dest.display()
            )));
        }
        let dest = dest.to_owned();
        self.on_pool(
            move |this| async move { this.materialize_local(hash, format, &dest, mode).await },
        )
        .await
    }

    async fn materialize_local(
        &self,
        hash: Hash,
        format: StoredFormat,
        dest: &Path,
        mode: MaterializeMode,
    ) -> Result<(u64, bool), OpError> {
        let allow_reflink = matches!(mode, MaterializeMode::ReflinkOrCopy);
        match format {
            StoredFormat::Raw => self.export_blob(hash, dest, allow_reflink).await,
            StoredFormat::HashSeq => {
                let collection = Collection::load_db(&self.store, &hash)
                    .await
                    .map_err(|e| OpError::NotFound(format!("loading collection: {e}")))?;
                let mut total = 0u64;
                let mut any_reflink = false;
                tokio::fs::create_dir_all(dest).await?;
                for (name, child) in collection.iter() {
                    let rel = sanitize_collection_name(name)?;
                    let (bytes, reflinked) = self
                        .export_blob(*child, &dest.join(rel), allow_reflink)
                        .await?;
                    total += bytes;
                    any_reflink |= reflinked;
                }
                Ok((total, any_reflink))
            }
        }
    }

    /// Export one blob. Tries `FICLONE` via reflink first (zero-copy on
    /// btrfs/XFS when store and dest share a filesystem), falls back to a
    /// streaming copy. Never hardlinks: store files must stay immutable.
    async fn export_blob(
        &self,
        hash: Hash,
        dest: &Path,
        allow_reflink: bool,
    ) -> Result<(u64, bool), OpError> {
        let entry =
            self.store.get(&hash).await?.ok_or_else(|| {
                OpError::NotFound(format!("{} is not in the store", hash.to_hex()))
            })?;
        if !entry.is_complete() {
            return Err(OpError::NotFound(format!(
                "{} is only partially present",
                hash.to_hex()
            )));
        }
        let size = entry.size().value();
        if let Some(parent) = dest.parent() {
            tokio::fs::create_dir_all(parent).await?;
        }
        // Materialize replaces the destination, it never appends.
        match tokio::fs::remove_file(dest).await {
            Ok(()) => {}
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
            Err(e) => return Err(e.into()),
        }

        if allow_reflink {
            // Complete blobs above the inline threshold live as plain
            // files at data/<hex>.data; those we can clone directly.
            // Inline (small) blobs have no backing file and are copied.
            let src = self.data_dir.join(format!("{}.data", hash.to_hex()));
            let src_ok = std::fs::metadata(&src)
                .map(|m| m.len() == size)
                .unwrap_or(false);
            if src_ok {
                match reflink_copy::reflink(&src, dest) {
                    Ok(()) => {
                        debug!(hash = %hash.to_hex(), dest = %dest.display(), "reflinked");
                        return Ok((size, true));
                    }
                    Err(e) => {
                        debug!(error = %e, "reflink failed, falling back to copy");
                    }
                }
            }
        }

        let mut reader = entry.data_reader();
        let mut file = tokio::fs::File::create(dest).await?;
        let mut offset = 0u64;
        while offset < size {
            let len = (size - offset).min(COPY_CHUNK) as usize;
            let chunk = reader.read_at(offset, len).await?;
            if chunk.is_empty() {
                return Err(OpError::Io(std::io::Error::new(
                    std::io::ErrorKind::UnexpectedEof,
                    format!("blob {} truncated at {offset}", hash.to_hex()),
                )));
            }
            file.write_all(&chunk).await?;
            offset += chunk.len() as u64;
        }
        file.flush().await?;
        Ok((size, false))
    }

    /// Presence of `hash`: (have_bytes, total_bytes, complete).
    pub async fn presence(&self, hash: Hash) -> Result<(u64, Option<u64>, bool), OpError> {
        match self.store.entry_status(&hash).await? {
            EntryStatus::NotFound => Ok((0, None, false)),
            EntryStatus::Partial => {
                let total = self.store.get(&hash).await?.map(|e| e.size().value());
                // Valid-range accounting for partials arrives with the
                // transfer milestone; absence of data is the safe report.
                Ok((0, total, false))
            }
            EntryStatus::Complete => {
                let entry = self
                    .store
                    .get(&hash)
                    .await?
                    .ok_or_else(|| OpError::NotFound(hash.to_hex().to_string()))?;
                let size = entry.size().value();
                Ok((size, Some(size), true))
            }
        }
    }

    /// Drop every blob not reachable from `roots` (mark and sweep).
    /// In-flight imports are protected by their temp tags; tags stored in
    /// the blob database are honored too.
    pub async fn gc(&self, roots: Vec<HashAndFormat>) -> Result<u64, OpError> {
        self.on_pool(move |this| async move { this.gc_local(roots).await })
            .await
    }

    async fn gc_local(&self, roots: Vec<HashAndFormat>) -> Result<u64, OpError> {
        let mut live: BTreeSet<Hash> = BTreeSet::new();
        let mut all_roots = roots;
        all_roots.extend(self.store.temp_tags());
        for item in self.store.tags(None, None).await.context(TAGS_CONTEXT)? {
            let (_name, haf) = item.context(TAGS_CONTEXT)?;
            all_roots.push(haf);
        }

        for HashAndFormat { hash, format } in all_roots {
            if !live.insert(hash) || format.is_raw() {
                continue;
            }
            // HashSeq root: its children are live too. A partial root
            // can't be expanded; its bytes are still protected.
            let Some(entry) = self.store.get(&hash).await? else {
                continue;
            };
            if !entry.is_complete() {
                continue;
            }
            let mut reader = entry.data_reader();
            let bytes = reader.read_to_end().await?;
            let seq = HashSeq::try_from(bytes)
                .map_err(|e| OpError::Internal(anyhow::anyhow!("invalid hashseq: {e}")))?;
            live.extend(seq.iter());
        }

        let mut doomed = Vec::new();
        for hash in self
            .store
            .blobs()
            .await?
            .chain(self.store.partial_blobs().await?)
        {
            let hash = hash?;
            if !live.contains(&hash) {
                doomed.push(hash);
            }
        }
        let removed = doomed.len() as u64;
        if !doomed.is_empty() {
            self.store.delete(doomed).await?;
        }
        info!(removed, live = live.len(), "gc done");
        Ok(removed)
    }

    /// Flush and shut down the store actor.
    pub async fn shutdown(&self) {
        self.store.shutdown().await;
    }
}

const TAGS_CONTEXT: &str = "listing tags";

/// Recursively collect regular files under `dir` as (relative-name, path),
/// sorted for deterministic collection hashes. Symlinks are followed for
/// files; symlinked directories are rejected to avoid cycles.
fn collect_files(root: &Path, dir: &Path, out: &mut Vec<(String, PathBuf)>) -> Result<(), OpError> {
    let mut entries: Vec<_> =
        std::fs::read_dir(dir)?.collect::<Result<Vec<_>, std::io::Error>>()?;
    entries.sort_by_key(|e| e.file_name());
    for entry in entries {
        let path = entry.path();
        let file_type = entry.file_type()?;
        if file_type.is_dir() {
            collect_files(root, &path, out)?;
        } else if file_type.is_symlink() && std::fs::metadata(&path)?.is_dir() {
            return Err(OpError::InvalidArgument(format!(
                "symlinked directory not supported: {}",
                path.display()
            )));
        } else {
            let rel = path
                .strip_prefix(root)
                .expect("walked path is under root")
                .components()
                .map(|c| c.as_os_str().to_string_lossy())
                .collect::<Vec<_>>()
                .join("/");
            out.push((rel, path));
        }
    }
    Ok(())
}

/// Turn a collection entry name into a safe relative path.
fn sanitize_collection_name(name: &str) -> Result<PathBuf, OpError> {
    let mut path = PathBuf::new();
    for comp in name.split('/') {
        if comp.is_empty() || comp == "." || comp == ".." {
            return Err(OpError::InvalidArgument(format!(
                "unsafe name in collection: {name:?}"
            )));
        }
        path.push(comp);
    }
    Ok(path)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn collection_names_are_sanitized() {
        assert!(sanitize_collection_name("a/b.txt").is_ok());
        assert!(sanitize_collection_name("../etc/passwd").is_err());
        assert!(sanitize_collection_name("/abs").is_err());
        assert!(sanitize_collection_name("a//b").is_err());
        assert!(sanitize_collection_name("a/./b").is_err());
    }
}