Rework transfer example to use clap

since the others use it as well

Author: rklaehn

examples/transfer.rs

@@ -1,10 +1,34 @@
 use std::path::PathBuf;
 
-use iroh::{protocol::Router, Endpoint};
+use clap::Parser;
+use iroh::{discovery::pkarr::PkarrResolver, protocol::Router, Endpoint};
 use iroh_blobs::{store::mem::MemStore, ticket::BlobTicket, BlobsProtocol};
+mod common;
+use common::setup_logging;
+#[derive(Debug, Parser)]
+#[command(version, about)]
+pub struct Cli {
+    #[clap(subcommand)]
+    command: Command,
+}
 
-#[tokio::main]
-async fn main() -> anyhow::Result<()> {
+#[derive(Parser, Debug)]
+pub enum Command {
+    /// Send a file to the network
+    Send {
+        /// Path to the file to send
+        file: PathBuf,
+    },
+    /// Receive a file from the network
+    Receive {
+        /// Ticket describing the content to fetch
+        ticket: BlobTicket,
+        /// Path to save the received file
+        filename: PathBuf,
+    },
+}
+
+async fn send(filename: PathBuf) -> anyhow::Result<()> {
     // Create an endpoint, it allows creating and accepting
     // connections in the iroh p2p world
     let endpoint = Endpoint::builder().discovery_n0().bind().await?;
@@ -14,80 +38,82 @@ async fn main() -> anyhow::Result<()> {
     // Then we initialize a struct that can accept blobs requests over iroh connections
     let blobs = BlobsProtocol::new(&store, endpoint.clone(), None);
 
-    // Grab all passed in arguments, the first one is the binary itself, so we skip it.
-    let args: Vec<String> = std::env::args().skip(1).collect();
-    // Convert to &str, so we can pattern-match easily:
-    let arg_refs: Vec<&str> = args.iter().map(String::as_str).collect();
+    let abs_path = std::path::absolute(&filename)?;
+    let tag = store.blobs().add_path(abs_path).await?;
+
+    let node_id = endpoint.node_id();
+    let ticket = BlobTicket::new(node_id.into(), tag.hash, tag.format);
 
-    match arg_refs.as_slice() {
-        ["send", filename] => {
-            let filename: PathBuf = filename.parse()?;
-            let abs_path = std::path::absolute(&filename)?;
+    println!("File hashed. Fetch this file by running:");
+    println!(
+        "cargo run --example transfer -- receive {ticket} {}",
+        filename.display()
+    );
 
-            println!("Hashing file.");
+    // For sending files we build a router that accepts blobs connections & routes them
+    // to the blobs protocol.
+    let router = Router::builder(endpoint)
+        .accept(iroh_blobs::ALPN, blobs)
+        .spawn();
 
-            // When we import a blob, we get back a "tag" that refers to said blob in the store
-            // and allows us to control when/if it gets garbage-collected
-            let tag = store.blobs().add_path(abs_path).await?;
+    tokio::signal::ctrl_c().await?;
 
-            let node_id = endpoint.node_id();
-            let ticket = BlobTicket::new(node_id.into(), tag.hash, tag.format);
+    // Gracefully shut down the node
+    println!("Shutting down.");
+    router.shutdown().await?;
+    Ok(())
+}
 
-            println!("File hashed. Fetch this file by running:");
-            println!(
-                "cargo run --example transfer -- receive {ticket} {}",
-                filename.display()
-            );
+async fn receive(ticket: BlobTicket, filename: PathBuf) -> anyhow::Result<()> {
+    // Create a store to download blobs into
+    let store = MemStore::new();
 
-            // For sending files we build a router that accepts blobs connections & routes them
-            // to the blobs protocol.
-            let router = Router::builder(endpoint)
-                .accept(iroh_blobs::ALPN, blobs)
-                .spawn();
+    // Create an endpoint, it allows creating and accepting
+    // connections in the iroh p2p world.
+    //
+    // Since we just want to receive files, we don't need a stable node address
+    // or to publish our discovery information.
+    let endpoint = Endpoint::builder()
+        .discovery(PkarrResolver::n0_dns())
+        .bind()
+        .await?;
 
-            tokio::signal::ctrl_c().await?;
+    // For receiving files, we create a "downloader" that allows us to fetch files
+    // from other nodes via iroh connections
+    let downloader = store.downloader(&endpoint);
 
-            // Gracefully shut down the node
-            println!("Shutting down.");
-            router.shutdown().await?;
-        }
-        ["receive", ticket, filename] => {
-            let filename: PathBuf = filename.parse()?;
-            let abs_path = std::path::absolute(filename)?;
-            let ticket: BlobTicket = ticket.parse()?;
+    println!("Starting download.");
 
-            // For receiving files, we create a "downloader" that allows us to fetch files
-            // from other nodes via iroh connections
-            let downloader = store.downloader(&endpoint);
+    downloader
+        .download(ticket.hash(), [ticket.node_addr().node_id])
+        .await?;
 
-            println!("Starting download.");
+    println!("Finished download.");
+    println!("Copying to destination.");
 
-            downloader
-                .download(ticket.hash(), Some(ticket.node_addr().node_id))
-                .await?;
+    store.export(ticket.hash(), filename).await?;
 
-            println!("Finished download.");
-            println!("Copying to destination.");
+    println!("Finished copying.");
 
-            store.blobs().export(ticket.hash(), abs_path).await?;
+    // Gracefully shut down the endpoint and the store
+    println!("Shutting down.");
+    endpoint.close().await;
+    store.shutdown().await?;
+    Ok(())
+}
 
-            println!("Finished copying.");
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    setup_logging();
+    let cli = Cli::parse();
 
-            // Gracefully shut down the node
-            println!("Shutting down.");
-            endpoint.close().await;
+    match cli.command {
+        Command::Send { file } => {
+            send(file).await?;
         }
-        _ => {
-            println!("Couldn't parse command line arguments: {args:?}");
-            println!("Usage:");
-            println!("    # to send:");
-            println!("    cargo run --example transfer -- send [FILE]");
-            println!("    # this will print a ticket.");
-            println!();
-            println!("    # to receive:");
-            println!("    cargo run --example transfer -- receive [TICKET] [FILE]");
+        Command::Receive { ticket, filename } => {
+            receive(ticket, filename).await?;
         }
     }
-
     Ok(())
 }