feat(audio-client): accept a client request in connect

crates/buzz-audio-client/src/connection.rs

@@ -5,7 +5,7 @@ use serde_json::{json, Value};
 use std::{collections::BTreeSet, future::Future, time::Duration};
 use tokio_tungstenite::{
     connect_async_tls_with_config,
-    tungstenite::{protocol::WebSocketConfig, Message},
+    tungstenite::{client::IntoClientRequest, protocol::WebSocketConfig, Message},
     Connector, MaybeTlsStream, WebSocketStream,
 };
 
@@ -99,11 +99,12 @@ async fn control(socket: &mut AudioSocket) -> Result<Value, String> {
 /// Connect and authenticate using a host-provided asynchronous signer.
 ///
 /// The host validates the destination before calling, and constructs the NIP-42
-/// event (including its exact relay tag) from the challenge. No private key is
+/// event (including its exact relay tag) from the challenge. `request` is the
+/// audio URL, or a request carrying host-supplied upgrade headers. No private key is
 /// accepted. Dropping this future cancels connection/admission; the entire
 /// operation is bounded to 15 seconds, including signing.
 pub async fn connect<F, Fut>(
-    audio_url: &str,
+    request: impl IntoClientRequest + Unpin,
     parent: Option<&str>,
     connector: Option<Connector>,
     sign: F,
@@ -117,7 +118,7 @@ where
             .max_message_size(Some(65_536))
             .max_frame_size(Some(65_536));
         let (mut socket, _) =
-            connect_async_tls_with_config(audio_url, Some(config), false, connector)
+            connect_async_tls_with_config(request, Some(config), false, connector)
                 .await
                 .map_err(|e| format!("audio WS connect failed: {e}"))?;
         let challenge = tokio::time::timeout(Duration::from_secs(5), control(&mut socket))

crates/buzz-audio-client/tests/connection.rs

@@ -3,7 +3,14 @@ use futures_util::{SinkExt, StreamExt};
 use serde_json::{json, Value};
 use std::time::Duration;
 use tokio::{net::TcpListener, sync::oneshot};
-use tokio_tungstenite::{accept_async, tungstenite::Message};
+use tokio_tungstenite::{
+    accept_async, accept_hdr_async,
+    tungstenite::{
+        client::IntoClientRequest,
+        handshake::server::{Request, Response},
+        Message,
+    },
+};
 
 #[tokio::test]
 async fn authenticates_with_host_signer_and_preserves_v2_media() {
@@ -131,3 +138,52 @@ fn rejects_invalid_rosters_without_truncation() {
     )
     .is_err());
 }
+
+#[tokio::test]
+// tungstenite fixes the accept callback's `Result<Response, ErrorResponse>` signature.
+#[allow(clippy::result_large_err)]
+async fn request_headers_reach_the_upgrade() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let url = format!("ws://{}/huddle/room/audio", listener.local_addr().unwrap());
+    let server = tokio::spawn(async move {
+        let (stream, _) = listener.accept().await.unwrap();
+        let (seen, header) = oneshot::channel();
+        let mut socket = accept_hdr_async(stream, |request: &Request, response: Response| {
+            let _ = seen.send(request.headers().get("authorization").cloned());
+            Ok(response)
+        })
+        .await
+        .unwrap();
+        socket
+            .send(Message::Text(
+                json!({"type":"challenge","challenge":"c"})
+                    .to_string()
+                    .into(),
+            ))
+            .await
+            .unwrap();
+        socket.next().await.unwrap().unwrap();
+        socket
+            .send(Message::Text(
+                json!({"type":"joined","peer_index":0,"peers":[]})
+                    .to_string()
+                    .into(),
+            ))
+            .await
+            .unwrap();
+        header.await.unwrap()
+    });
+    let mut request = url.into_client_request().unwrap();
+    request
+        .headers_mut()
+        .insert("authorization", "Bearer upgrade-token".parse().unwrap());
+    connect(request, None, None, |challenge| async move {
+        Ok(json!({ "signed": challenge }))
+    })
+    .await
+    .unwrap();
+    assert_eq!(
+        server.await.unwrap().unwrap().to_str().unwrap(),
+        "Bearer upgrade-token"
+    );
+}