Add query_dht cli arg

Author: rklaehn

content-discovery/Cargo.lock

@@ -18,9 +18,8 @@ serde = { version = "1", features = ["derive"] }
 # Optional features for the client functionality
 tracing = { version = "0.1", optional = true }
 quinn = { version = "0.10", optional = true }
-iroh-pkarr-node-discovery = { path = "../../iroh-pkarr-node-discovery", optional = true }
+iroh-pkarr-node-discovery = { version = "0.1.2", optional = true }
 mainline = { version = "1.0.0", optional = true }
-pkarr = { version = "1.0.1", features = ["async"], optional = true }
 anyhow = { version = "1", features = ["backtrace"], optional = true }
 postcard = { version = "1", default-features = false, features = ["alloc", "use-std"], optional = true }
 ed25519-dalek = { version = "2.1.0", optional = true }
@@ -37,9 +36,9 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = tr
 tokio = { version = "1", features = ["io-util", "rt"], optional = true }
 
 [features]
-client = ["iroh-pkarr-node-discovery", "mainline", "pkarr", "quinn", "tracing", "anyhow", "rcgen", "genawaiter", "rustls", "futures", "postcard"]
+client = ["iroh-pkarr-node-discovery", "mainline", "quinn", "tracing", "anyhow", "rcgen", "genawaiter", "rustls", "futures", "postcard"]
 cli = ["client", "clap", "tempfile", "derive_more", "tracing-subscriber", "tokio"]
-default = []
+default = ["client", "cli"]
 
 [[bin]]
 name = "iroh-mainline-content-discovery"

content-discovery/iroh-mainline-content-discovery/Cargo.toml

@@ -15,6 +15,7 @@ pub struct Args {
 pub enum Commands {
     Announce(AnnounceArgs),
     Query(QueryArgs),
+    QueryDht(QueryDhtArgs),
 }
 
 /// Various ways to specify content.
@@ -80,10 +81,6 @@ pub struct AnnounceArgs {
     #[clap(long)]
     pub tracker: NodeId,
 
-    /// the port to use for announcing
-    #[clap(long)]
-    pub port: Option<u16>,
-
     /// The host to announce. Not needed if content is a ticket.
     #[clap(long)]
     pub host: Option<NodeId>,
@@ -98,17 +95,36 @@ pub struct AnnounceArgs {
     /// Announce that the peer has only partial data.
     #[clap(long)]
     pub partial: bool,
+
+    /// the port to use for announcing
+    #[clap(long)]
+    pub magic_port: Option<u16>,
 }
 
 #[derive(Parser, Debug)]
 pub struct QueryArgs {
+    /// the tracker to query
     #[clap(long)]
     pub tracker: TrackerId,
 
+    /// The content to find hosts for.
+    pub content: ContentArg,
+
+    /// Ask for hosts that were announced as having just partial data
+    #[clap(long)]
+    pub partial: bool,
+
+    /// Ask for hosts that were recently checked and found to have some data
+    #[clap(long)]
+    pub verified: bool,
+
     /// the port to use for querying
     #[clap(long)]
-    pub port: Option<u16>,
+    pub magic_port: Option<u16>,
+}
 
+#[derive(Parser, Debug)]
+pub struct QueryDhtArgs {
     /// The content to find hosts for.
     pub content: ContentArg,
 
@@ -119,4 +135,12 @@ pub struct QueryArgs {
     /// Ask for hosts that were recently checked and found to have some data
     #[clap(long)]
     pub verified: bool,
+
+    /// Parallelism for querying the dht
+    #[clap(long)]
+    pub query_parallelism: Option<usize>,
+
+    /// the port to use for querying
+    #[clap(long)]
+    pub quinn_port: Option<u16>,
 }

content-discovery/iroh-mainline-content-discovery/src/args.rs

@@ -11,7 +11,6 @@ use iroh_bytes::HashAndFormat;
 use iroh_net::{MagicEndpoint, NodeId};
 use iroh_pkarr_node_discovery::PkarrNodeDiscovery;
 use mainline::common::{GetPeerResponse, StoreQueryMetdata};
-use pkarr::PkarrClient;
 
 use crate::protocol::{
     Announce, Query, QueryResponse, Request, Response, ALPN, REQUEST_SIZE_LIMIT,
@@ -52,7 +51,7 @@ fn unique_tracker_addrs(
     Gen::new(|co| async move {
         let mut found = HashSet::new();
         while let Some(response) = response.next_async().await {
-            println!("got get_peers response: {:?}", response);
+            tracing::info!("got get_peers response: {:?}", response);
             let tracker = response.peer;
             if !found.insert(tracker) {
                 continue;
@@ -185,7 +184,7 @@ pub async fn query(connection: quinn::Connection, args: Query) -> anyhow::Result
 }
 
 /// Create a quinn client endpoint.
-fn create_quinn_client(
+pub fn create_quinn_client(
     bind_addr: SocketAddr,
     alpn_protocols: Vec<Vec<u8>>,
     keylog: bool,
@@ -207,9 +206,11 @@ async fn create_endpoint(
     port: u16,
     publish: bool,
 ) -> anyhow::Result<MagicEndpoint> {
-    let pkarr = PkarrClient::new();
-    let discovery_key = if publish { Some(&key) } else { None };
-    let mainline_discovery = PkarrNodeDiscovery::new(pkarr, discovery_key);
+    let mainline_discovery = if publish {
+        PkarrNodeDiscovery::builder().secret_key(&key).build()
+    } else {
+        PkarrNodeDiscovery::default()
+    };
     iroh_net::MagicEndpoint::builder()
         .secret_key(key)
         .discovery(Box::new(mainline_discovery))

content-discovery/iroh-mainline-content-discovery/src/client.rs

@@ -1,16 +1,24 @@
 pub mod args;
 
-use std::collections::BTreeSet;
+use std::{
+    collections::BTreeSet,
+    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
+};
 
 use anyhow::Context;
+use args::QueryDhtArgs;
 use clap::Parser;
-use iroh_mainline_content_discovery::protocol::{Announce, AnnounceKind, Query, QueryFlags, ALPN};
+use futures::StreamExt;
+use iroh_mainline_content_discovery::{
+    create_quinn_client,
+    protocol::{Announce, AnnounceKind, Query, QueryFlags, ALPN},
+    to_infohash,
+};
 use iroh_net::{
     magic_endpoint::{get_alpn, get_remote_node_id},
     MagicEndpoint, NodeId,
 };
 use iroh_pkarr_node_discovery::PkarrNodeDiscovery;
-use pkarr::PkarrClient;
 use tokio::io::AsyncWriteExt;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
 
@@ -21,9 +29,11 @@ async fn create_endpoint(
     port: u16,
     publish: bool,
 ) -> anyhow::Result<MagicEndpoint> {
-    let pkarr = PkarrClient::new();
-    let discovery_key = if publish { Some(&key) } else { None };
-    let mainline_discovery = PkarrNodeDiscovery::new(pkarr, discovery_key);
+    let mainline_discovery = if publish {
+        PkarrNodeDiscovery::builder().secret_key(&key).build()
+    } else {
+        PkarrNodeDiscovery::default()
+    };
     iroh_net::MagicEndpoint::builder()
         .secret_key(key)
         .discovery(Box::new(mainline_discovery))
@@ -68,7 +78,7 @@ async fn announce(args: AnnounceArgs) -> anyhow::Result<()> {
     for content in &content {
         println!("    {}", content);
     }
-    let endpoint = create_endpoint(key, 11112, false).await?;
+    let endpoint = create_endpoint(key, args.magic_port.unwrap_or_default(), false).await?;
     let connection = endpoint.connect_by_node_id(&args.tracker, ALPN).await?;
     println!("connected to {:?}", connection.remote_address());
     let kind = if args.partial {
@@ -87,9 +97,11 @@ async fn announce(args: AnnounceArgs) -> anyhow::Result<()> {
 }
 
 async fn query(args: QueryArgs) -> anyhow::Result<()> {
-    let connection =
-        iroh_mainline_content_discovery::connect(&args.tracker, args.port.unwrap_or_default())
-            .await?;
+    let connection = iroh_mainline_content_discovery::connect(
+        &args.tracker,
+        args.magic_port.unwrap_or_default(),
+    )
+    .await?;
     let q = Query {
         content: args.content.hash_and_format(),
         flags: QueryFlags {
@@ -108,6 +120,36 @@ async fn query(args: QueryArgs) -> anyhow::Result<()> {
     Ok(())
 }
 
+async fn query_dht(args: QueryDhtArgs) -> anyhow::Result<()> {
+    let bind_addr = SocketAddr::V4(SocketAddrV4::new(
+        Ipv4Addr::UNSPECIFIED,
+        args.quinn_port.unwrap_or_default(),
+    ));
+    let endpoint = create_quinn_client(bind_addr, vec![ALPN.to_vec()], false)?;
+    let dht = mainline::Dht::default();
+    let q = Query {
+        content: args.content.hash_and_format(),
+        flags: QueryFlags {
+            complete: !args.partial,
+            verified: args.verified,
+        },
+    };
+    println!("content corresponds to infohash {}", to_infohash(q.content));
+    let mut stream = iroh_mainline_content_discovery::query_dht(
+        endpoint,
+        dht,
+        q,
+        args.query_parallelism.unwrap_or(4),
+    );
+    while let Some(item) = stream.next().await {
+        match item {
+            Ok(provider) => println!("found provider {}", provider),
+            Err(e) => println!("error: {}", e),
+        }
+    }
+    Ok(())
+}
+
 // set the RUST_LOG env var to one of {debug,info,warn} to see logging info
 pub fn setup_logging() {
     tracing_subscriber::registry()
@@ -124,6 +166,7 @@ async fn main() -> anyhow::Result<()> {
     match args.command {
         Commands::Announce(args) => announce(args).await,
         Commands::Query(args) => query(args).await,
+        Commands::QueryDht(args) => query_dht(args).await,
     }
 }
 

content-discovery/iroh-mainline-content-discovery/src/main.rs

@@ -1,8 +1,8 @@
 [package]
 name = "iroh-mainline-tracker"
-version = "0.3.0"
+version = "0.1.0"
 edition = "2021"
-description = "Content discovery for iroh, using the bittorrent mainline DHT"
+description = "Content tracker for iroh, using the bittorrent mainline DHT"
 license = "MIT OR Apache-2.0"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -11,7 +11,6 @@ license = "MIT OR Apache-2.0"
 anyhow = { version = "1", features = ["backtrace"] }
 bao-tree = { version = "0.9.1", features = ["tokio_fsm"], default-features = false }
 bytes = "1"
-clap = { version = "4", features = ["derive"], optional = true }
 derive_more = { version = "1.0.0-beta.1", features = ["debug", "display", "from", "try_into"] }
 dirs-next = "2"
 ed25519-dalek = "2.1.0"
@@ -20,10 +19,10 @@ hex = "0.4.3"
 humantime = "2.1.0"
 iroh-net = "0.12.0"
 iroh-bytes = "0.12.0"
-iroh-pkarr-node-discovery = { path = "../../iroh-pkarr-node-discovery" }
+iroh-pkarr-node-discovery = { path = "../../iroh-pkarr-node-discovery", default-features = false }
 mainline = "1.0.0"
 pkarr = { version = "1.0.1", features = ["async"] }
-postcard = { version = "1", default-features = false, features = ["alloc", "use-std", "experimental-derive"] }
+postcard = { version = "1", default-features = false, features = ["alloc", "use-std"] }
 quinn = "0.10"
 rand = "0.8"
 rcgen = "0.12.0"
@@ -44,6 +43,8 @@ flume = "0.11.0"
 genawaiter = { version = "0.99.1", features = ["futures03"] }
 iroh-mainline-content-discovery = { path = "../iroh-mainline-content-discovery", features = ["client"] }
 
+clap = { version = "4", features = ["derive"], optional = true }
+
 [features]
 cli = ["clap"]
 default = ["cli"]

content-discovery/iroh-mainline-tracker/Cargo.toml

@@ -25,7 +25,6 @@ use iroh_net::{
     MagicEndpoint, NodeId,
 };
 use iroh_pkarr_node_discovery::PkarrNodeDiscovery;
-use pkarr::PkarrClient;
 use tokio::io::AsyncWriteExt;
 use tokio_util::task::LocalPoolHandle;
 
@@ -71,9 +70,11 @@ async fn create_endpoint(
     port: u16,
     publish: bool,
 ) -> anyhow::Result<MagicEndpoint> {
-    let pkarr = PkarrClient::new();
-    let discovery_key = if publish { Some(&key) } else { None };
-    let mainline_discovery = PkarrNodeDiscovery::new(pkarr, discovery_key);
+    let mainline_discovery = if publish {
+        PkarrNodeDiscovery::builder().secret_key(&key).build()
+    } else {
+        PkarrNodeDiscovery::default()
+    };
     iroh_net::MagicEndpoint::builder()
         .secret_key(key)
         .discovery(Box::new(mainline_discovery))