Handle chunked http response

Author: Ma233

Cargo.lock

@@ -8,13 +8,15 @@ base64 = "0.22.1"
 bytes = "1.6.0"
 clap = "4.5.4"
 futures = "0.3.30"
+futures-util = "0.3.30"
 hex = "0.4.3"
 http = "1.1.0"
+http-body = "1.0.0"
 http-body-util = "0.1.1"
 httparse = "1.8.0"
 hyper = { version = "1.3.1", features = ["server", "http1"] }
 hyper-util = { version = "0.1.5", features = ["tokio"] }
-litep2p = "0.6.0"
+litep2p = { git = "https://github.com/Ma233/litep2p.git", branch = "keepalive" }
 # Do not upgrade multiaddr, see: https://github.com/paritytech/litep2p/pull/91
 multiaddr = "0.17.1"
 percent-encoding = "2.3.1"

Cargo.toml

@@ -1,11 +1,14 @@
+use std::convert::Infallible;
 use std::net::SocketAddr;
 use std::str::FromStr;
 
 use base64::prelude::*;
 use bytes::Bytes;
+use http_body::Frame;
 use http_body_util::combinators::BoxBody;
 use http_body_util::BodyExt;
 use http_body_util::Full;
+use http_body_util::StreamBody;
 use hyper::server::conn::http1;
 use hyper::service::service_fn;
 use hyper::Request;
@@ -99,6 +102,16 @@ async fn gateway(
     };
 
     let (parts, _) = resp.into_parts();
+
+    // Handle chunked encoding
+    if let Some(true) = parts.headers.get(http::header::TRANSFER_ENCODING).map(|v| {
+        let Ok(v) = v.to_str() else { return false };
+        v.to_lowercase().contains("chunked")
+    }) {
+        let chunks = split_chunked_body(trailer);
+        return Ok(Response::from_parts(parts, stream_body(chunks)));
+    }
+
     let t = trailer.to_vec();
     Ok(Response::from_parts(parts, full(t)))
 }
@@ -109,6 +122,12 @@ fn full<T: Into<Bytes>>(chunk: T) -> BoxBody<Bytes, hyper::Error> {
         .boxed()
 }
 
+fn stream_body(chunks: Vec<Result<Frame<Bytes>, Infallible>>) -> BoxBody<Bytes, hyper::Error> {
+    StreamBody::new(futures_util::stream::iter(chunks))
+        .map_err(|never| match never {})
+        .boxed()
+}
+
 fn error_response(
     msg: &'static str,
     status: http::StatusCode,
@@ -228,3 +247,33 @@ fn parse_response_header_easy(
     let mut h = [httparse::EMPTY_HEADER; 50];
     parse_response_header(buf, h.as_mut())
 }
+
+fn split_chunked_body(input: &[u8]) -> Vec<Result<Frame<Bytes>, Infallible>> {
+    use std::io::BufRead;
+    use std::io::BufReader;
+    use std::io::Cursor;
+    use std::io::Read;
+
+    let mut reader = BufReader::new(Cursor::new(input));
+    let mut chunks = Vec::new();
+
+    loop {
+        let mut chunk_size = String::new();
+        let read = reader.read_line(&mut chunk_size).unwrap();
+        if read == 0 {
+            break;
+        }
+
+        let chunk_size = usize::from_str_radix(chunk_size.trim_end(), 16).unwrap();
+
+        let mut chunk = vec![0; chunk_size];
+        reader.read_exact(&mut chunk).unwrap();
+
+        let mut crlf = [0; 2];
+        reader.read_exact(&mut crlf).unwrap();
+
+        chunks.push(Ok(Frame::data(Bytes::from(chunk))));
+    }
+
+    chunks
+}

src/gateway.rs

@@ -66,6 +66,7 @@ pub struct PProxyHandle {
 pub enum FullLength {
     NotParsed,
     NotSet,
+    Chunked,
     Parsed(usize),
 }
 
@@ -74,8 +75,8 @@ impl FullLength {
         matches!(self, FullLength::NotParsed)
     }
 
-    pub fn not_set(&self) -> bool {
-        matches!(self, FullLength::NotSet)
+    pub fn chunked(&self) -> bool {
+        matches!(self, FullLength::Chunked)
     }
 }
 
@@ -139,7 +140,7 @@ impl PProxy {
                     return Err(Error::IncompleteHttpRequest);
                 }
 
-                let mut stream = tcp_connect_with_timeout(&proxy_addr, 1).await?;
+                let mut stream = tcp_connect_with_timeout(&proxy_addr, 2).await?;
                 stream.write_all(request).await?;
 
                 let mut response = Vec::new();
@@ -149,7 +150,7 @@ impl PProxy {
                     let mut buf = [0u8; 30000];
 
                     let Ok(Ok(n)) =
-                        timeout(std::time::Duration::from_secs(5), stream.read(&mut buf)).await
+                        timeout(std::time::Duration::from_secs(2), stream.read(&mut buf)).await
                     else {
                         break;
                     };
@@ -175,12 +176,25 @@ impl PProxy {
                                 value.parse::<usize>().ok()
                             });
 
-                            match content_length {
-                                Some(content_length) => {
+                            let transfor_encoding = resp_checker.headers.iter().find_map(|h| {
+                                if h.name.to_lowercase() != "transfer-encoding" {
+                                    return None;
+                                }
+                                let Ok(value) = std::str::from_utf8(h.value) else {
+                                    return None;
+                                };
+                                Some(value)
+                            });
+
+                            match (content_length, transfor_encoding) {
+                                (Some(content_length), _) => {
                                     let header_length = res.unwrap();
                                     full_length = FullLength::Parsed(header_length + content_length)
                                 }
-                                None => {
+                                (None, Some(value)) if value.to_lowercase().contains("chunked") => {
+                                    full_length = FullLength::Chunked;
+                                }
+                                _ => {
                                     full_length = FullLength::NotSet;
                                 }
                             }
@@ -192,6 +206,10 @@ impl PProxy {
                             break;
                         }
                     }
+
+                    if full_length.chunked() && response.ends_with(b"0\r\n\r\n") {
+                        break;
+                    }
                 }
 
                 if response.is_empty() {