first pass of merging in warp (doesn't build)

crates/websocket/src/lib.rs

@@ -14,6 +14,8 @@ mod sink_map_err;
 
 use anyhow::anyhow;
 #[cfg(not(target_family = "wasm"))]
+pub use async_tungstenite::tungstenite;
+#[cfg(not(target_family = "wasm"))]
 pub use async_tungstenite::tungstenite::client::IntoClientRequest;
 #[cfg(not(target_family = "wasm"))]
 use async_tungstenite::tungstenite::http::HeaderValue;
@@ -22,9 +24,6 @@ use futures_util::{future, SinkExt, TryStreamExt};
 use itertools::Itertools;
 use thiserror::Error;
 
-#[cfg(not(target_family = "wasm"))]
-pub use async_tungstenite::tungstenite;
-
 use crate::sink_map_err::map_err;
 
 // Unfortunately, `anyhow::Error` does not implement `std::error::Error`, which is required by the
@@ -128,11 +127,70 @@ impl WebSocket {
         Ok(Self(socket))
     }
 
+    /// Connect to `url`, attaching the provided extra request headers to the
+    /// HTTP upgrade handshake.
+    ///
+    /// Custom handshake headers are not supported on wasm websockets, so
+    /// `headers` is ignored there.
+    pub async fn connect_with_headers<'a>(
+        url: &str,
+        protocols: impl IntoIterator<Item = &'a str>,
+        headers: Vec<(&'a str, String)>,
+    ) -> anyhow::Result<Self> {
+        cfg_if::cfg_if! {
+            if #[cfg(not(target_family = "wasm"))] {
+                // `connect` accepts any `IntoClientRequest` and handles
+                // protocol negotiation + wrapping, so build the enriched
+                // request and hand it off.
+                let request = build_request_with_headers(url, headers)?;
+                Self::connect(request, protocols).await
+            } else {
+                let _ = headers;
+                Self::connect(url, protocols).await
+            }
+        }
+    }
+
     pub async fn into_graphql_client_builder(self) -> graphql_ws_client::ClientBuilder {
         self.0.into_graphql_client_builder().await
     }
 }
 
+/// If `err` originated from a websocket handshake that received a non-101 HTTP
+/// response (e.g. an auth or proxy challenge), returns that response so callers
+/// can inspect its status and headers (for example, to detect a GCP IAP
+/// challenge). Native-only: wasm websockets do not surface the handshake
+/// response on error.
+#[cfg(not(target_family = "wasm"))]
+pub fn connect_error_http_response(
+    err: &anyhow::Error,
+) -> Option<&tungstenite::http::Response<Option<Vec<u8>>>> {
+    err.chain()
+        .find_map(|cause| match cause.downcast_ref::<tungstenite::Error>() {
+            Some(tungstenite::Error::Http(response)) => Some(response),
+            _ => None,
+        })
+}
+
+/// Builds a native websocket client request for `url`, attaching the provided
+/// extra request headers to the handshake.
+#[cfg(not(target_family = "wasm"))]
+fn build_request_with_headers(
+    url: &str,
+    headers: Vec<(&str, String)>,
+) -> anyhow::Result<tungstenite::handshake::client::Request> {
+    let mut request = url.into_client_request()?;
+    let request_headers = request.headers_mut();
+    for (name, value) in headers {
+        let header_name = http::HeaderName::from_bytes(name.as_bytes())
+            .map_err(|err| anyhow!("invalid websocket header name `{name}`: {err}"))?;
+        let header_value = http::HeaderValue::from_str(&value)
+            .map_err(|err| anyhow!("invalid websocket header value: {err}"))?;
+        request_headers.insert(header_name, header_value);
+    }
+    Ok(request)
+}
+
 /// Trait that defines a [`Sink`] returned by the websocket.
 pub trait Sink: futures::Sink<Message, Error = Error> + Send + Unpin + 'static {}
 

crates/websocket/src/native.rs

@@ -2,13 +2,11 @@
 
 use std::sync::Arc;
 
-use async_tungstenite::{
-    tokio::{
-        client_async_tls_with_connector_and_config, connect_async_with_tls_connector, ClientStream,
-    },
-    tungstenite::client::IntoClientRequest,
-    WebSocketStream,
+use async_tungstenite::tokio::{
+    client_async_tls_with_connector_and_config, connect_async_with_tls_connector, ClientStream,
 };
+use async_tungstenite::tungstenite::client::IntoClientRequest;
+use async_tungstenite::WebSocketStream;
 use futures::{Sink, Stream};
 use futures_util::StreamExt as _;
 use rustls_platform_verifier::ConfigVerifierExt;

crates/websocket/src/proxy.rs

@@ -10,7 +10,8 @@ use std::env;
 use std::time::Duration;
 
 use anyhow::{bail, Context};
-use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
+use base64::engine::general_purpose::STANDARD as BASE64;
+use base64::Engine as _;
 use http_body_util::Empty;
 use hyper::body::Bytes;
 use hyper_util::rt::TokioIo;

crates/websocket/src/sink_map_err.rs

@@ -1,8 +1,9 @@
 use core::pin::Pin;
+use std::task::{Context, Poll};
+
 use futures::{Sink, Stream};
 use futures_util::stream::FusedStream;
 use pin_project::pin_project;
-use std::task::{Context, Poll};
 
 /// Maps the error returned by the [`Sink`] using the provided `err` function.
 pub fn map_err<S, I, E>(sink: S, err: impl FnMut(S::Error) -> E) -> impl Sink<I, Error = E>
@@ -13,7 +14,7 @@ where
 }
 
 /// Helper struct to map the [`Err`] of an underlying [`Sink`].
-/// This is a fork of the the `SinkMapErr` defined within `futures-util`
+/// This is a fork of the `SinkMapErr` defined within `futures-util`
 /// (https://docs.rs/futures/latest/futures/sink/struct.SinkMapErr.html) except that it does _not_
 /// panic if the caller tries to write to the sink after a previous attempt to write returned an
 /// error. See <https://github.com/rust-lang/futures-rs/issues/2108> for more details about the

crates/websocket/src/sink_map_err_tests.rs

@@ -1,8 +1,9 @@
-use super::*;
 use futures_test_sink::SinkMock;
 use futures_util::SinkExt;
 use thiserror::Error;
 
+use super::*;
+
 #[derive(Error, PartialEq, Debug, Copy, Clone)]
 #[error("Unable to send to sink")]
 struct UnmappedError(u8);

crates/websocket/src/wasm.rs

@@ -1,8 +1,7 @@
 use futures::{Sink, Stream, StreamExt};
 use itertools::Itertools;
-use ws_stream_wasm::{WsErr, WsMessage, WsMeta};
-
 pub use ws_stream_wasm::WsMessage as Message;
+use ws_stream_wasm::{WsErr, WsMessage, WsMeta};
 
 use crate::WebsocketMessage;