From 77925a4d65327f3984e01dc073077b50c81f43e1 Mon Sep 17 00:00:00 2001
From: hatoo <hato2000@gmail.com>
Date: Sun, 6 Oct 2024 16:06:38 +0900
Subject: [PATCH] Add doc for upgraded connection

---
 README.md             | 32 +++++++++++++++++++++++++++++---
 examples/proxy.rs     | 32 +++++++++++++++++++++++++++++---
 src/default_client.rs |  1 +
 3 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index c9292ff..6562952 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,8 @@ A HTTP proxy server library intended to be a backend of application like Burp pr
 use std::path::PathBuf;
 
 use clap::{Args, Parser};
-use http_mitm_proxy::{DefaultClient, MitmProxy};
+use futures::StreamExt;
+use http_mitm_proxy::{default_client::Upgrade, DefaultClient, MitmProxy};
 use moka::sync::Cache;
 use tracing_subscriber::EnvFilter;
 
@@ -100,11 +101,36 @@ async fn main() {
                 let uri = req.uri().clone();
 
                 // You can modify request here
-                // or You can just return response anyware
+                // or You can just return response anywhere
 
-                let (res, _upgrade) = client.send_request(req).await?;
+                let (res, upgrade) = client.send_request(req).await?;
 
                 println!("{} -> {}", uri, res.status());
+                if let Some(upgrade) = upgrade {
+                    // If the response is an upgrade, e.g. Websocket, you can see traffic.
+                    // Modifying upgraded traffic is not supported yet.
+
+                    // You can try https://echo.websocket.org/.ws to test websocket.
+                    // But you need to disable alpn of DefaultClient to disable HTTP2 because echo.websocket.org does not support HTTP/2 for Websocket.
+                    // It should be match incoming and outgoing HTTP version on DefaultClient, I'll fix this later. #54
+                    println!("Upgrade connection");
+                    let Upgrade {
+                        mut client_to_server,
+                        mut server_to_client,
+                    } = upgrade;
+                    let url = uri.to_string();
+                    tokio::spawn(async move {
+                        while let Some(data) = client_to_server.next().await {
+                            println!("Client -> Server: {} {:?}", url, data);
+                        }
+                    });
+                    let url = uri.to_string();
+                    tokio::spawn(async move {
+                        while let Some(data) = server_to_client.next().await {
+                            println!("Server -> Client: {} {:?}", url, data);
+                        }
+                    });
+                }
 
                 // You can modify response here
 
diff --git a/examples/proxy.rs b/examples/proxy.rs
index 2d69c3a..35752bb 100644
--- a/examples/proxy.rs
+++ b/examples/proxy.rs
@@ -1,7 +1,8 @@
 use std::path::PathBuf;
 
 use clap::{Args, Parser};
-use http_mitm_proxy::{DefaultClient, MitmProxy};
+use futures::StreamExt;
+use http_mitm_proxy::{default_client::Upgrade, DefaultClient, MitmProxy};
 use moka::sync::Cache;
 use tracing_subscriber::EnvFilter;
 
@@ -87,11 +88,36 @@ async fn main() {
                 let uri = req.uri().clone();
 
                 // You can modify request here
-                // or You can just return response anyware
+                // or You can just return response anywhere
 
-                let (res, _upgrade) = client.send_request(req).await?;
+                let (res, upgrade) = client.send_request(req).await?;
 
                 println!("{} -> {}", uri, res.status());
+                if let Some(upgrade) = upgrade {
+                    // If the response is an upgrade, e.g. Websocket, you can see traffic.
+                    // Modifying upgraded traffic is not supported yet.
+
+                    // You can try https://echo.websocket.org/.ws to test websocket.
+                    // But you need to disable alpn of DefaultClient to disable HTTP2 because echo.websocket.org does not support HTTP/2 for Websocket.
+                    // It should be match incoming and outgoing HTTP version on DefaultClient, I'll fix this later. #54
+                    println!("Upgrade connection");
+                    let Upgrade {
+                        mut client_to_server,
+                        mut server_to_client,
+                    } = upgrade;
+                    let url = uri.to_string();
+                    tokio::spawn(async move {
+                        while let Some(data) = client_to_server.next().await {
+                            println!("Client -> Server: {} {:?}", url, data);
+                        }
+                    });
+                    let url = uri.to_string();
+                    tokio::spawn(async move {
+                        while let Some(data) = server_to_client.next().await {
+                            println!("Server -> Client: {} {:?}", url, data);
+                        }
+                    });
+                }
 
                 // You can modify response here
 
diff --git a/src/default_client.rs b/src/default_client.rs
index 6bc86e8..877c31f 100644
--- a/src/default_client.rs
+++ b/src/default_client.rs
@@ -44,6 +44,7 @@ impl DefaultClient {
     }
 
     /// Send a request and return a response.
+    /// If the response is an upgrade (= if status code is 101 Switching Protocols), it will return a response and an Upgrade struct.
     /// Request should have a full URL including scheme.
     pub async fn send_request<B>(
         &self,