feat(iroh-net): Implement websocket protocol upgrade in iroh-relay

iroh-blobs/src/downloader/test.rs

@@ -10,7 +10,7 @@ use iroh_net::key::SecretKey;
 
 use crate::{
     get::{db::BlobId, progress::TransferState},
-    util::progress::{FlumeProgressSender, IdGenerator, ProgressSender},
+    util::progress::{FlumeProgressSender, IdGenerator},
 };
 
 use super::*;

iroh-blobs/src/downloader/test/dialer.rs

@@ -1,9 +1,6 @@
 //! Implementation of [`super::Dialer`] used for testing.
 
-use std::{
-    collections::HashSet,
-    task::{Context, Poll},
-};
+use std::task::{Context, Poll};
 
 use parking_lot::RwLock;
 

iroh-blobs/src/downloader/test/getter.rs

@@ -2,7 +2,6 @@
 
 use futures_lite::{future::Boxed as BoxFuture, FutureExt};
 use parking_lot::RwLock;
-use std::{sync::Arc, time::Duration};
 
 use super::*;
 

iroh-blobs/src/store/bao_file.rs

@@ -732,7 +732,6 @@ pub mod test_support {
         BlockSize, ChunkRanges,
     };
     use futures_lite::{Stream, StreamExt};
-    use iroh_base::hash::Hash;
     use iroh_io::AsyncStreamReader;
     use rand::RngCore;
     use range_collections::RangeSet2;

iroh-cli/src/commands/doctor.rs

@@ -44,7 +44,6 @@ use iroh::{
 };
 use portable_atomic::AtomicU64;
 use postcard::experimental::max_size::MaxSize;
-use ratatui::backend::Backend;
 use serde::{Deserialize, Serialize};
 use tokio::{io::AsyncWriteExt, sync};
 

iroh-net/Cargo.toml

@@ -88,6 +88,9 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = tr
 # metrics
 iroh-metrics = { version = "0.18.0", path = "../iroh-metrics", default-features = false }
 strum = { version = "0.26.2", features = ["derive"] }
+tungstenite = "0.23.0"
+tokio-tungstenite = "0.23.1"
+tokio-tungstenite-wasm = "0.3.1"
 
 [target.'cfg(any(target_os = "linux", target_os = "android"))'.dependencies]
 netlink-packet-core = "0.7.0"

iroh-net/src/magicsock.rs

@@ -2738,7 +2738,6 @@ impl NetInfo {
 #[cfg(test)]
 mod tests {
     use anyhow::Context;
-    use futures_lite::StreamExt;
     use iroh_test::CallOnDrop;
     use rand::RngCore;
 

iroh-net/src/relay.rs

@@ -13,7 +13,7 @@
 pub(crate) mod client;
 pub(crate) mod client_conn;
 pub(crate) mod clients;
-mod codec;
+pub(crate) mod codec;
 pub mod http;
 pub mod iroh_relay;
 mod map;

iroh-net/src/relay/client.rs

@@ -1,15 +1,18 @@
 //! based on tailscale/derp/derp_client.go
 use std::net::SocketAddr;
+use std::pin::Pin;
 use std::sync::Arc;
+use std::task::{Context, Poll};
 use std::time::Duration;
 
 use anyhow::{anyhow, bail, ensure, Result};
 use bytes::Bytes;
-use futures_lite::StreamExt;
+use futures_lite::Stream;
 use futures_sink::Sink;
-use futures_util::sink::SinkExt;
-use tokio::io::AsyncWrite;
+use futures_util::stream::{SplitSink, SplitStream, StreamExt};
+use futures_util::SinkExt;
 use tokio::sync::mpsc;
+use tokio_tungstenite_wasm::WebSocketStream;
 use tokio_util::codec::{FramedRead, FramedWrite};
 use tracing::{debug, info_span, trace, Instrument};
 
@@ -64,12 +67,12 @@ impl ClientReceiver {
     }
 }
 
-type RelayReader = FramedRead<MaybeTlsStreamReader, DerpCodec>;
-
 #[derive(derive_more::Debug)]
 pub struct InnerClient {
-    // our local addrs
-    local_addr: SocketAddr,
+    /// Our local address, if known.
+    ///
+    /// Is `None` in tests or when using websockets (because we don't control connection establishment in browsers).
+    local_addr: Option<SocketAddr>,
     /// Channel on which to communicate to the server. The associated [`mpsc::Receiver`] will close
     /// if there is ever an error writing to the server.
     writer_channel: mpsc::Sender<ClientWriterMessage>,
@@ -123,8 +126,10 @@ impl Client {
     }
 
     /// The local address that the [`Client`] is listening on.
-    pub fn local_addr(&self) -> Result<SocketAddr> {
-        Ok(self.inner.local_addr)
+    ///
+    /// `None`, when run in a testing environment or when using websockets.
+    pub fn local_addr(&self) -> Option<SocketAddr> {
+        self.inner.local_addr
     }
 
     /// Whether or not this [`Client`] is closed.
@@ -209,13 +214,13 @@ enum ClientWriterMessage {
 ///
 /// Shutsdown when you send a [`ClientWriterMessage::Shutdown`], or if there is an error writing to
 /// the server.
-struct ClientWriter<W: AsyncWrite + Unpin + Send + 'static> {
+struct ClientWriter {
     recv_msgs: mpsc::Receiver<ClientWriterMessage>,
-    writer: FramedWrite<W, DerpCodec>,
+    writer: RelayConnWriter,
     rate_limiter: Option<RateLimiter>,
 }
 
-impl<W: AsyncWrite + Unpin + Send + 'static> ClientWriter<W> {
+impl ClientWriter {
     async fn run(mut self) -> Result<()> {
         while let Some(msg) = self.recv_msgs.recv().await {
             match msg {
@@ -244,25 +249,92 @@ impl<W: AsyncWrite + Unpin + Send + 'static> ClientWriter<W> {
     }
 }
 
-/// The Builder returns a [`Client`] starts a [`ClientWriter`] run task.
+/// The Builder returns a [`Client`] and a started [`ClientWriter`] run task.
 pub struct ClientBuilder {
     secret_key: SecretKey,
-    reader: RelayReader,
-    writer: FramedWrite<MaybeTlsStreamWriter, DerpCodec>,
-    local_addr: SocketAddr,
+    reader: RelayConnReader,
+    writer: RelayConnWriter,
+    local_addr: Option<SocketAddr>,
+}
+
+pub(crate) enum RelayConnReader {
+    Derp(FramedRead<MaybeTlsStreamReader, DerpCodec>),
+    Ws(SplitStream<WebSocketStream>),
+}
+
+pub(crate) enum RelayConnWriter {
+    Derp(FramedWrite<MaybeTlsStreamWriter, DerpCodec>),
+    Ws(SplitSink<WebSocketStream, tokio_tungstenite_wasm::Message>),
+}
+
+fn tung_wasm_to_io_err(e: tokio_tungstenite_wasm::Error) -> std::io::Error {
+    std::io::Error::new(std::io::ErrorKind::Other, e.to_string())
+}
+
+impl Stream for RelayConnReader {
+    type Item = Result<Frame>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match *self {
+            Self::Derp(ref mut ws) => Pin::new(ws).poll_next(cx),
+            Self::Ws(ref mut ws) => match Pin::new(ws).poll_next(cx) {
+                Poll::Ready(Some(item)) => match Frame::from_wasm_ws_message(item) {
+                    Some(frame) => Poll::Ready(Some(frame)),
+                    // We filter websocket messages that don't carry `Frame`s.
+                    None => Poll::Pending,
+                },
+                Poll::Ready(None) => Poll::Ready(None),
+                Poll::Pending => Poll::Pending,
+            },
+        }
+    }
+}
+
+impl Sink<Frame> for RelayConnWriter {
+    type Error = std::io::Error;
+
+    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        match *self {
+            Self::Derp(ref mut ws) => Pin::new(ws).poll_ready(cx),
+            Self::Ws(ref mut ws) => Pin::new(ws).poll_ready(cx).map_err(tung_wasm_to_io_err),
+        }
+    }
+
+    fn start_send(mut self: Pin<&mut Self>, item: Frame) -> Result<(), Self::Error> {
+        match *self {
+            Self::Derp(ref mut ws) => Pin::new(ws).start_send(item),
+            Self::Ws(ref mut ws) => Pin::new(ws)
+                .start_send(item.into_wasm_ws_message()?)
+                .map_err(tung_wasm_to_io_err),
+        }
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        match *self {
+            Self::Derp(ref mut ws) => Pin::new(ws).poll_flush(cx),
+            Self::Ws(ref mut ws) => Pin::new(ws).poll_flush(cx).map_err(tung_wasm_to_io_err),
+        }
+    }
+
+    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        match *self {
+            Self::Derp(ref mut ws) => Pin::new(ws).poll_close(cx),
+            Self::Ws(ref mut ws) => Pin::new(ws).poll_close(cx).map_err(tung_wasm_to_io_err),
+        }
+    }
 }
 
 impl ClientBuilder {
     pub fn new(
         secret_key: SecretKey,
-        local_addr: SocketAddr,
-        reader: MaybeTlsStreamReader,
-        writer: MaybeTlsStreamWriter,
+        local_addr: Option<SocketAddr>,
+        reader: RelayConnReader,
+        writer: RelayConnWriter,
     ) -> Self {
         Self {
             secret_key,
-            reader: FramedRead::new(reader, DerpCodec),
-            writer: FramedWrite::new(writer, DerpCodec),
+            reader,
+            writer,
             local_addr,
         }
     }