bee-check 0.5.2

//! `bee-check` — retrievability checker for Swarm references.
//!
//! Multi-vantage probe of `GET /stewardship/{ref}` across one or more
//! Bee nodes, with optional per-chunk drill-down and re-seed.

use anyhow::{Context, Result};
use clap::{Parser, ValueEnum};

use bee_check::{
    DEFAULT_GATEWAY, OutputFormat, ParsedInput, ReseedRequest, annotate_target_overlay,
    check_gateways, check_multi_vantage, check_stamp, cold_download_all, drill_down,
    merge_gateways, parse_input, render_report, render_stamp_status, reseed, resolve_feed,
};

#[derive(Parser, Debug)]
#[command(
    name = "bee-check",
    version,
    about = "Retrievability checker for Swarm references",
    long_about = "Probe one or more Bee nodes to determine whether a Swarm \
                  reference is retrievable from the network. Supports \
                  per-chunk drill-down and one-shot re-seed via stewardship."
)]
struct Cli {
    /// Swarm reference (64- or 128-hex) or feed reference
    /// `feed:OWNER:TOPIC` (40-hex owner, 64-hex topic). Feed inputs are
    /// resolved via `GET /feeds/{owner}/{topic}` on the first --bee
    /// before probing.
    #[arg(value_name = "INPUT")]
    reference: String,

    /// Bee API URL(s) to probe. Repeat for multi-vantage. Defaults to
    /// $BEE_API_URL or http://localhost:1633.
    #[arg(short = 'b', long = "bee", value_name = "URL")]
    bee: Vec<String>,

    /// Public gateway URL(s) to HEAD-probe via `{gw}/bzz/{ref}/`.
    /// Repeat for multiple. Default: api.gateway.ethswarm.org unless
    /// --no-gateway is set.
    #[arg(long = "gateway", value_name = "URL", conflicts_with = "no_gateway")]
    gateway: Vec<String>,

    /// Skip public-gateway probing entirely.
    #[arg(long)]
    no_gateway: bool,

    /// Compute and display the proximity order between each vantage
    /// and this target overlay (hex). Vantages are re-sorted closest-
    /// first. Useful for "from a node near neighborhood X, is this
    /// retrievable?" questions.
    #[arg(long, value_name = "HEX")]
    target_overlay: Option<String>,

    /// Walk the manifest and probe each chunk per vantage.
    #[arg(long)]
    per_chunk: bool,

    /// Cold end-to-end download probe. For each --bee, issues
    /// `GET /bytes/{ref}` and streams the body to EOF. For each
    /// --gateway, issues `GET /bzz/{ref}/`. Records bytes_downloaded
    /// + elapsed_ms. Complements the stewardship probe (which only
    /// walks chunks) by exercising the HTTP body transport.
    #[arg(long)]
    cold: bool,

    /// After probing, re-upload the reference via PUT /stewardship/{ref}.
    /// Requires --stamp.
    #[arg(long)]
    reseed: bool,

    /// Postage batch ID for re-seed.
    #[arg(long, value_name = "ID", requires = "reseed")]
    stamp: Option<String>,

    /// Per-call timeout in seconds.
    #[arg(long, default_value_t = 60)]
    timeout: u64,

    /// Max concurrent chunk probes during drill-down.
    #[arg(long, default_value_t = 8)]
    concurrency: usize,

    /// Output format.
    #[arg(long, value_enum, default_value_t = OutputKind::Text)]
    output: OutputKind,

    /// Increase log verbosity on stderr. Repeat for more: `-v` info,
    /// `-vv` debug, `-vvv` trace. The default level is `warn` so
    /// stderr stays quiet unless something is wrong. Overridden by
    /// `RUST_LOG` if set.
    #[arg(short = 'v', long, action = clap::ArgAction::Count)]
    verbose: u8,
}

#[derive(Copy, Clone, Debug, ValueEnum)]
enum OutputKind {
    Text,
    Json,
}

impl From<OutputKind> for OutputFormat {
    fn from(k: OutputKind) -> Self {
        match k {
            OutputKind::Text => OutputFormat::Text,
            OutputKind::Json => OutputFormat::Json,
        }
    }
}

fn default_bees() -> Vec<String> {
    std::env::var("BEE_API_URL")
        .ok()
        .into_iter()
        .chain(std::iter::once("http://localhost:1633".to_string()))
        .take(1)
        .collect()
}

fn init_tracing(verbosity: u8) {
    use tracing_subscriber::EnvFilter;
    let default = match verbosity {
        0 => "warn",
        1 => "bee_check=info,bee=info",
        2 => "bee_check=debug,bee=debug",
        _ => "bee_check=trace,bee=trace",
    };
    let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(default));
    tracing_subscriber::fmt()
        .with_env_filter(filter)
        .with_writer(std::io::stderr)
        .with_target(false)
        .compact()
        .init();
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    init_tracing(cli.verbose);

    // Validate flag combinations BEFORE running any probe so a typo
    // like `--reseed` without `--stamp` fails fast instead of wasting
    // a multi-vantage probe.
    if cli.reseed && cli.stamp.is_none() {
        anyhow::bail!("--reseed requires --stamp <batch-id>");
    }

    let bees: Vec<String> = if cli.bee.is_empty() {
        default_bees()
    } else {
        cli.bee.clone()
    };

    let timeout = std::time::Duration::from_secs(cli.timeout);

    // Resolve the input. A `feed:OWNER:TOPIC` is turned into a bare
    // reference via the first vantage; everything else is taken as a
    // reference verbatim.
    let (reference, resolution) = match parse_input(&cli.reference) {
        // Normalize to lowercase so the report's `reference` matches
        // chunk addresses (which `Reference::to_hex` always emits as
        // lowercase) and so two reports for the same content compare
        // equal regardless of how the user typed the hex.
        ParsedInput::Reference(r) => (r.to_lowercase(), None),
        ParsedInput::Feed { owner, topic } => {
            let first_bee = bees
                .first()
                .context("feed resolution requires at least one --bee URL")?;
            let (r, res) = resolve_feed(first_bee, &owner, &topic, timeout)
                .await
                .context("feed resolution failed")?;
            eprintln!("resolved feed -> {r}");
            (r, Some(res))
        }
    };

    let gateways: Vec<String> = if cli.no_gateway {
        Vec::new()
    } else if cli.gateway.is_empty() {
        vec![DEFAULT_GATEWAY.to_string()]
    } else {
        cli.gateway.clone()
    };

    let (vantage_report, gateway_results) = tokio::join!(
        check_multi_vantage(&reference, &bees, timeout),
        check_gateways(&reference, &gateways, timeout),
    );
    let mut report = vantage_report.context("multi-vantage check failed")?;
    let gateways_out = gateway_results.context("gateway probe failed")?;
    report = merge_gateways(report, gateways_out);
    report.resolution = resolution;

    let report = if cli.per_chunk {
        drill_down(report, &bees, timeout, cli.concurrency)
            .await
            .context("per-chunk drill-down failed")?
    } else {
        report
    };

    let report = match cli.target_overlay.as_deref() {
        Some(target) => annotate_target_overlay(report, target),
        None => report,
    };

    let mut report = report;
    if cli.cold {
        let cold = cold_download_all(&bees, &gateways, &reference, timeout)
            .await
            .context("cold-download probe failed")?;
        report.cold_downloads = cold;
    }

    print!("{}", render_report(&report, cli.output.into()));

    if cli.reseed {
        let stamp = cli
            .stamp
            .as_ref()
            .context("--reseed requires --stamp <batch-id>")?;
        let target_bee = bees.first().context("no bee URL for --reseed")?;
        let status = check_stamp(target_bee, stamp, timeout)
            .await
            .context("stamp pre-flight failed")?;
        eprint!("{}", render_stamp_status(&status));
        if !status.exists || !status.usable {
            anyhow::bail!("refusing to re-seed: stamp is not usable; see warnings above");
        }
        let req = ReseedRequest {
            reference: reference.clone(),
            bee_url: target_bee.clone(),
            batch_id: stamp.clone(),
            timeout,
        };
        reseed(req).await.context("re-seed failed")?;
        eprintln!("re-seeded {} via {}", reference, target_bee);
    }

    let any_retrievable = report.vantages.iter().any(|v| v.retrievable == Some(true))
        || report.gateways.iter().any(|g| g.retrievable == Some(true));
    if !any_retrievable {
        std::process::exit(2);
    }
    Ok(())
}