diff --git a/crates/buzz-audio-client/src/connection.rs b/crates/buzz-audio-client/src/connection.rs index 193a32c5fe1..6413b41742b 100644 --- a/crates/buzz-audio-client/src/connection.rs +++ b/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 { /// 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( - audio_url: &str, + request: impl IntoClientRequest + Unpin, parent: Option<&str>, connector: Option, 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)) diff --git a/crates/buzz-audio-client/tests/connection.rs b/crates/buzz-audio-client/tests/connection.rs index c67c24b7645..31e9cb66f5a 100644 --- a/crates/buzz-audio-client/tests/connection.rs +++ b/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` 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" + ); +}