iroh_relay/protos/
streams.rs

1//! Implements logic for abstracting over a websocket stream that allows sending only [`Bytes`]-based
2//! messages.
3use std::{
4    pin::Pin,
5    task::{Context, Poll},
6};
7
8use bytes::Bytes;
9use n0_error::AnyError;
10#[cfg(not(wasm_browser))]
11use n0_error::anyerr;
12use n0_future::{Sink, Stream, ready};
13#[cfg(not(wasm_browser))]
14use tokio::io::{AsyncRead, AsyncWrite};
15use tracing::warn;
16
17use crate::ExportKeyingMaterial;
18
19#[cfg(not(wasm_browser))]
20#[derive(derive_more::Debug)]
21pub(crate) struct WsBytesFramed<T> {
22    #[debug("WebSocketStream<T>")]
23    pub(crate) io: tokio_websockets::WebSocketStream<T>,
24}
25
26#[cfg(wasm_browser)]
27#[derive(derive_more::Debug)]
28pub(crate) struct WsBytesFramed {
29    #[debug("WebSocketStream")]
30    pub(crate) io: ws_stream_wasm::WsStream,
31}
32
33/// Error type for WebSocket stream operations.
34///
35/// This is a generic [`AnyError`] and represents errors that can occur during WebSocket
36/// communication. The underlying error type depends on the platform:
37/// - On non-browser platforms: `tokio_websockets::Error`
38/// - On browser WASM: `ws_stream_wasm::WsErr`
39///
40/// You can use [`AnyError::downcast_ref`] to downcast to the concrete error. However, the
41/// downcast is not covered by any semver guarantees and the underlying error may change.
42pub type StreamError = AnyError;
43
44/// Shorthand for a type that implements both a websocket-based stream & sink for [`Bytes`].
45pub trait BytesStreamSink:
46    Stream<Item = Result<Bytes, StreamError>> + Sink<Bytes, Error = StreamError> + Unpin
47{
48}
49
50impl<T> BytesStreamSink for T where
51    T: Stream<Item = Result<Bytes, StreamError>> + Sink<Bytes, Error = StreamError> + Unpin
52{
53}
54
55#[cfg(not(wasm_browser))]
56impl<IO: ExportKeyingMaterial + AsyncRead + AsyncWrite + Unpin> ExportKeyingMaterial
57    for WsBytesFramed<IO>
58{
59    fn export_keying_material<T: AsMut<[u8]>>(
60        &self,
61        output: T,
62        label: &[u8],
63        context: Option<&[u8]>,
64    ) -> Option<T> {
65        self.io
66            .get_ref()
67            .export_keying_material(output, label, context)
68    }
69}
70
71#[cfg(wasm_browser)]
72impl ExportKeyingMaterial for WsBytesFramed {
73    fn export_keying_material<T: AsMut<[u8]>>(
74        &self,
75        _output: T,
76        _label: &[u8],
77        _context: Option<&[u8]>,
78    ) -> Option<T> {
79        None
80    }
81}
82
83#[cfg(not(wasm_browser))]
84impl<T: AsyncRead + AsyncWrite + Unpin> Stream for WsBytesFramed<T> {
85    type Item = Result<Bytes, StreamError>;
86
87    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
88        loop {
89            match ready!(Pin::new(&mut self.io).poll_next(cx)) {
90                None => return Poll::Ready(None),
91                Some(Err(e)) => return Poll::Ready(Some(Err(anyerr!(e)))),
92                Some(Ok(msg)) => {
93                    if msg.is_close() {
94                        // Indicate the stream is done when we receive a close message.
95                        // Note: We don't have to poll the stream to completion for it to close gracefully.
96                        return Poll::Ready(None);
97                    }
98                    if msg.is_ping() || msg.is_pong() {
99                        continue; // Responding appropriately to these is done inside of tokio_websockets/browser impls
100                    }
101                    if !msg.is_binary() {
102                        warn!(?msg, "Got websocket message of unsupported type, skipping.");
103                        continue;
104                    }
105                    return Poll::Ready(Some(Ok(msg.into_payload().into())));
106                }
107            }
108        }
109    }
110}
111
112#[cfg(wasm_browser)]
113impl Stream for WsBytesFramed {
114    type Item = Result<Bytes, StreamError>;
115
116    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
117        loop {
118            match ready!(Pin::new(&mut self.io).poll_next(cx)) {
119                None => return Poll::Ready(None),
120                Some(ws_stream_wasm::WsMessage::Binary(msg)) => {
121                    return Poll::Ready(Some(Ok(msg.into())));
122                }
123                Some(msg) => {
124                    warn!(?msg, "Got websocket message of unsupported type, skipping.");
125                    continue;
126                }
127            }
128        }
129    }
130}
131
132#[cfg(not(wasm_browser))]
133impl<T: AsyncRead + AsyncWrite + Unpin> Sink<Bytes> for WsBytesFramed<T> {
134    type Error = StreamError;
135
136    fn start_send(mut self: Pin<&mut Self>, bytes: Bytes) -> Result<(), Self::Error> {
137        let msg = tokio_websockets::Message::binary(tokio_websockets::Payload::from(bytes));
138        Pin::new(&mut self.io)
139            .start_send(msg)
140            .map_err(AnyError::from_std)
141    }
142
143    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
144        Pin::new(&mut self.io)
145            .poll_ready(cx)
146            .map_err(AnyError::from_std)
147    }
148
149    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
150        Pin::new(&mut self.io)
151            .poll_flush(cx)
152            .map_err(AnyError::from_std)
153    }
154
155    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
156        Pin::new(&mut self.io)
157            .poll_close(cx)
158            .map_err(AnyError::from_std)
159    }
160}
161
162#[cfg(wasm_browser)]
163impl Sink<Bytes> for WsBytesFramed {
164    type Error = StreamError;
165
166    fn start_send(mut self: Pin<&mut Self>, bytes: Bytes) -> Result<(), Self::Error> {
167        let msg = ws_stream_wasm::WsMessage::Binary(Vec::from(bytes));
168        Pin::new(&mut self.io)
169            .start_send(msg)
170            .map_err(AnyError::from_std)
171    }
172
173    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
174        Pin::new(&mut self.io)
175            .poll_ready(cx)
176            .map_err(AnyError::from_std)
177    }
178
179    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
180        Pin::new(&mut self.io)
181            .poll_flush(cx)
182            .map_err(AnyError::from_std)
183    }
184
185    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
186        Pin::new(&mut self.io)
187            .poll_close(cx)
188            .map_err(AnyError::from_std)
189    }
190}
iroh_relay/protos/streams.rs

iroh_relay/protos/
streams.rs
