Implementing DHT publishing

Author: rklaehn

iroh-mainline-content-discovery/Cargo.lock

@@ -40,6 +40,8 @@ tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 ttl_cache = "0.5.1"
 url = "2.5.0"
+flume = "0.11.0"
+genawaiter = { version = "0.99.1", features = ["futures03"] }
 
 [features]
 cli = ["clap"]

iroh-mainline-content-discovery/Cargo.toml

@@ -1,11 +1,20 @@
-use std::{net::SocketAddr, sync::Arc, time::Duration};
+use genawaiter::sync::Gen;
+use std::{
+    collections::{BTreeSet, HashSet},
+    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
+    sync::Arc,
+    time::Duration,
+};
 
 use anyhow::Context;
+use futures::{future::BoxFuture, FutureExt, Stream, StreamExt};
+use iroh_bytes::HashAndFormat;
 use iroh_net::{
     magic_endpoint::{get_alpn, get_remote_node_id},
     MagicEndpoint, NodeId,
 };
 use iroh_pkarr_node_discovery::PkarrNodeDiscovery;
+use mainline::common::{GetPeerResponse, StoreQueryMetdata};
 use pkarr::PkarrClient;
 use protocol::QueryResponse;
 use tokio::io::AsyncWriteExt;
@@ -18,36 +27,6 @@ pub mod options;
 pub mod protocol;
 pub mod tracker;
 
-/// A tracker id for queries - either a node id or an address.
-#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
-pub enum TrackerId {
-    NodeId(NodeId),
-    Addr(std::net::SocketAddr),
-}
-
-impl std::fmt::Display for TrackerId {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            TrackerId::NodeId(node_id) => write!(f, "{}", node_id),
-            TrackerId::Addr(addr) => write!(f, "{}", addr),
-        }
-    }
-}
-
-impl std::str::FromStr for TrackerId {
-    type Err = anyhow::Error;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        if let Ok(node_id) = s.parse() {
-            return Ok(TrackerId::NodeId(node_id));
-        }
-        if let Ok(addr) = s.parse() {
-            return Ok(TrackerId::Addr(addr));
-        }
-        anyhow::bail!("invalid tracker id")
-    }
-}
-
 /// Accept an incoming connection and extract the client-provided [`NodeId`] and ALPN protocol.
 pub(crate) async fn accept_conn(
     mut conn: quinn::Connecting,
@@ -78,6 +57,100 @@ pub async fn announce(connection: quinn::Connection, args: Announce) -> anyhow::
     Ok(())
 }
 
+fn to_infohash(haf: HashAndFormat) -> mainline::Id {
+    let mut data = [0u8; 20];
+    data.copy_from_slice(&haf.hash.as_bytes()[..20]);
+    mainline::Id::from_bytes(data).unwrap()
+}
+
+fn unique_tracker_addrs(
+    mut response: mainline::common::Response<GetPeerResponse>,
+) -> impl Stream<Item = SocketAddr> {
+    Gen::new(|co| async move {
+        let mut found = HashSet::new();
+        while let Some(response) = response.next_async().await {
+            let tracker = response.peer;
+            if !found.insert(tracker) {
+                continue;
+            }
+            co.yield_(tracker).await;
+        }
+    })
+}
+
+async fn query_one(
+    endpoint: impl ConnectionProvider,
+    addr: SocketAddr,
+    args: Query,
+) -> anyhow::Result<Vec<NodeId>> {
+    let connection = endpoint.connect(addr).await?;
+    let result = query(connection, args).await?;
+    Ok(result.hosts)
+}
+
+/// A connection provider that can be used to connect to a tracker.
+///
+/// This can either be a [`quinn::Endpoint`] where connections are created on demand,
+/// or some sort of connection pool.
+pub trait ConnectionProvider: Clone {
+    fn connect(&self, addr: SocketAddr) -> BoxFuture<anyhow::Result<quinn::Connection>>;
+}
+
+impl ConnectionProvider for quinn::Endpoint {
+    fn connect(&self, addr: SocketAddr) -> BoxFuture<anyhow::Result<quinn::Connection>> {
+        async move { Ok(self.connect(addr, "localhost")?.await?) }.boxed()
+    }
+}
+
+/// Query the mainline DHT for trackers for the given content, then query each tracker for peers.
+pub async fn query_dht(
+    endpoint: impl ConnectionProvider,
+    dht: mainline::dht::Dht,
+    args: Query,
+    query_parallelism: usize,
+) -> impl Stream<Item = anyhow::Result<NodeId>> {
+    let dht = dht.as_async();
+    let info_hash = to_infohash(args.content);
+    let response: mainline::common::Response<GetPeerResponse> = dht.get_peers(info_hash);
+    let unique_tracker_addrs = unique_tracker_addrs(response);
+    unique_tracker_addrs
+        .map(move |addr| {
+            let endpoint = endpoint.clone();
+            async move {
+                let hosts = match query_one(endpoint, addr, args).await {
+                    Ok(hosts) => hosts.into_iter().map(anyhow::Ok).collect(),
+                    Err(cause) => vec![Err(cause)],
+                };
+                futures::stream::iter(hosts)
+            }
+        })
+        .buffer_unordered(query_parallelism)
+        .flatten()
+}
+
+/// Announce to the mainline DHT in parallel.
+///
+/// Note that this should only be called from a publicly reachable node, where port is the port
+/// on which the tracker protocol is reachable.
+pub fn announce_dht(
+    dht: mainline::dht::Dht,
+    content: BTreeSet<HashAndFormat>,
+    port: u16,
+    announce_parallelism: usize,
+) -> impl Stream<Item = (HashAndFormat, mainline::Result<StoreQueryMetdata>)> {
+    let dht = dht.as_async();
+    futures::stream::iter(content)
+        .map(move |content| {
+            let dht = dht.clone();
+            async move {
+                let info_hash = to_infohash(content);
+                let res = dht.announce_peer(info_hash, Some(port)).await;
+                (content, res)
+            }
+        })
+        .buffer_unordered(announce_parallelism)
+}
+
 /// Assume an existing connection to a tracker and query it for peers for some content.
 pub async fn query(connection: quinn::Connection, args: Query) -> anyhow::Result<QueryResponse> {
     tracing::info!("connected to {:?}", connection.remote_address());
@@ -193,3 +266,69 @@ async fn create_endpoint(
         .bind(port)
         .await
 }
+
+/// A tracker id for queries - either a node id or an address.
+#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
+pub enum TrackerId {
+    NodeId(NodeId),
+    Addr(std::net::SocketAddr),
+}
+
+impl std::fmt::Display for TrackerId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            TrackerId::NodeId(node_id) => write!(f, "{}", node_id),
+            TrackerId::Addr(addr) => write!(f, "{}", addr),
+        }
+    }
+}
+
+impl std::str::FromStr for TrackerId {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if let Ok(node_id) = s.parse() {
+            return Ok(TrackerId::NodeId(node_id));
+        }
+        if let Ok(addr) = s.parse() {
+            return Ok(TrackerId::Addr(addr));
+        }
+        anyhow::bail!("invalid tracker id")
+    }
+}
+
+/// Connect to a tracker using the [protocol::TRACKER_ALPN] protocol, using either
+/// a node id or an address.
+///
+/// Note that this is less efficient than using an existing endpoint when doing multiple requests.
+/// It is provided as a convenience function for short lived utilities.
+pub async fn connect(tracker: &TrackerId, local_port: u16) -> anyhow::Result<quinn::Connection> {
+    match tracker {
+        TrackerId::Addr(tracker) => connect_socket(*tracker, local_port).await,
+        TrackerId::NodeId(tracker) => connect_magic(&tracker, local_port).await,
+    }
+}
+
+/// Create a magic endpoint and connect to a tracker using the [protocol::TRACKER_ALPN] protocol.
+pub async fn connect_magic(tracker: &NodeId, local_port: u16) -> anyhow::Result<quinn::Connection> {
+    // todo: uncomment once the connection problems are fixed
+    // for now, a random node id is more reliable.
+    // let key = load_secret_key(tracker_path(CLIENT_KEY)?).await?;
+    let key = iroh_net::key::SecretKey::generate();
+    let endpoint = create_endpoint(key, local_port, false).await?;
+    tracing::info!("trying to connect to tracker at {:?}", tracker);
+    let connection = endpoint.connect_by_node_id(tracker, TRACKER_ALPN).await?;
+    Ok(connection)
+}
+
+/// Create a quinn endpoint and connect to a tracker using the [protocol::TRACKER_ALPN] protocol.
+pub async fn connect_socket(
+    tracker: SocketAddr,
+    local_port: u16,
+) -> anyhow::Result<quinn::Connection> {
+    let bind_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, local_port));
+    let endpoint = create_quinn_client(bind_addr, vec![TRACKER_ALPN.to_vec()], false)?;
+    tracing::info!("trying to connect to tracker at {:?}", tracker);
+    let connection = endpoint.connect(tracker, "localhost")?.await?;
+    Ok(connection)
+}

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

@@ -12,17 +12,13 @@ use std::{
 
 use anyhow::Context;
 use clap::Parser;
-use iroh_mainline_content_discovery::TrackerId;
 use iroh_mainline_content_discovery::{
     io::{
         self, load_from_file, setup_logging, tracker_home, tracker_path, CONFIG_DEFAULTS_FILE,
         CONFIG_FILE, SERVER_KEY_FILE,
     },
     options::Options,
-    protocol::{
-        Announce, AnnounceKind, Query, QueryFlags, Request, Response, REQUEST_SIZE_LIMIT,
-        TRACKER_ALPN,
-    },
+    protocol::{Announce, AnnounceKind, Query, QueryFlags, TRACKER_ALPN},
     tracker::Tracker,
 };
 use iroh_net::{
@@ -127,16 +123,17 @@ async fn server(args: ServerArgs) -> anyhow::Result<()> {
         "tracker query --tracker {} <hash> or <ticket>",
         addr.node_id
     );
-    log!();
     let db2 = db.clone();
     let db3 = db.clone();
+    let db4 = db.clone();
     let endpoint2 = endpoint.clone();
     let _probe_task = tpc.spawn_pinned(move || db2.probe_loop(endpoint2));
+    let _announce_task = tpc.spawn_pinned(move || db3.dht_announce_loop(args.quinn_port));
     let magic_accept_task = tokio::spawn(db.magic_accept_loop(endpoint));
     let server_config = configure_server(&key)?;
     let bind_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, args.quinn_port));
     let quinn_endpoint = quinn::Endpoint::server(server_config, bind_addr)?;
-    let quinn_accept_task = tokio::spawn(db3.quinn_accept_loop(quinn_endpoint));
+    let quinn_accept_task = tokio::spawn(db4.quinn_accept_loop(quinn_endpoint));
     magic_accept_task.await??;
     quinn_accept_task.await??;
     Ok(())
@@ -189,7 +186,9 @@ async fn announce(args: AnnounceArgs) -> anyhow::Result<()> {
 }
 
 async fn query(args: QueryArgs) -> anyhow::Result<()> {
-    let connection = connect(&args.tracker).await?;
+    let connection =
+        iroh_mainline_content_discovery::connect(&args.tracker, args.port.unwrap_or_default())
+            .await?;
     let q = Query {
         content: args.content.hash_and_format(),
         flags: QueryFlags {
@@ -208,32 +207,6 @@ async fn query(args: QueryArgs) -> anyhow::Result<()> {
     Ok(())
 }
 
-async fn connect(tracker: &TrackerId) -> anyhow::Result<quinn::Connection> {
-    match tracker {
-        TrackerId::Addr(tracker) => connect_socket(*tracker).await,
-        TrackerId::NodeId(tracker) => connect_magic(tracker.clone()).await,
-    }
-}
-
-async fn connect_magic(tracker: NodeId) -> anyhow::Result<quinn::Connection> {
-    // todo: uncomment once the connection problems are fixed
-    // for now, a random node id is more reliable.
-    // let key = load_secret_key(tracker_path(CLIENT_KEY)?).await?;
-    let key = iroh_net::key::SecretKey::generate();
-    let endpoint = create_endpoint(key, 0, false).await?;
-    tracing::info!("trying to connect to tracker at {:?}", tracker);
-    let connection = endpoint.connect_by_node_id(&tracker, TRACKER_ALPN).await?;
-    Ok(connection)
-}
-
-async fn connect_socket(tracker: SocketAddr) -> anyhow::Result<quinn::Connection> {
-    let bind_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0));
-    let endpoint = create_quinn_client(bind_addr, vec![TRACKER_ALPN.to_vec()], false)?;
-    tracing::info!("trying to connect to tracker at {:?}", tracker);
-    let connection = endpoint.connect(tracker, "localhost")?.await?;
-    Ok(connection)
-}
-
 #[tokio::main(flavor = "multi_thread")]
 async fn main() -> anyhow::Result<()> {
     setup_logging();
@@ -251,23 +224,6 @@ fn configure_server(secret_key: &iroh_net::key::SecretKey) -> anyhow::Result<qui
     make_server_config(secret_key, 8, 1024, vec![TRACKER_ALPN.to_vec()])
 }
 
-fn create_quinn_client(
-    bind_addr: SocketAddr,
-    alpn_protocols: Vec<Vec<u8>>,
-    keylog: bool,
-) -> anyhow::Result<quinn::Endpoint> {
-    let secret_key = iroh_net::key::SecretKey::generate();
-    let tls_client_config =
-        iroh_net::tls::make_client_config(&secret_key, None, alpn_protocols, keylog)?;
-    let mut client_config = quinn::ClientConfig::new(Arc::new(tls_client_config));
-    let mut endpoint = quinn::Endpoint::client(bind_addr)?;
-    let mut transport_config = quinn::TransportConfig::default();
-    transport_config.keep_alive_interval(Some(Duration::from_secs(1)));
-    client_config.transport_config(Arc::new(transport_config));
-    endpoint.set_default_client_config(client_config);
-    Ok(endpoint)
-}
-
 /// Create a [`quinn::ServerConfig`] with the given secret key and limits.
 pub fn make_server_config(
     secret_key: &iroh_net::key::SecretKey,

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

@@ -29,6 +29,10 @@ pub struct Options {
     pub announce_data_path: Option<PathBuf>,
     // number of peers to probe in parallel
     pub probe_parallelism: usize,
+
+    pub dht_announce_parallelism: usize,
+
+    pub dht_announce_interval: Duration,
 }
 
 impl Default for Options {
@@ -43,6 +47,8 @@ impl Default for Options {
             probe_log: Some("probe.log".into()),
             announce_data_path: Some("announce.data.toml".into()),
             probe_parallelism: 4,
+            dht_announce_parallelism: 4,
+            dht_announce_interval: Duration::from_secs(10),
         }
     }
 }

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

@@ -45,7 +45,7 @@ pub struct Announce {
 }
 
 ///
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
 pub struct QueryFlags {
     /// Only return peers that supposedly have the complete data.
     ///
@@ -63,7 +63,7 @@ pub struct QueryFlags {
 }
 
 /// Query a peer for a blob or set of blobs.
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
 pub struct Query {
     /// The content we want to find.
     ///