From c8a96f2616958a25835c57cd1e4fcec3d91fa04e Mon Sep 17 00:00:00 2001
From: MeexReay <meexreay@gmail.com>
Date: Sun, 8 Dec 2024 15:51:43 +0300
Subject: [PATCH] fix chunked data

---
 src/flowgate/server.rs | 59 ++++++++++++++++++++++++++++++++++++------
 1 file changed, 51 insertions(+), 8 deletions(-)

diff --git a/src/flowgate/server.rs b/src/flowgate/server.rs
index e225247..920c414 100644
--- a/src/flowgate/server.rs
+++ b/src/flowgate/server.rs
@@ -425,6 +425,7 @@ impl FlowgateServer {
                 if read >= content_length { break }
             }
         } else if is_chunked {
+            
             loop {
                 let mut length = Vec::new();
                 {
@@ -440,16 +441,17 @@ impl FlowgateServer {
                             (1, b'\n') => break,
                             _ => 0,
                         };
-                        conn.stream.write_all(&buf).ok()?;
                     }
+                    conn.stream.write_all(&length).ok()?;
 
                     length.truncate(length.len() - 2);
                 }
-                let length = usize::from_str_radix(String::from_utf8(length).ok()?.as_str(), 16).ok()?;
-                let mut data = vec![0; length+2];
+                let length = String::from_utf8(length).ok()?;
+                let length = usize::from_str_radix(length.as_str(), 16).ok()?;
+                let mut data = vec![0u8; length+2];
                 stream.read_exact(&mut data).ok()?;
+
                 conn.stream.write_all(&data).ok()?;
-                data.truncate(length);
                 if length == 0 {
                     break;
                 }
@@ -486,16 +488,27 @@ impl FlowgateServer {
             let head_str = String::from_utf8(head.clone()).ok()?;
             let head_str = head_str.trim_matches(char::from(0));
 
-            let content_length = head_str.split("\r\n")
+            let headers = head_str.split("\r\n")
                 .skip(1)
                 .filter(|l| l.contains(": "))
                 .map(|l| l.split_once(": ").unwrap())
-                .filter(|(k, _)| k.to_lowercase() == "content-length")
-                .next()
+                .map(|(k,v)| (k.to_lowercase(),v.to_string()))
+                .collect::<Vec<(String,String)>>();
+
+            let content_length = headers.iter()
+                .find(|(k, _)| k == "content-length")
                 .map(|o| o.1.parse().ok())
                 .flatten()
                 .unwrap_or(0usize);
+            
+            let is_chunked = headers.iter()
+                .find(|o| o.0.to_lowercase() == "transfer-encoding")
+                .map(|o| o.1.split(",").map(|x| x.trim_matches(' ').to_string()).collect::<Vec<String>>())
+                .map(|o| o.contains(&"chunked".to_string()))
+                .unwrap_or(false);
 
+            println!("{content_length} {is_chunked}");
+            
             if content_length > 0 {
                 let mut read = 0usize;
                 let mut buf = vec![0; 4096];
@@ -507,6 +520,37 @@ impl FlowgateServer {
                     buf = vec![0; 4096];
                     if read == content_length { break }
                 }
+            } else if is_chunked {
+                loop {
+                    let mut length = Vec::new();
+                    {
+                        let mut buf = [0; 1];
+                        let mut counter = 0;
+    
+                        while let Ok(1) = conn.stream.read(&mut buf) {
+                            let byte = buf[0];
+                            length.push(byte);
+    
+                            counter = match (counter, byte) {
+                                (0, b'\r') => 1,
+                                (1, b'\n') => break,
+                                _ => 0,
+                            };
+                        }
+                        stream.write_all(&length).ok()?;
+    
+                        length.truncate(length.len() - 2);
+                    }
+                    let length = String::from_utf8(length).ok()?;
+                    let length = usize::from_str_radix(length.as_str(), 16).ok()?;
+                    let mut data = vec![0u8; length+2];
+                    conn.stream.read_exact(&mut data).ok()?;
+
+                    stream.write_all(&data).ok()?;
+                    if length == 0 {
+                        break;
+                    }
+                }
             }
         } else {
             let mut buf = vec![0;1024];
@@ -514,7 +558,6 @@ impl FlowgateServer {
                 if n == 0 { break }
                 buf.truncate(n);
                 stream.write_all(&buf).ok()?;
-                if n < 1024 { break }
                 buf = vec![0;1024];
             }
         }