iroh_relay/client/
conn.rs

1//! Manages client-side connections to the relay server.
2//!
3//! based on tailscale/derp/derp_client.go
4
5use std::{
6    pin::Pin,
7    task::{Context, Poll, ready},
8};
9
10use iroh_base::SecretKey;
11use n0_error::{AnyError, anyerr, ensure, stack_error};
12use n0_future::{Sink, Stream};
13use tracing::trace;
14
15use super::KeyCache;
16#[cfg(not(wasm_browser))]
17use crate::client::streams::{MaybeTlsStream, ProxyStream};
18use crate::{
19    MAX_PACKET_SIZE,
20    http::ProtocolVersion,
21    protos::{
22        handshake,
23        relay::{ClientToRelayMsg, Error as ProtoError, RelayToClientMsg},
24        streams::WsBytesFramed,
25    },
26};
27
28/// Error for sending messages to the relay server.
29#[stack_error(derive, add_meta, from_sources, std_sources)]
30#[allow(missing_docs)]
31#[non_exhaustive]
32pub enum SendError {
33    /// Error returned from the underlying WebSocket stream while sending.
34    ///
35    /// The concrete error type is `tokio_websockets::Error` on native targets and
36    /// `ws_stream_wasm::WsErr` on `wasm_browser` targets. Use [`AnyError::downcast_ref`] to
37    /// recover it. Note that the concrete downcast type is not covered by any semver
38    /// guarantees and may change between releases.
39    #[error("Stream error")]
40    StreamError { source: AnyError },
41    #[error("Exceeds max packet size ({MAX_PACKET_SIZE}): {size}")]
42    ExceedsMaxPacketSize { size: usize },
43    #[error("Attempted to send empty packet")]
44    EmptyPacket {},
45}
46
47/// Errors when receiving messages from the relay server.
48#[stack_error(derive, add_meta, from_sources, std_sources)]
49#[allow(missing_docs)]
50#[non_exhaustive]
51pub enum RecvError {
52    #[error(transparent)]
53    Protocol { source: ProtoError },
54    /// Error returned from the underlying WebSocket stream while receiving.
55    ///
56    /// The concrete error type is `tokio_websockets::Error` on native targets and
57    /// `ws_stream_wasm::WsErr` on `wasm_browser` targets. Use [`AnyError::downcast_ref`] to
58    /// recover it. Note that the concrete downcast type is not covered by any semver
59    /// guarantees and may change between releases.
60    #[error("Stream error")]
61    StreamError { source: AnyError },
62}
63
64/// A connection to a relay server.
65///
66/// This holds a connection to a relay server.  It is:
67///
68/// - A [`Stream`] for [`RelayToClientMsg`] to receive from the server.
69/// - A [`Sink`] for [`ClientToRelayMsg`] to send to the server.
70#[derive(derive_more::Debug)]
71pub(crate) struct Conn {
72    #[cfg(not(wasm_browser))]
73    #[debug("tokio_websockets::WebSocketStream")]
74    pub(crate) conn: WsBytesFramed<MaybeTlsStream<ProxyStream>>,
75    #[cfg(wasm_browser)]
76    #[debug("ws_stream_wasm::WsStream")]
77    pub(crate) conn: WsBytesFramed,
78    pub(crate) key_cache: KeyCache,
79    pub(crate) protocol_version: ProtocolVersion,
80}
81
82impl Conn {
83    /// Constructs a new websocket connection, including the initial server handshake.
84    pub(crate) async fn new(
85        #[cfg(not(wasm_browser))] io: tokio_websockets::WebSocketStream<
86            MaybeTlsStream<ProxyStream>,
87        >,
88        #[cfg(wasm_browser)] io: ws_stream_wasm::WsStream,
89        key_cache: KeyCache,
90        secret_key: &SecretKey,
91        protocol_version: ProtocolVersion,
92    ) -> Result<Self, handshake::Error> {
93        let mut conn = WsBytesFramed { io };
94
95        // exchange information with the server
96        trace!("server_handshake: started");
97        handshake::clientside(&mut conn, secret_key).await?;
98        trace!("server_handshake: done");
99
100        Ok(Self {
101            conn,
102            key_cache,
103            protocol_version,
104        })
105    }
106
107    #[cfg(all(test, feature = "server"))]
108    pub(crate) fn test(io: tokio::io::DuplexStream, protocol_version: ProtocolVersion) -> Self {
109        use crate::protos::relay::MAX_FRAME_SIZE;
110        Self {
111            conn: WsBytesFramed {
112                io: tokio_websockets::ClientBuilder::new()
113                    .limits(
114                        tokio_websockets::Limits::default().max_payload_len(Some(MAX_FRAME_SIZE)),
115                    )
116                    .take_over(MaybeTlsStream::Test(io)),
117            },
118            key_cache: KeyCache::test(),
119            protocol_version,
120        }
121    }
122}
123
124impl Stream for Conn {
125    type Item = Result<RelayToClientMsg, RecvError>;
126
127    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
128        match ready!(Pin::new(&mut self.conn).poll_next(cx)) {
129            Some(Ok(msg)) => {
130                let message =
131                    RelayToClientMsg::from_bytes(msg, &self.key_cache, self.protocol_version);
132                Poll::Ready(Some(message.map_err(Into::into)))
133            }
134            Some(Err(e)) => Poll::Ready(Some(Err(anyerr!(e).into()))),
135            None => Poll::Ready(None),
136        }
137    }
138}
139
140impl Sink<ClientToRelayMsg> for Conn {
141    type Error = SendError;
142
143    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
144        Pin::new(&mut self.conn).poll_ready(cx).map_err(Into::into)
145    }
146
147    fn start_send(mut self: Pin<&mut Self>, frame: ClientToRelayMsg) -> Result<(), Self::Error> {
148        let size = frame.encoded_len();
149        ensure!(
150            size <= MAX_PACKET_SIZE,
151            SendError::ExceedsMaxPacketSize { size }
152        );
153        if let ClientToRelayMsg::Datagrams { datagrams, .. } = &frame {
154            ensure!(!datagrams.contents.is_empty(), SendError::EmptyPacket);
155        }
156
157        Pin::new(&mut self.conn)
158            .start_send(frame.to_bytes().freeze())
159            .map_err(Into::into)
160    }
161
162    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
163        Pin::new(&mut self.conn).poll_flush(cx).map_err(Into::into)
164    }
165
166    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
167        Pin::new(&mut self.conn).poll_close(cx).map_err(Into::into)
168    }
169}