iroh_relay/client/
conn.rs1use 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#[stack_error(derive, add_meta, from_sources, std_sources)]
30#[allow(missing_docs)]
31#[non_exhaustive]
32pub enum SendError {
33 #[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#[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("Stream error")]
61 StreamError { source: AnyError },
62}
63
64#[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 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 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}