feat: pass stream and sink to websocket handler

omjadas · omjadas · commit a8f97f0b8c6a · 2023-04-24T18:25:55.000+10:00
diff --git a/README.md b/README.md
@@ -8,7 +8,7 @@ Hudsucker is a MITM HTTP/S proxy written in Rust that allows you to:
 
 - Modify HTTP/S requests
 - Modify HTTP/S responses
-- Modify websocket messages
+- Modify WebSocket messages
 
 ## Features
 
diff --git a/src/lib.rs b/src/lib.rs
@@ -4,7 +4,7 @@
 //!
 //! - Modify HTTP/S requests
 //! - Modify HTTP/S responses
-//! - Modify websocket messages
+//! - Modify WebSocket messages
 //!
 //! ## Features
 //!
@@ -25,14 +25,16 @@ mod rewind;
 
 pub mod certificate_authority;
 
+use futures::{Sink, SinkExt, Stream, StreamExt};
 use hyper::{Body, Request, Response, StatusCode, Uri};
 use std::net::SocketAddr;
-use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::tungstenite::{self, Message};
 use tracing::error;
 
 pub(crate) use rewind::Rewind;
 
 pub use async_trait;
+pub use futures;
 pub use hyper;
 #[cfg(feature = "openssl-ca")]
 pub use openssl;
@@ -98,7 +100,7 @@ pub enum WebSocketContext {
 /// Each request/response pair is passed to the same instance of the handler.
 #[async_trait::async_trait]
 pub trait HttpHandler: Clone + Send + Sync + 'static {
-    /// The handler will be called for each HTTP request. It can either return a modified request,
+    /// This handler will be called for each HTTP request. It can either return a modified request,
     /// or a response. If a request is returned, it will be sent to the upstream server. If a
     /// response is returned, it will be sent to the client.
     async fn handle_request(
@@ -109,13 +111,13 @@ pub trait HttpHandler: Clone + Send + Sync + 'static {
         req.into()
     }
 
-    /// The handler will be called for each HTTP response. It can modify a response before it is
+    /// This handler will be called for each HTTP response. It can modify a response before it is
     /// forwarded to the client.
     async fn handle_response(&mut self, _ctx: &HttpContext, res: Response<Body>) -> Response<Body> {
         res
     }
 
-    /// The handler will be called if a proxy request fails. Default response is a 502 Bad Gateway.
+    /// This handler will be called if a proxy request fails. Default response is a 502 Bad Gateway.
     async fn handle_error(&mut self, _ctx: &HttpContext, err: hyper::Error) -> Response<Body> {
         error!("Failed to forward request: {}", err);
         Response::builder()
@@ -130,12 +132,48 @@ pub trait HttpHandler: Clone + Send + Sync + 'static {
     }
 }
 
-/// Handler for websocket messages.
+/// Handler for WebSocket messages.
 ///
-/// Messages sent over the same websocket stream are passed to the same instance of the handler.
+/// Messages sent over the same WebSocket Stream are passed to the same instance of the handler.
 #[async_trait::async_trait]
 pub trait WebSocketHandler: Clone + Send + Sync + 'static {
-    /// The handler will be called for each websocket message. It can return an optional modified
+    /// This handler is responsible for forwarding WebSocket messages from a Stream to a Sink and
+    /// recovering from any potential errors.
+    async fn handle_websocket(
+        mut self,
+        ctx: WebSocketContext,
+        mut stream: impl Stream<Item = Result<Message, tungstenite::Error>> + Unpin + Send + 'static,
+        mut sink: impl Sink<Message, Error = tungstenite::Error> + Unpin + Send + 'static,
+    ) {
+        while let Some(message) = stream.next().await {
+            match message {
+                Ok(message) => {
+                    let Some(message) = self.handle_message(&ctx, message).await else {
+                        continue
+                    };
+
+                    match sink.send(message).await {
+                        Err(tungstenite::Error::ConnectionClosed) => (),
+                        Err(e) => error!("WebSocket send error: {}", e),
+                        _ => (),
+                    }
+                }
+                Err(e) => {
+                    error!("WebSocket message error: {}", e);
+
+                    match sink.send(Message::Close(None)).await {
+                        Err(tungstenite::Error::ConnectionClosed) => (),
+                        Err(e) => error!("WebSocket close error: {}", e),
+                        _ => (),
+                    };
+
+                    break;
+                }
+            }
+        }
+    }
+
+    /// This handler will be called for each WebSocket message. It can return an optional modified
     /// message. If None is returned the message will not be forwarded.
     async fn handle_message(
         &mut self,
diff --git a/src/noop.rs b/src/noop.rs
@@ -2,7 +2,7 @@ use crate::{HttpHandler, WebSocketHandler};
 
 /// A No-op handler.
 ///
-/// When using this handler, HTTP requests and responses and websocket messages will not be
+/// When using this handler, HTTP requests and responses and WebSocket messages will not be
 /// modified.
 #[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
 pub struct NoopHandler(());
diff --git a/src/proxy/internal.rs b/src/proxy/internal.rs
@@ -2,7 +2,7 @@ use crate::{
     certificate_authority::CertificateAuthority, HttpContext, HttpHandler, RequestOrResponse,
     Rewind, WebSocketContext, WebSocketHandler,
 };
-use futures::{Sink, SinkExt, Stream, StreamExt};
+use futures::{Sink, Stream, StreamExt};
 use http::uri::{Authority, Scheme};
 use hyper::{
     client::connect::Connect, header::Entry, server::conn::Http, service::service_fn,
@@ -153,7 +153,7 @@ where
                                     if let Err(e) =
                                         self.serve_stream(upgraded, Scheme::HTTP, authority).await
                                     {
-                                        error!("Websocket connect error: {}", e);
+                                        error!("WebSocket connect error: {}", e);
                                     }
 
                                     return;
@@ -252,11 +252,11 @@ where
                     match websocket.await {
                         Ok(ws) => {
                             if let Err(e) = self.handle_websocket(ws, req).await {
-                                error!("Failed to handle websocket: {}", e);
+                                error!("Failed to handle WebSocket: {}", e);
                             }
                         }
                         Err(e) => {
-                            error!("Failed to upgrade to websocket: {}", e);
+                            error!("Failed to upgrade to WebSocket: {}", e);
                         }
                     }
                 };
@@ -350,41 +350,13 @@ where
 }
 
 fn spawn_message_forwarder(
-    mut stream: impl Stream<Item = Result<Message, tungstenite::Error>> + Unpin + Send + 'static,
-    mut sink: impl Sink<Message, Error = tungstenite::Error> + Unpin + Send + 'static,
-    mut handler: impl WebSocketHandler,
+    stream: impl Stream<Item = Result<Message, tungstenite::Error>> + Unpin + Send + 'static,
+    sink: impl Sink<Message, Error = tungstenite::Error> + Unpin + Send + 'static,
+    handler: impl WebSocketHandler,
     ctx: WebSocketContext,
 ) {
     let span = info_span!("message_forwarder", context = ?ctx);
-    let fut = async move {
-        while let Some(message) = stream.next().await {
-            match message {
-                Ok(message) => {
-                    let Some(message) = handler.handle_message(&ctx, message).await else {
-                        continue
-                    };
-
-                    match sink.send(message).await {
-                        Err(tungstenite::Error::ConnectionClosed) => (),
-                        Err(e) => error!("Websocket send error: {}", e),
-                        _ => (),
-                    }
-                }
-                Err(e) => {
-                    error!("Websocket message error: {}", e);
-
-                    match sink.send(Message::Close(None)).await {
-                        Err(tungstenite::Error::ConnectionClosed) => (),
-                        Err(e) => error!("Websocket close error: {}", e),
-                        _ => (),
-                    };
-
-                    break;
-                }
-            }
-        }
-    };
-
+    let fut = handler.handle_websocket(ctx, stream, sink);
     spawn_with_trace(fut, span);
 }
 
