Add websocket proxying support

Cargo.toml

@@ -68,6 +68,7 @@ alloy-transport-http = { version = "1.4.3", features = [
   "hyper",
 ], optional = true }
 url = { version = "2.5.7", optional = true }
+hyper-tungstenite = "0.19.0"
 
 [dev-dependencies]
 tempfile = "3.23.0"

src/attested_get.rs

@@ -1,5 +1,7 @@
 //! A one-shot attested TLS proxy client which sends a single GET request and returns the response
-use crate::{AttestationGenerator, AttestationVerifier, ProxyClient, ProxyError};
+use crate::{
+    http_version::HttpVersion, AttestationGenerator, AttestationVerifier, ProxyClient, ProxyError,
+};
 use tokio_rustls::rustls::pki_types::CertificateDer;
 
 /// Start a proxy-client, send a single HTTP GET request to the given path and return the
@@ -17,6 +19,7 @@ pub async fn attested_get(
         AttestationGenerator::with_no_attestation(),
         attestation_verifier,
         remote_certificate,
+        HttpVersion::Http2,
     )
     .await?;
 
@@ -103,6 +106,7 @@ mod tests {
             AttestationGenerator::with_no_attestation(),
             AttestationVerifier::mock(),
             None,
+            HttpVersion::Http2,
         )
         .await
         .unwrap();

src/file_server.rs

@@ -52,7 +52,7 @@ pub(crate) async fn static_file_server(path: PathBuf) -> Result<SocketAddr, Prox
 
 #[cfg(test)]
 mod tests {
-    use crate::{attestation::AttestationType, ProxyClient};
+    use crate::{attestation::AttestationType, http_version::HttpVersion, ProxyClient};
 
     use super::*;
     use crate::test_helpers::{generate_certificate_chain, generate_tls_config};
@@ -122,6 +122,7 @@ mod tests {
             AttestationGenerator::with_no_attestation(),
             AttestationVerifier::mock(),
             None,
+            HttpVersion::Http2,
         )
         .await
         .unwrap();

src/http_version.rs

@@ -1,8 +1,7 @@
 //! HTTP Version support and negotiation
 use hyper::Response;
-use hyper_util::rt::TokioIo;
+use std::future::Future;
 use std::pin::Pin;
-use std::task::{Context, Poll};
 
 pub const ALPN_H2: &[u8] = b"h2";
 pub const ALPN_HTTP11: &[u8] = b"http/1.1";
@@ -47,17 +46,6 @@ impl HttpVersion {
 type Http1Sender = hyper::client::conn::http1::SendRequest<hyper::body::Incoming>;
 type Http2Sender = hyper::client::conn::http2::SendRequest<hyper::body::Incoming>;
 
-type Http1Connection = hyper::client::conn::http1::Connection<
-    TokioIo<tokio_rustls::client::TlsStream<tokio::net::TcpStream>>,
-    hyper::body::Incoming,
->;
-
-type Http2Connection = hyper::client::conn::http2::Connection<
-    TokioIo<tokio_rustls::client::TlsStream<tokio::net::TcpStream>>,
-    hyper::body::Incoming,
-    crate::TokioExecutor,
->;
-
 /// A protocol version agnostic HTTP sender
 pub enum HttpSender {
     Http1(Http1Sender),
@@ -88,34 +76,5 @@ impl HttpSender {
     }
 }
 
-pin_project_lite::pin_project! {
-    /// A protocol version agnostic HTTP connection
-    #[project = HttpConnectionProj]
-    pub enum HttpConnection {
-        Http1 { #[pin] inner: Http1Connection },
-        Http2 { #[pin] inner: Http2Connection },
-    }
-}
-
-impl From<Http1Connection> for HttpConnection {
-    fn from(inner: Http1Connection) -> Self {
-        Self::Http1 { inner }
-    }
-}
-
-impl From<Http2Connection> for HttpConnection {
-    fn from(inner: Http2Connection) -> Self {
-        Self::Http2 { inner }
-    }
-}
-
-impl Future for HttpConnection {
-    type Output = Result<(), hyper::Error>;
-
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        match self.project() {
-            HttpConnectionProj::Http1 { inner } => inner.poll(cx),
-            HttpConnectionProj::Http2 { inner } => inner.poll(cx),
-        }
-    }
-}
+/// A protocol version agnostic HTTP connection future
+pub type HttpConnection = Pin<Box<dyn Future<Output = Result<(), hyper::Error>> + Send>>;