bee_check/

lib.rs

//! Core check engine for `bee-check`.
//!
//! The CLI in `main.rs` and (in future) the `bee-check-web` SPA both
//! consume the same [`Report`] shape produced here. See `SPEC.md` for
//! the JSON shape.

use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::{Context, Result, anyhow};
use bee::Client;
use bee::manifest::{is_null_address, unmarshal};
use bee::swarm::gsoc::proximity;
use bee::swarm::{BatchId, Reference};
use futures::stream::{FuturesUnordered, StreamExt};
use serde::{Deserialize, Serialize};
use tokio::sync::Semaphore;

const MAX_CHUNKS: usize = 1000;
/// Stamp is "low TTL" below this many seconds (~24h). A re-seed
/// against a stamp below this threshold may not outlive the batch.
const STAMP_LOW_TTL_SECS: i64 = 86_400;

#[derive(Copy, Clone, Debug)]
pub enum OutputFormat {
    Text,
    Json,
}

/// Outcome enum reported per vantage and aggregated for the whole check.
#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum Status {
    Retrievable,
    Unretrievable,
    Partial,
    Error,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Report {
    pub reference: String,
    pub status: Status,
    pub vantages: Vec<VantageResult>,
    /// Populated only when `--per-chunk` was requested.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub chunks: Option<Vec<ChunkProbe>>,
    pub spec_version: u32,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct VantageResult {
    pub bee_url: String,
    /// `None` means the call errored (see `error`).
    pub retrievable: Option<bool>,
    pub elapsed_ms: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
    /// Hex overlay address of the probed Bee, from `GET /addresses`. The
    /// first 2 hex chars are the neighborhood the node sits in.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub overlay: Option<String>,
    /// Bee semver, from `GET /health` (Bee's `version` field).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub bee_version: Option<String>,
    /// Proximity order between the probe node's overlay and the target
    /// reference. Higher = closer to where the chunk is stored. PO 0
    /// means the probe is not in the chunk's neighborhood at all.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub proximity_to_root: Option<u32>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct ChunkProbe {
    pub address: String,
    /// First 2 hex chars — neighborhood the chunk should land in.
    pub neighborhood: String,
    pub per_vantage: BTreeMap<String, ChunkVantage>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct ChunkVantage {
    pub found: bool,
    pub elapsed_ms: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
    /// Proximity between this vantage's overlay and the chunk address.
    /// Unset when the vantage's overlay couldn't be fetched.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub proximity: Option<u32>,
}

pub struct ReseedRequest {
    pub reference: String,
    pub bee_url: String,
    pub batch_id: String,
    pub timeout: Duration,
}

/// Result of the `--reseed --stamp <id>` pre-flight check. Mirrors
/// ipfs-check's "stale records" hint in spirit: surface upstream-data
/// problems before doing the operation.
#[derive(Debug, Serialize, Deserialize)]
pub struct StampStatus {
    pub batch_id: String,
    pub exists: bool,
    pub usable: bool,
    pub batch_ttl: i64,
    /// `usable && exists && batch_ttl >= STAMP_LOW_TTL_SECS` (or
    /// `batch_ttl < 0`, meaning unknown/infinite).
    pub healthy: bool,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub warnings: Vec<String>,
}

const SPEC_VERSION: u32 = 1;

fn parse_reference(s: &str) -> Result<Reference> {
    Reference::from_hex(s).map_err(|e| anyhow!("invalid reference {s}: {e}"))
}

fn make_bee(url: &str, timeout: Duration) -> Result<Client> {
    let http = reqwest::Client::builder()
        .timeout(timeout)
        .build()
        .context("building http client")?;
    Client::with_http_client(url, http).map_err(|e| anyhow!("invalid bee url {url}: {e}"))
}

/// Probe `/stewardship/{ref}` across all vantages in parallel. Alongside
/// the retrievability call, fetches `/addresses` (overlay) and
/// `/health` (Bee version) so the rendered report can surface which
/// neighborhood the probe came from. These ancillary calls are
/// best-effort: failure leaves `overlay` / `bee_version` unset rather
/// than failing the whole vantage.
pub async fn check_multi_vantage(
    reference: &str,
    bees: &[String],
    timeout: Duration,
) -> Result<Report> {
    let r = parse_reference(reference)?;
    let root_bytes = first_32(&r);

    let mut futs = FuturesUnordered::new();
    for bee_url in bees {
        let bee_url = bee_url.clone();
        let r = r.clone();
        futs.push(async move {
            let bee = match make_bee(&bee_url, timeout) {
                Ok(b) => b,
                Err(e) => {
                    return VantageResult {
                        bee_url,
                        retrievable: None,
                        elapsed_ms: 0,
                        error: Some(format!("{e:#}")),
                        overlay: None,
                        bee_version: None,
                        proximity_to_root: None,
                    };
                }
            };
            let started = Instant::now();
            let api = bee.api();
            let debug = bee.debug();
            let (stew_res, addr_res, health_res) = tokio::join!(
                api.is_retrievable(&r),
                debug.addresses(),
                debug.health(),
            );
            let elapsed_ms = started.elapsed().as_millis() as u64;

            let overlay = addr_res.ok().map(|a| a.overlay);
            let bee_version = health_res.ok().map(|h| h.version);
            let proximity_to_root = overlay
                .as_deref()
                .and_then(decode_overlay)
                .map(|o| proximity(&o, &root_bytes));

            match stew_res {
                Ok(ok) => VantageResult {
                    bee_url,
                    retrievable: Some(ok),
                    elapsed_ms,
                    error: None,
                    overlay,
                    bee_version,
                    proximity_to_root,
                },
                Err(e) => VantageResult {
                    bee_url,
                    retrievable: None,
                    elapsed_ms,
                    error: Some(format!("{e}")),
                    overlay,
                    bee_version,
                    proximity_to_root,
                },
            }
        });
    }

    let mut vantages = Vec::with_capacity(bees.len());
    while let Some(v) = futs.next().await {
        vantages.push(v);
    }
    vantages.sort_by(|a, b| a.bee_url.cmp(&b.bee_url));

    let status = aggregate_status(&vantages);
    Ok(Report {
        reference: reference.to_string(),
        status,
        vantages,
        chunks: None,
        spec_version: SPEC_VERSION,
    })
}

/// Strip the optional `0x` prefix, decode hex, return the first 32
/// bytes (overlay/reference length). Returns `None` on malformed hex
/// or short input.
fn decode_overlay(hex: &str) -> Option<[u8; 32]> {
    let s = hex.strip_prefix("0x").unwrap_or(hex);
    if s.len() < 64 {
        return None;
    }
    let mut out = [0u8; 32];
    for (i, b) in out.iter_mut().enumerate() {
        let h = u8::from_str_radix(&s[i * 2..i * 2 + 2], 16).ok()?;
        *b = h;
    }
    Some(out)
}

fn first_32(r: &Reference) -> [u8; 32] {
    let mut out = [0u8; 32];
    out.copy_from_slice(&r.as_bytes()[..32]);
    out
}

fn aggregate_status(vantages: &[VantageResult]) -> Status {
    let total = vantages.len();
    if total == 0 {
        return Status::Error;
    }
    let mut retr = 0usize;
    let mut unret = 0usize;
    let mut err = 0usize;
    for v in vantages {
        match v.retrievable {
            Some(true) => retr += 1,
            Some(false) => unret += 1,
            None => err += 1,
        }
    }
    if err == total {
        Status::Error
    } else if retr == total {
        Status::Retrievable
    } else if retr == 0 && unret + err == total {
        Status::Unretrievable
    } else {
        Status::Partial
    }
}

/// Walk the manifest at `report.reference`, probe each leaf chunk via
/// `GET /chunks/{addr}` from every vantage, and attach the result to
/// the report. When the per-vantage overlay is known (already fetched
/// in [`check_multi_vantage`]), each [`ChunkVantage`] also carries the
/// proximity between that vantage and the chunk.
pub async fn drill_down(
    mut report: Report,
    bees: &[String],
    timeout: Duration,
    concurrency: usize,
) -> Result<Report> {
    let r = parse_reference(&report.reference)?;
    // Use the first vantage as the source-of-truth for manifest walking.
    let walker_bee = bees.first().context("no bee URL for drill-down")?;
    let walker = make_bee(walker_bee, timeout)?;

    let addresses = collect_chunk_addresses(&walker, &r).await?;

    // Map vantage URL → its (already-fetched) overlay bytes, for
    // per-chunk proximity tagging without re-hitting `/addresses`.
    let overlays: BTreeMap<String, [u8; 32]> = report
        .vantages
        .iter()
        .filter_map(|v| {
            v.overlay
                .as_deref()
                .and_then(decode_overlay)
                .map(|o| (v.bee_url.clone(), o))
        })
        .collect();

    let clients: Vec<(String, Client)> = bees
        .iter()
        .map(|u| make_bee(u, timeout).map(|b| (u.clone(), b)))
        .collect::<Result<_>>()?;

    let sem = Arc::new(Semaphore::new(concurrency.max(1)));
    let mut probes = Vec::with_capacity(addresses.len());

    let mut futs = FuturesUnordered::new();
    for addr in addresses {
        let sem = sem.clone();
        let clients = clients.clone();
        let overlays = overlays.clone();
        futs.push(async move {
            let chunk_bytes = first_32_of_ref(&addr);
            let mut per_vantage = BTreeMap::new();
            for (url, bee) in &clients {
                let _permit = sem.acquire().await.expect("semaphore not closed");
                let started = Instant::now();
                let res = bee.file().download_chunk(&addr, None).await;
                let elapsed_ms = started.elapsed().as_millis() as u64;
                let prox = overlays
                    .get(url)
                    .map(|o| proximity(o, &chunk_bytes));
                let cv = match res {
                    Ok(_) => ChunkVantage {
                        found: true,
                        elapsed_ms,
                        error: None,
                        proximity: prox,
                    },
                    Err(e) => ChunkVantage {
                        found: false,
                        elapsed_ms,
                        error: Some(format!("{e}")),
                        proximity: prox,
                    },
                };
                per_vantage.insert(url.clone(), cv);
            }
            let hex = addr.to_hex();
            let neighborhood = hex.chars().take(2).collect::<String>();
            ChunkProbe {
                address: hex,
                neighborhood,
                per_vantage,
            }
        });
    }

    while let Some(p) = futs.next().await {
        probes.push(p);
    }
    probes.sort_by(|a, b| a.address.cmp(&b.address));
    report.chunks = Some(probes);
    Ok(report)
}

fn first_32_of_ref(r: &Reference) -> [u8; 32] {
    let mut out = [0u8; 32];
    out.copy_from_slice(&r.as_bytes()[..32]);
    out
}

/// BFS-walk the manifest starting at `root`, collecting both child
/// manifest chunk addresses and content (`target_address`) addresses.
/// Capped at [`MAX_CHUNKS`] to bound work for pathological cases. If
/// `root` isn't a manifest, returns just `[root]`.
async fn collect_chunk_addresses(bee: &Client, root: &Reference) -> Result<Vec<Reference>> {
    let mut addresses: Vec<Reference> = vec![root.clone()];
    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
    seen.insert(root.to_hex());

    let mut queue: std::collections::VecDeque<Reference> = std::collections::VecDeque::new();
    queue.push_back(root.clone());

    while let Some(addr) = queue.pop_front() {
        if addresses.len() >= MAX_CHUNKS {
            break;
        }
        let bytes = match bee.file().download_chunk(&addr, None).await {
            Ok(b) => b,
            // If we can't fetch a manifest node, we can't walk deeper. The
            // outer probe loop will still record the chunk-level miss.
            Err(_) => continue,
        };
        let Ok(node) = unmarshal(&bytes, addr.as_bytes()) else {
            // Root might be raw content; deeper nodes that don't parse
            // are leaves — both fine to skip.
            continue;
        };
        // Content reference at this node — probe it but don't recurse.
        if !is_null_address(&node.target_address) {
            if let Ok(r) = Reference::new(&node.target_address) {
                if seen.insert(r.to_hex()) {
                    addresses.push(r);
                }
            }
        }
        // Child manifest chunks — probe and recurse.
        for fork in node.forks.values() {
            if let Some(sa) = fork.node.self_address {
                if let Ok(r) = Reference::new(&sa) {
                    if seen.insert(r.to_hex()) {
                        addresses.push(r.clone());
                        queue.push_back(r);
                    }
                }
            }
        }
    }
    Ok(addresses)
}

pub async fn reseed(req: ReseedRequest) -> Result<()> {
    let bee = make_bee(&req.bee_url, req.timeout)?;
    let r = parse_reference(&req.reference)?;
    let batch = BatchId::from_hex(&req.batch_id)
        .map_err(|e| anyhow!("invalid batch id {}: {e}", req.batch_id))?;
    bee.api().reupload(&r, &batch).await?;
    Ok(())
}

/// Pre-flight check before `--reseed`: look up `GET /stamps/{id}` on the
/// target Bee and surface usable/expiry concerns. Mirrors the spirit of
/// ipfs-check's "stale records" UX hint — flag freshness problems
/// before doing the operation.
pub async fn check_stamp(
    bee_url: &str,
    batch_id: &str,
    timeout: Duration,
) -> Result<StampStatus> {
    let bee = make_bee(bee_url, timeout)?;
    let batch = BatchId::from_hex(batch_id)
        .map_err(|e| anyhow!("invalid batch id {batch_id}: {e}"))?;
    let pb = bee
        .postage()
        .get_postage_batch(&batch)
        .await
        .map_err(anyhow::Error::from)?;

    let mut warnings = Vec::new();
    if !pb.exists {
        warnings.push("batch not known to this Bee".to_string());
    }
    if !pb.usable {
        warnings.push("batch not usable yet (chain may be syncing)".to_string());
    }
    if pb.batch_ttl >= 0 && pb.batch_ttl < STAMP_LOW_TTL_SECS {
        warnings.push(format!(
            "batch TTL low: ~{} (re-seed may not outlive the batch)",
            humanize_secs(pb.batch_ttl)
        ));
    }

    let healthy = pb.exists && pb.usable && (pb.batch_ttl < 0 || pb.batch_ttl >= STAMP_LOW_TTL_SECS);

    Ok(StampStatus {
        batch_id: batch_id.to_string(),
        exists: pb.exists,
        usable: pb.usable,
        batch_ttl: pb.batch_ttl,
        healthy,
        warnings,
    })
}

fn humanize_secs(s: i64) -> String {
    if s < 0 {
        return "unknown".to_string();
    }
    let s = s as u64;
    if s >= 86_400 {
        format!("{} day(s)", s / 86_400)
    } else if s >= 3_600 {
        format!("{} hour(s)", s / 3_600)
    } else if s >= 60 {
        format!("{} min", s / 60)
    } else {
        format!("{}s", s)
    }
}

pub fn render_report(report: &Report, fmt: OutputFormat) -> String {
    match fmt {
        OutputFormat::Json => {
            serde_json::to_string_pretty(report).expect("report serialization") + "\n"
        }
        OutputFormat::Text => render_text(report),
    }
}

fn render_text(r: &Report) -> String {
    use std::fmt::Write;
    let mut out = String::new();
    let _ = writeln!(out, "ref     {}", r.reference);
    let _ = writeln!(out, "status  {:?}", r.status);
    let _ = writeln!(out);
    let _ = writeln!(out, "vantages:");
    let url_w = r.vantages.iter().map(|v| v.bee_url.len()).max().unwrap_or(20);
    for v in &r.vantages {
        let state = match (v.retrievable, &v.error) {
            (Some(true), _) => "retrievable",
            (Some(false), _) => "unretrievable",
            (None, Some(_)) => "error",
            (None, None) => "unknown",
        };
        let meta = vantage_meta(v);
        let _ = writeln!(
            out,
            "  {:<url_w$}  {:<14} {:>6} ms{}{}",
            v.bee_url,
            state,
            v.elapsed_ms,
            if meta.is_empty() { String::new() } else { format!("  {meta}") },
            v.error
                .as_deref()
                .map(|e| format!("  ({e})"))
                .unwrap_or_default(),
            url_w = url_w
        );
    }
    if let Some(chunks) = &r.chunks {
        let _ = writeln!(out);
        let _ = writeln!(out, "chunks: {} probed", chunks.len());
        let mut missing = 0usize;
        for c in chunks {
            let missing_in: Vec<String> = c
                .per_vantage
                .iter()
                .filter(|(_, cv)| !cv.found)
                .map(|(u, cv)| match cv.proximity {
                    Some(p) => format!("{u} (PO {p})"),
                    None => u.clone(),
                })
                .collect();
            if !missing_in.is_empty() {
                missing += 1;
                let _ = writeln!(
                    out,
                    "  [{}] {}  missing on: {}",
                    c.neighborhood,
                    short(&c.address),
                    missing_in.join(", ")
                );
            }
        }
        if missing == 0 {
            let _ = writeln!(out, "  all chunks present on all vantages");
        } else {
            let _ = writeln!(out, "  {missing} chunk(s) missing on at least one vantage");
        }
    }
    out
}

/// Single-line metadata trailer for a vantage: overlay neighborhood,
/// proximity to root, Bee version. Compactly formatted; pieces that
/// weren't fetched are silently dropped.
fn vantage_meta(v: &VantageResult) -> String {
    let mut parts: Vec<String> = Vec::new();
    if let Some(o) = &v.overlay {
        let neigh = o.chars().take(2).collect::<String>();
        let short_overlay = short_overlay(o);
        parts.push(format!("overlay {short_overlay} (nb {neigh})"));
    }
    if let Some(p) = v.proximity_to_root {
        parts.push(format!("PO {p}"));
    }
    if let Some(ver) = &v.bee_version {
        parts.push(format!("v{ver}"));
    }
    if parts.is_empty() {
        String::new()
    } else {
        format!("· {}", parts.join(" · "))
    }
}

fn short_overlay(hex: &str) -> String {
    let s = hex.strip_prefix("0x").unwrap_or(hex);
    if s.len() > 12 {
        format!("{}…{}", &s[..6], &s[s.len() - 4..])
    } else {
        s.to_string()
    }
}

/// Human-readable summary of a stamp pre-flight check, suitable for
/// stderr before a `--reseed` operation.
pub fn render_stamp_status(s: &StampStatus) -> String {
    use std::fmt::Write;
    let mut out = String::new();
    let ttl = if s.batch_ttl < 0 {
        "unknown".to_string()
    } else {
        humanize_secs(s.batch_ttl)
    };
    let header = if s.healthy { "stamp OK" } else { "stamp warning" };
    let _ = writeln!(
        out,
        "{header}: batch {} · usable={} · ttl={}",
        short_overlay(&s.batch_id),
        s.usable,
        ttl,
    );
    for w in &s.warnings {
        let _ = writeln!(out, "  · {w}");
    }
    out
}

fn short(hex: &str) -> String {
    if hex.len() > 16 {
        format!("{}…{}", &hex[..8], &hex[hex.len() - 4..])
    } else {
        hex.to_string()
    }
}