francoisgib_webserver 1.0.3

use std::str::FromStr;

use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncReadExt};

use crate::utils::buffer::Buffer;

use super::{BodyType, HttpRequest};

/// Reads the body of an HTTP request based on its `BodyType`.
///
/// If the body type is `Fixed`, it reads the specified number of bytes.
/// If the body type is `Chunked`, it reads the body using chunked transfer encoding.
/// If the body type is `None`, it returns `None`.
///
/// # Arguments
///
/// * `reader` - A mutable reference to an async buffered reader.
/// * `request` - The parsed `HttpRequest` that contains information about the body type.
///
/// # Returns
///
/// * `Ok(Some(Buffer))` if the body was successfully read.
/// * `Ok(None)` if there is no body.
/// * `Err(String)` if an error occurs during reading.
///
/// # Errors
/// Returns a string error if the reading fails.
pub async fn read_request_body<R>(
    reader: &mut R,
    request: &HttpRequest,
) -> Result<Option<Buffer>, String>
where
    R: AsyncBufRead + Unpin,
{
    match request.body_type {
        Some(BodyType::Fixed(length)) => {
            let mut body = Buffer::new(length as usize);
            body.async_read(reader).await.map_err(|e| e.to_string())?;
            Ok(Some(Buffer::from(body)))
        }
        Some(BodyType::Chunked) => Ok(read_chunked_body(reader).await.map(|body| Some(body))?),
        None => Ok(None),
    }
}

/// Reads an HTTP body using chunked transfer encoding.
///
/// Parses each chunk size, reads the corresponding data, and collects it into a `Buffer`.
/// Handles optional chunk extensions and validates the final 0-size chunk.
///
/// # Arguments
///
/// * `reader` - A mutable reference to an async buffered reader.
///
/// # Returns
///
/// * `Ok(Buffer)` containing the full assembled body if successful.
/// * `Err(String)` if an error occurs while parsing or reading the chunks.
///
/// # Errors
///
/// Returns a string error if the chunk size is invalid,
/// if unexpected EOF is encountered, or if any read operation fails.
pub async fn read_chunked_body<R>(reader: &mut R) -> Result<Buffer, String>
where
    R: AsyncBufRead + AsyncBufReadExt + Unpin,
{
    let mut result = String::new();

    loop {
        let mut size_line = String::new();
        reader
            .read_line(&mut size_line)
            .await
            .map_err(|e| format!("Error reading chunk size: {}", e))?;

        let size_str = size_line.trim_end();
        let semicolon_pos = size_str.find(';').unwrap_or(size_str.len());
        let chunk_size = u64::from_str_radix(&size_str[..semicolon_pos], 16)
            .map_err(|e| format!("Invalid chunk size format: {}", e))?;

        if chunk_size == 0 {
            let mut final_line = String::new();
            reader
                .read_line(&mut final_line)
                .await
                .map_err(|e| format!("Error reading final chunk delimiter: {}", e))?;
            break;
        }

        let mut chunk_buffer = vec![0u8; chunk_size as usize];
        let mut bytes_read = 0;

        while bytes_read < chunk_size as usize {
            let n = reader
                .read(&mut chunk_buffer[bytes_read..])
                .await
                .map_err(|e| format!("Error reading chunk data: {}", e))?;

            if n == 0 {
                return Err("Unexpected EOF while reading chunk data".to_string());
            }

            bytes_read += n;
        }

        let chunk_str = String::from_utf8_lossy(&chunk_buffer[..bytes_read]).to_string();
        result.push_str(&chunk_str);

        let mut chunk_end = String::new();
        reader
            .read_line(&mut chunk_end)
            .await
            .map_err(|e| format!("Error reading chunk delimiter: {}", e))?;
    }

    Buffer::from_str(result.as_str())
}

#[cfg(test)]
mod tests {
    use crate::http::methods::HttpMethod;

    use super::*;
    use smallvec::smallvec;
    use tokio::io::BufReader;

    use super::{BodyType, HttpRequest};

    #[tokio::test]
    async fn test_read_fixed_body() {
        let data = b"Hello world!";
        let mut reader = BufReader::new(&data[..]);

        let mut request = mock_request();
        request.body_type = Some(BodyType::Fixed(data.len() as u64));
        let body = read_request_body(&mut reader, &request)
            .await
            .unwrap()
            .unwrap();

        assert_eq!(body.to_string().as_bytes(), data);
    }

    #[tokio::test]
    async fn test_read_chunked_body() {
        let chunked_data = b"5\r\nHello\r\n7\r\n world!\r\n0\r\n\r\n";
        let mut reader = BufReader::new(&chunked_data[..]);

        let mut request = mock_request();
        request.body =
            Buffer::from_str(String::from_utf8_lossy(chunked_data).to_string().as_str()).ok();
        request.body_type = Some(BodyType::Chunked);

        let body = read_request_body(&mut reader, &request)
            .await
            .unwrap()
            .unwrap();

        assert_eq!(body.to_string(), "Hello world!");
    }

    #[tokio::test]
    async fn test_read_empty_chunked_body() {
        let chunked_data = b"0\r\n\r\n";
        let mut reader = BufReader::new(&chunked_data[..]);

        let mut request = mock_request();
        request.body =
            Buffer::from_str(String::from_utf8_lossy(chunked_data).to_string().as_str()).ok();
        request.body_type = Some(BodyType::Chunked);

        let body = read_request_body(&mut reader, &request)
            .await
            .unwrap()
            .unwrap();

        assert_eq!(body.to_string(), "");
    }

    #[tokio::test]
    async fn test_invalid_chunk_size() {
        let invalid_data = b"Hello\r\nOworld!\r\n";
        let mut reader = BufReader::new(&invalid_data[..]);

        let result = read_chunked_body(&mut reader).await;

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Invalid chunk size format"));
    }

    #[tokio::test]
    async fn test_no_body() {
        let data = b"";
        let mut reader = BufReader::new(&data[..]);

        let mut request = mock_request();
        request.body = None;

        let body = read_request_body(&mut reader, &request).await.unwrap();
        assert!(body.is_none());
    }

    fn mock_request() -> HttpRequest {
        HttpRequest::new(
            HttpMethod::GET,
            "/".into(),
            (1, 1),
            smallvec![],
            Some(Buffer::from_str("Hello world!").unwrap()),
        )
    }
}