From 39c31d965157cee6b1694f744ee3771bb57d699c Mon Sep 17 00:00:00 2001
From: DarkCat09 <gh@dc09.ru>
Date: Tue, 6 Aug 2024 18:33:11 +0400
Subject: [PATCH] important fix for buffered stream + code cleanup

a `TlsStream<TcpStream>` was being wrapped in a `BufReader` during header parsing,
but then the pure stream (without BufReader) was being passed into a Response constructor,
discarding all content in an internal BufReader's buffer,
resulting in an incomplete body (truncated at the beginning)
or empty at all
---
 src/client.rs   | 61 ++++++++++++++++++++++++++-----------------------
 src/response.rs |  2 +-
 2 files changed, 34 insertions(+), 29 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index e9c3953..fead612 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -85,39 +85,44 @@ impl Client {
         stream.write_all(url_str.as_bytes()).await?;
         stream.write_all(b"\r\n").await?;
 
-        let mut buf: [u8; 3] = [0, 0, 0]; // 2 digits, space
-        stream.read_exact(&mut buf).await?;
-        let status = Status::parse_status(&buf)?;
+        let status = {
+            let mut buf: [u8; 3] = [0, 0, 0]; // 2 digits, space
+            stream.read_exact(&mut buf).await?;
+            Status::parse_status(&buf)?
+        };
 
-        let mut message: Vec<u8> = Vec::new();
-        let mut buf_reader = tokio::io::BufReader::new(&mut stream);
-        let mut buf: [u8; 1] = [0]; // buffer for LF (\n)
+        let mut stream = tokio::io::BufReader::new(stream);
 
-        // reading message after status code
-        // until CRLF (\r\n)
-        loop {
-            // until CR
-            buf_reader.read_until(b'\r', &mut message).await?;
-            // now read next char...
-            buf_reader.read_exact(&mut buf).await?;
-            if buf[0] == b'\n' {
-                // ...and check if it's LF
-                break;
-            } else {
-                // ...otherwise, CR is a part of message, not a CRLF terminator,
-                // so append that one byte that's supposed to be LF (but not LF)
-                // to the message buffer
-                message.push(buf[0].into());
+        let message = {
+            let mut result: Vec<u8> = Vec::new();
+            let mut buf = [0u8]; // buffer for LF (\n)
+
+            // reading message after status code
+            // until CRLF (\r\n)
+            loop {
+                // until CR
+                stream.read_until(b'\r', &mut result).await?;
+                // now read next char...
+                stream.read_exact(&mut buf).await?;
+                if buf[0] == b'\n' {
+                    // ...and check if it's LF
+                    break;
+                } else {
+                    // ...otherwise, CR is a part of message, not a CRLF terminator,
+                    // so append that one byte that's supposed to be LF (but not LF)
+                    // to the message buffer
+                    result.push(buf[0].into());
+                }
             }
-        }
 
-        // trim last CR
-        if message.last().is_some_and(|c| c == &b'\r') {
-            message.pop();
-        }
+            // trim last CR
+            if result.last().is_some_and(|c| c == &b'\r') {
+                result.pop();
+            }
 
-        // Vec<u8> -> ASCII or UTF-8 String
-        let message = String::from_utf8(message)?;
+            // Vec<u8> -> ASCII or UTF-8 String
+            String::from_utf8(result)?
+        };
 
         Ok(Response::new(status, message, stream))
     }
diff --git a/src/response.rs b/src/response.rs
index 32eb0db..d0384eb 100644
--- a/src/response.rs
+++ b/src/response.rs
@@ -3,7 +3,7 @@ use crate::{status::Status, LibError, ReplyType};
 use bytes::Bytes;
 use tokio::io::AsyncReadExt;
 
-type BodyStream = tokio_rustls::client::TlsStream<tokio::net::TcpStream>;
+type BodyStream = tokio::io::BufReader<tokio_rustls::client::TlsStream<tokio::net::TcpStream>>;
 
 #[derive(Debug)]
 pub struct Response {