g3proxy: send small http body along with header

bytedance · Nov 17, 2024 · ce06bfb · ce06bfb
1 parent 35b344e
commit ce06bfb
Show file tree

Hide file tree

Showing 14 changed files with 370 additions and 114 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/g3proxy/src/escape/direct_fixed/http_forward/writer.rs b/g3proxy/src/escape/direct_fixed/http_forward/writer.rs
@@ -104,7 +104,11 @@ where
         }
     }
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
-        send_req_header_to_origin(&mut self.inner, req).await
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()> {
+        send_req_header_to_origin(&mut self.inner, req, body).await
     }
 }
diff --git a/g3proxy/src/escape/direct_float/http_forward/writer.rs b/g3proxy/src/escape/direct_float/http_forward/writer.rs
@@ -109,11 +109,15 @@ where
         }
     }
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()> {
         if self.bind.is_expired() {
             Err(io::Error::other("connection has expired"))
         } else {
-            send_req_header_to_origin(&mut self.inner, req).await
+            send_req_header_to_origin(&mut self.inner, req, body).await
         }
     }
 }
diff --git a/g3proxy/src/escape/proxy_float/peer/http/http_forward/writer.rs b/g3proxy/src/escape/proxy_float/peer/http/http_forward/writer.rs
@@ -114,7 +114,11 @@ where
         }
     }
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()> {
         if let Some(expire) = &self.config.expire_instant {
             let now = Instant::now();
             if expire.checked_duration_since(now).is_none() {
@@ -124,6 +128,7 @@ where
         send_req_header_via_proxy(
             &mut self.inner,
             req,
+            body,
             &self.upstream,
             &self.config.append_http_headers,
             None,
@@ -205,13 +210,17 @@ where
         }
     }
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()> {
         if let Some(expire) = &self.config.expire_instant {
             let now = Instant::now();
             if expire.checked_duration_since(now).is_none() {
                 return Err(io::Error::other("connection has expired"));
             }
         }
-        send_req_header_to_origin(&mut self.inner, req).await
+        send_req_header_to_origin(&mut self.inner, req, body).await
     }
 }
diff --git a/g3proxy/src/escape/proxy_float/peer/https/http_forward/writer.rs b/g3proxy/src/escape/proxy_float/peer/https/http_forward/writer.rs
@@ -105,7 +105,11 @@ where
         self.inner.reset_stats(Arc::new(wrapper_stats));
     }
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()> {
         if let Some(expire) = &self.config.expire_instant {
             let now = Instant::now();
             if expire.checked_duration_since(now).is_none() {
@@ -115,6 +119,7 @@ where
         send_req_header_via_proxy(
             &mut self.inner,
             req,
+            body,
             &self.upstream,
             &self.config.append_http_headers,
             None,
@@ -184,13 +189,17 @@ where
         self.inner.reset_stats(Arc::new(wrapper_stats));
     }
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()> {
         if let Some(expire) = &self.config.expire_instant {
             let now = Instant::now();
             if expire.checked_duration_since(now).is_none() {
                 return Err(io::Error::other("connection has expired"));
             }
         }
-        send_req_header_to_origin(&mut self.inner, req).await
+        send_req_header_to_origin(&mut self.inner, req, body).await
     }
 }
diff --git a/g3proxy/src/escape/proxy_float/peer/socks5/http_forward/writer.rs b/g3proxy/src/escape/proxy_float/peer/socks5/http_forward/writer.rs
@@ -109,13 +109,17 @@ where
         }
     }
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()> {
         if let Some(expire) = &self.config.expire_instant {
             let now = Instant::now();
             if expire.checked_duration_since(now).is_none() {
                 return Err(io::Error::other("connection has expired"));
             }
         }
-        send_req_header_to_origin(&mut self.inner, req).await
+        send_req_header_to_origin(&mut self.inner, req, body).await
     }
 }
diff --git a/g3proxy/src/escape/proxy_float/peer/socks5s/http_forward/writer.rs b/g3proxy/src/escape/proxy_float/peer/socks5s/http_forward/writer.rs
@@ -97,13 +97,17 @@ where
         self.inner.reset_stats(Arc::new(wrapper_stats));
     }
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()> {
         if let Some(expire) = &self.config.expire_instant {
             let now = Instant::now();
             if expire.checked_duration_since(now).is_none() {
                 return Err(io::Error::other("connection has expired"));
             }
         }
-        send_req_header_to_origin(&mut self.inner, req).await
+        send_req_header_to_origin(&mut self.inner, req, body).await
     }
 }
diff --git a/g3proxy/src/escape/proxy_http/http_forward/writer.rs b/g3proxy/src/escape/proxy_http/http_forward/writer.rs
@@ -116,11 +116,16 @@ where
         }
     }
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()> {
         let userid = self.pass_userid.as_deref();
         send_req_header_via_proxy(
             &mut self.inner,
             req,
+            body,
             &self.upstream,
             &self.config.append_http_headers,
             userid,
@@ -202,7 +207,11 @@ where
         }
     }
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
-        send_req_header_to_origin(&mut self.inner, req).await
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()> {
+        send_req_header_to_origin(&mut self.inner, req, body).await
     }
 }
diff --git a/g3proxy/src/escape/proxy_https/http_forward/writer.rs b/g3proxy/src/escape/proxy_https/http_forward/writer.rs
@@ -107,11 +107,16 @@ where
         self.inner.reset_stats(Arc::new(wrapper_stats));
     }
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()> {
         let userid = self.pass_userid.as_deref();
         send_req_header_via_proxy(
             &mut self.inner,
             req,
+            body,
             &self.upstream,
             &self.config.append_http_headers,
             userid,
@@ -181,7 +186,11 @@ where
         self.inner.reset_stats(Arc::new(wrapper_stats));
     }
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
-        send_req_header_to_origin(&mut self.inner, req).await
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()> {
+        send_req_header_to_origin(&mut self.inner, req, body).await
     }
 }
diff --git a/g3proxy/src/module/http_forward/connection/mod.rs b/g3proxy/src/module/http_forward/connection/mod.rs
@@ -51,7 +51,11 @@ pub(crate) trait HttpForwardWrite: AsyncWrite {
         user_stats: Vec<Arc<UserUpstreamTrafficStats>>,
     );
 
-    async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()>;
+    async fn send_request_header(
+        &mut self,
+        req: &HttpProxyClientRequest,
+        body: Option<&[u8]>,
+    ) -> io::Result<()>;
 }
 
 #[async_trait]
@@ -107,6 +111,6 @@ impl AsyncWrite for HttpForwardWriterForAdaptation<'_> {
 
 impl HttpRequestUpstreamWriter<HttpProxyClientRequest> for HttpForwardWriterForAdaptation<'_> {
     async fn send_request_header(&mut self, req: &HttpProxyClientRequest) -> io::Result<()> {
-        self.inner.send_request_header(req).await
+        self.inner.send_request_header(req, None).await
     }
 }
diff --git a/g3proxy/src/module/http_forward/connection/writer.rs b/g3proxy/src/module/http_forward/connection/writer.rs
@@ -14,19 +14,21 @@
  * limitations under the License.
  */
 
-use std::io;
+use std::io::{self, IoSlice};
+
+use g3_io_ext::LimitedWriteExt;
+use g3_types::net::UpstreamAddr;
 
 use bytes::BufMut;
 use tokio::io::{AsyncWrite, AsyncWriteExt};
 
-use g3_types::net::UpstreamAddr;
-
 use super::HttpProxyClientRequest;
 use crate::module::http_header;
 
 pub(crate) async fn send_req_header_via_proxy<W>(
     writer: &mut W,
     req: &HttpProxyClientRequest,
+    body: Option<&[u8]>,
     upstream: &UpstreamAddr,
     append_header_lines: &[String],
     pass_userid: Option<&str>,
@@ -45,16 +47,35 @@ where
     }
     buf.put_slice(b"\r\n");
 
-    writer.write_all(buf.as_ref()).await
+    send_request_header(writer, buf.as_slice(), body).await
 }
 
 pub(crate) async fn send_req_header_to_origin<W>(
     writer: &mut W,
     req: &HttpProxyClientRequest,
+    body: Option<&[u8]>,
 ) -> io::Result<()>
 where
     W: AsyncWrite + Unpin,
 {
     let buf = req.serialize_for_origin();
-    writer.write_all(buf.as_ref()).await
+    send_request_header(writer, buf.as_slice(), body).await
+}
+
+async fn send_request_header<W>(
+    writer: &mut W,
+    header: &[u8],
+    body: Option<&[u8]>,
+) -> io::Result<()>
+where
+    W: AsyncWrite + Unpin,
+{
+    if let Some(body) = body {
+        writer
+            .write_all_vectored([IoSlice::new(header), IoSlice::new(body)])
+            .await?;
+        Ok(())
+    } else {
+        writer.write_all(header).await
+    }
 }
diff --git a/g3proxy/src/module/http_forward/mod.rs b/g3proxy/src/module/http_forward/mod.rs
@@ -22,8 +22,8 @@ mod task;
 
 pub(crate) use connection::{
     send_req_header_to_origin, send_req_header_via_proxy, BoxHttpForwardConnection,
-    BoxHttpForwardReader, BoxHttpForwardWriter, HttpConnectionEofPoller, HttpForwardRead,
-    HttpForwardWrite, HttpForwardWriterForAdaptation,
+    BoxHttpForwardReader, HttpConnectionEofPoller, HttpForwardRead, HttpForwardWrite,
+    HttpForwardWriterForAdaptation,
 };
 pub(crate) use context::{
     BoxHttpForwardContext, DirectHttpForwardContext, FailoverHttpForwardContext,