use std::net::SocketAddr;
use anyhow::{Context, Result};
use iroh_io::ConcatenateSliceWriter;
use tracing_subscriber::{prelude::*, EnvFilter};
use iroh_bytes::{
get::fsm::{AtInitial, ConnectedNext, EndBlobNext},
hashseq::HashSeq,
protocol::GetRequest,
Hash,
};
mod connect;
use connect::{load_certs, make_client_endpoint};
pub fn setup_logging() {
tracing_subscriber::registry()
.with(tracing_subscriber::fmt::layer().with_writer(std::io::stderr))
.with(EnvFilter::from_default_env())
.try_init()
.ok();
}
#[tokio::main]
async fn main() -> Result<()> {
println!("\nfetch bytes example!");
setup_logging();
let args: Vec<_> = std::env::args().collect();
if args.len() != 4 {
anyhow::bail!("usage: fetch-bytes [HASH] [SOCKET_ADDR] [FORMAT]");
}
let hash: Hash = args[1].parse().context("unable to parse [HASH]")?;
let addr: SocketAddr = args[2].parse().context("unable to parse [SOCKET_ADDR]")?;
let format = {
if args[3] != "blob" && args[3] != "collection" {
anyhow::bail!(
"expected either 'blob' or 'collection' for FORMAT argument, got {}",
args[3]
);
}
args[3].clone()
};
let roots = load_certs().await?;
let endpoint = make_client_endpoint(roots)?;
println!("\nlistening on {}", endpoint.local_addr()?);
println!("fetching hash {hash} from {addr}");
let connection = endpoint.connect(addr, "localhost")?.await?;
if format == "collection" {
let request = GetRequest::all(hash);
let initial = iroh_bytes::get::fsm::start(connection, request);
write_collection(initial).await
} else {
let request = GetRequest::single(hash);
let initial = iroh_bytes::get::fsm::start(connection, request);
write_blob(initial).await
}
}
async fn write_blob(initial: AtInitial) -> Result<()> {
let connected = initial.next().await?;
let ConnectedNext::StartRoot(start_root) = connected.next().await? else {
panic!("expected start root")
};
let header = start_root.next();
let writer = ConcatenateSliceWriter::new(tokio::io::stdout());
println!();
let end = header.write_all(writer).await?;
let EndBlobNext::Closing(closing) = end.next() else {
panic!("expected closing")
};
let _stats = closing.next().await?;
Ok(())
}
async fn write_collection(initial: AtInitial) -> Result<()> {
let connected = initial.next().await?;
let ConnectedNext::StartRoot(start_root) = connected.next().await? else {
anyhow::bail!("failed to parse collection");
};
if !start_root.ranges().is_all() {
anyhow::bail!("collection was not requested completely");
}
let header: iroh_bytes::get::fsm::AtBlobHeader = start_root.next();
let (root_end, hashes_bytes) = header.concatenate_into_vec().await?;
let next = root_end.next();
let EndBlobNext::MoreChildren(at_meta) = next else {
anyhow::bail!("missing meta blob, got {next:?}");
};
let hashes = HashSeq::try_from(bytes::Bytes::from(hashes_bytes))
.context("failed to parse hashes")?
.into_iter()
.collect::<Vec<_>>();
let meta_hash = hashes.first().context("missing meta hash")?;
let (meta_end, _meta_bytes) = at_meta.next(*meta_hash).concatenate_into_vec().await?;
let mut curr = meta_end.next();
let closing = loop {
match curr {
EndBlobNext::MoreChildren(more) => {
let Some(hash) = hashes.get(more.child_offset() as usize) else {
break more.finish();
};
let header = more.next(*hash);
let writer = ConcatenateSliceWriter::new(tokio::io::stdout());
let end = header.write_all(writer).await?;
println!();
curr = end.next();
}
EndBlobNext::Closing(closing) => {
break closing;
}
}
};
let _stats = closing.next().await?;
Ok(())
}