iroh_content_discovery/
client.rs

1use std::{future::Future, result};
2
3use iroh::{
4    endpoint::{ConnectOptions, Connection},
5    Endpoint, NodeId,
6};
7use n0_future::{BufferedStreamExt, Stream, StreamExt};
8use snafu::prelude::*;
9use tracing::trace;
10
11use crate::protocol::{
12    Query, QueryResponse, Request, Response, SignedAnnounce, ALPN, REQUEST_SIZE_LIMIT,
13};
14
15#[derive(Debug, Snafu)]
16pub enum Error {
17    #[snafu(display("Failed to connect to tracker: {}", source))]
18    Connect {
19        source: iroh::endpoint::ConnectWithOptsError,
20        backtrace: snafu::Backtrace,
21    },
22
23    #[snafu(display("Failed connect to tracker using 1-rtt: {}", source))]
24    Connect1Rtt {
25        source: iroh::endpoint::ConnectionError,
26        backtrace: snafu::Backtrace,
27    },
28
29    #[snafu(display("Failed to open bidi stream to tracker: {}", source))]
30    OpenStream {
31        source: iroh::endpoint::ConnectionError,
32        backtrace: snafu::Backtrace,
33    },
34
35    #[snafu(display("Failed to serialize request: {}", source))]
36    SerializeRequest {
37        source: postcard::Error,
38        backtrace: snafu::Backtrace,
39    },
40
41    #[snafu(display("Failed to write data: {}", source))]
42    WriteRequest {
43        source: iroh::endpoint::WriteError,
44        backtrace: snafu::Backtrace,
45    },
46
47    #[snafu(display("Failed to finish: {}", source))]
48    FinishWrite {
49        source: iroh::endpoint::ClosedStream,
50        backtrace: snafu::Backtrace,
51    },
52
53    #[snafu(display("Failed to read response: {}", source))]
54    ReadResponse {
55        source: iroh::endpoint::ReadToEndError,
56        backtrace: snafu::Backtrace,
57    },
58
59    #[snafu(display("Failed to deserialize response: {}", source))]
60    DeserializeResponse {
61        source: postcard::Error,
62        backtrace: snafu::Backtrace,
63    },
64
65    #[snafu(display("Failed to get remote node id: {}", source))]
66    RemoteNodeId {
67        source: iroh::endpoint::RemoteNodeIdError,
68        backtrace: snafu::Backtrace,
69    },
70}
71
72pub type Result<T> = result::Result<T, Error>;
73
74/// Announce to multiple trackers in parallel.
75pub fn announce_all(
76    endpoint: Endpoint,
77    trackers: impl IntoIterator<Item = NodeId>,
78    signed_announce: SignedAnnounce,
79    announce_parallelism: usize,
80) -> impl Stream<Item = (NodeId, Result<()>)> {
81    n0_future::stream::iter(trackers)
82        .map(move |tracker| {
83            let endpoint = endpoint.clone();
84            async move {
85                let res = announce(&endpoint, tracker, signed_announce).await;
86                (tracker, res)
87            }
88        })
89        .buffered_unordered(announce_parallelism)
90}
91
92/// Announce to a tracker.
93///
94/// You can only announce content you yourself claim to have, to avoid spamming other nodes.
95///
96/// `endpoint` is the iroh endpoint to use for announcing.
97/// `tracker` is the node id of the tracker to announce to. It must understand the [crate::ALPN] protocol.
98/// `content` is the content to announce.
99/// `kind` is the kind of the announcement. We can claim to have the complete data or only some of it.
100pub async fn announce(
101    endpoint: &Endpoint,
102    node_id: NodeId,
103    signed_announce: SignedAnnounce,
104) -> Result<()> {
105    let connecting = endpoint
106        .connect_with_opts(node_id, ALPN, ConnectOptions::default())
107        .await
108        .context(ConnectSnafu)?;
109    match connecting.into_0rtt() {
110        Ok((connection, zero_rtt_accepted)) => {
111            trace!("connected to tracker using possibly 0-rtt: {node_id}");
112            announce_conn(&connection, signed_announce, zero_rtt_accepted).await?;
113            wait_for_session_ticket(connection);
114            Ok(())
115        }
116        Err(connecting) => {
117            let connection = connecting.await.context(Connect1RttSnafu)?;
118            trace!("connected to tracker using 1-rtt: {node_id}");
119            announce_conn(&connection, signed_announce, async { true }).await?;
120            connection.close(0u32.into(), b"");
121            Ok(())
122        }
123    }
124}
125
126/// Announce via an existing connection.
127///
128/// The proceed future can be used to reattempt the send, which can be useful if the connection
129/// was established using 0-rtt and the tracker does not support it. If you have an existing
130/// 1-rtt connection, you can pass `async { true }` to proceed immediately.
131pub async fn announce_conn(
132    connection: &Connection,
133    signed_announce: SignedAnnounce,
134    proceed: impl Future<Output = bool>,
135) -> Result<()> {
136    let (mut send, recv) = connection.open_bi().await.context(OpenStreamSnafu)?;
137    let request = Request::Announce(signed_announce);
138    let request = postcard::to_stdvec(&request).context(SerializeRequestSnafu)?;
139    trace!("sending announce");
140    send.write_all(&request).await.context(WriteRequestSnafu)?;
141    send.finish().context(FinishWriteSnafu)?;
142    let mut recv = if proceed.await {
143        recv
144    } else {
145        let (mut send, recv) = connection.open_bi().await.context(OpenStreamSnafu)?;
146        trace!("re-sending announce using 1-rtt");
147        send.write_all(&request).await.context(WriteRequestSnafu)?;
148        send.finish().context(FinishWriteSnafu)?;
149        recv
150    };
151    let _response = recv
152        .read_to_end(REQUEST_SIZE_LIMIT)
153        .await
154        .context(ReadResponseSnafu)?;
155    trace!("got response");
156    Ok(())
157}
158
159/// A single query to a tracker, using 0-rtt if possible.
160pub async fn query(
161    endpoint: &Endpoint,
162    node_id: NodeId,
163    args: Query,
164) -> Result<Vec<SignedAnnounce>> {
165    let connecting = endpoint
166        .connect_with_opts(node_id, ALPN, ConnectOptions::default())
167        .await
168        .context(ConnectSnafu)?;
169    let result = match connecting.into_0rtt() {
170        Ok((connection, zero_rtt_accepted)) => {
171            trace!("connected to tracker using possibly 0-rtt: {node_id}");
172            let res = query_conn(&connection, args, zero_rtt_accepted).await?;
173            wait_for_session_ticket(connection);
174            res
175        }
176        Err(connecting) => {
177            let connection = connecting.await.context(Connect1RttSnafu)?;
178            trace!("connected to tracker using 1-rtt: {node_id}");
179            let res = query_conn(&connection, args, async { true }).await?;
180            connection.close(0u32.into(), b"");
181            res
182        }
183    };
184    Ok(result.hosts)
185}
186
187/// Query multiple trackers in parallel and merge the results.
188///
189/// You will lose the information about which tracker the results came from, so if you need that,
190/// use [`query`] instead.
191pub fn query_all(
192    endpoint: Endpoint,
193    trackers: impl IntoIterator<Item = NodeId>,
194    args: Query,
195    query_parallelism: usize,
196) -> impl Stream<Item = Result<SignedAnnounce>> {
197    n0_future::stream::iter(trackers)
198        .map(move |tracker| {
199            let endpoint = endpoint.clone();
200            async move {
201                let hosts = match query(&endpoint, tracker, args).await {
202                    Ok(hosts) => hosts.into_iter().map(Ok).collect(),
203                    Err(cause) => vec![Err(cause)],
204                };
205                n0_future::stream::iter(hosts)
206            }
207        })
208        .buffered_unordered(query_parallelism)
209        .flatten()
210}
211
212/// Query via an existing connection.
213///
214/// The proceed future can be used to reattempt the send, which can be useful if the connection
215/// was established using 0-rtt and the tracker does not support it. If you have an existing
216/// 1-rtt connection, you can pass `async { true }` to proceed immediately.
217pub async fn query_conn(
218    connection: &Connection,
219    args: Query,
220    proceed: impl Future<Output = bool>,
221) -> Result<QueryResponse> {
222    let request = Request::Query(args);
223    let request = postcard::to_stdvec(&request).context(SerializeRequestSnafu)?;
224    trace!(
225        "connected to {:?}",
226        connection.remote_node_id().context(RemoteNodeIdSnafu)?
227    );
228    trace!("opened bi stream");
229    let (mut send, recv) = connection.open_bi().await.context(OpenStreamSnafu)?;
230    trace!("sending query");
231    send.write_all(&request).await.context(WriteRequestSnafu)?;
232    send.finish().context(FinishWriteSnafu)?;
233    let mut recv = if proceed.await {
234        recv
235    } else {
236        let (mut send, recv) = connection.open_bi().await.context(OpenStreamSnafu)?;
237        trace!("sending query again using 1-rtt");
238        send.write_all(&request).await.context(WriteRequestSnafu)?;
239        send.finish().context(FinishWriteSnafu)?;
240        recv
241    };
242    let response = recv
243        .read_to_end(REQUEST_SIZE_LIMIT)
244        .await
245        .context(ReadResponseSnafu)?;
246    let response = postcard::from_bytes::<Response>(&response).context(DeserializeResponseSnafu)?;
247    Ok(match response {
248        Response::QueryResponse(response) => response,
249    })
250}
251
252fn wait_for_session_ticket(connection: Connection) {
253    tokio::spawn(async move {
254        // todo: use a more precise API for waiting once it is available.
255        // See https://github.com/quinn-rs/quinn/pull/2257
256        tokio::time::sleep(connection.rtt() * 2).await;
257        connection.close(0u32.into(), b"");
258    });
259}
iroh_content_discovery/client.rs

iroh_content_discovery/
client.rs
