pingap-core 0.16.0

// Copyright 2024-2025 Tree xie.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use super::{
    HTTP_HEADER_CONTENT_HTML, HTTP_HEADER_CONTENT_JSON,
    HTTP_HEADER_CONTENT_TEXT, HTTP_HEADER_NO_CACHE, HTTP_HEADER_NO_STORE,
    HTTP_HEADER_TRANSFER_CHUNKED, HttpHeader, LOG_TARGET, get_super_ts,
    new_internal_error,
};
use bytes::{Bytes, BytesMut};
use http::StatusCode;
use http::header;
use http::{HeaderName, HeaderValue};
use pingora::http::ResponseHeader;
use pingora::proxy::Session;
use serde::Serialize;
use std::pin::Pin;
use tokio::io::AsyncReadExt;
use tracing::error;

/// A helper function to generate a `Cache-Control` header.
///
/// It determines the appropriate `Cache-Control` value based on the provided
/// `max_age` and `cache_private` settings. This is a performance-sensitive
/// function, so it's optimized to avoid allocations where possible.
///
/// # Arguments
/// * `max_age`: An `Option<u32>` specifying the max-age in seconds. If `Some(0)` or `None`,
///   a "no-cache" header is returned.
/// * `cache_private`: An `Option<bool>` indicating if the cache is private. Defaults to public.
///
/// # Returns
/// An `HttpHeader` tuple `(HeaderName, HeaderValue)` for the `Cache-Control` header.
fn new_cache_control_header(
    max_age: Option<u32>,
    cache_private: Option<bool>,
) -> HttpHeader {
    // Determine the max_age value. If it's 0 or not provided, return the static no-cache header immediately.
    let max_age = match max_age {
        Some(0) | None => return HTTP_HEADER_NO_CACHE.clone(),
        Some(age) => age,
    };

    // Determine the cache visibility ("private" or "public").
    let category: &[u8] = if cache_private.unwrap_or_default() {
        b"private"
    } else {
        b"public"
    };

    // Use a pre-allocated buffer (`BytesMut`) and the `itoa` crate to efficiently build the
    // header value without creating an intermediate `String` via `format!`.
    // The capacity is estimated to prevent reallocations.
    let mut buf = BytesMut::with_capacity(category.len() + 9 + 10); // e.g., "public, max-age=" + up to 10 digits for a u32
    buf.extend_from_slice(category);
    buf.extend_from_slice(b", max-age=");
    buf.extend_from_slice(itoa::Buffer::new().format(max_age).as_bytes());

    // Try to create a `HeaderValue` from the constructed bytes.
    if let Ok(value) = HeaderValue::from_bytes(&buf) {
        return (header::CACHE_CONTROL, value);
    }
    // If creation fails for any reason, fall back to the safe "no-cache" header.
    HTTP_HEADER_NO_CACHE.clone()
}

/// A builder for creating `HttpResponse` instances fluently using the builder pattern.
///
/// This allows for chaining method calls to configure a response, improving readability.
#[derive(Default, Debug)]
pub struct HttpResponseBuilder {
    /// The `HttpResponse` instance being built.
    response: HttpResponse,
}

impl HttpResponseBuilder {
    /// Creates a new builder with a given status code.
    pub fn new(status: StatusCode) -> Self {
        Self {
            response: HttpResponse {
                status,
                ..Default::default()
            },
        }
    }

    /// Sets the response body.
    ///
    /// This method is generic over `impl Into<Bytes>`, allowing various types like
    /// `Vec<u8>`, `String`, or `&'static str` to be passed as the body.
    pub fn body(mut self, body: impl Into<Bytes>) -> Self {
        self.response.body = body.into();
        self
    }

    /// Adds a single HTTP header to the response.
    ///
    /// If the headers vector doesn't exist yet, it will be created.
    pub fn header(mut self, header: HttpHeader) -> Self {
        self.response
            .headers
            .get_or_insert_with(Vec::new)
            .push(header);
        self
    }

    /// Appends multiple HTTP headers to the response.
    pub fn headers(mut self, headers: Vec<HttpHeader>) -> Self {
        self.response
            .headers
            .get_or_insert_with(Vec::new)
            .extend(headers);
        self
    }

    /// Sets the `Cache-Control` max-age and privacy directive.
    pub fn max_age(mut self, seconds: u32, is_private: bool) -> Self {
        self.response.max_age = Some(seconds);
        self.response.cache_private = Some(is_private);
        self
    }

    /// A convenience method to add a "no-store" `Cache-Control` header.
    pub fn no_store(self) -> Self {
        self.header(HTTP_HEADER_NO_STORE.clone())
    }

    /// Consumes the builder and returns the final `HttpResponse` instance.
    pub fn finish(self) -> HttpResponse {
        self.response
    }
}

/// Represents a complete HTTP response, including status, headers, and a body.
/// This struct is used for responses where the entire body is known upfront.
#[derive(Default, Clone, Debug)]
pub struct HttpResponse {
    /// The HTTP status code (e.g., 200 OK, 404 Not Found).
    pub status: StatusCode,
    /// The response body, stored as a `Bytes` object for efficiency.
    pub body: Bytes,
    /// The `max-age` directive for the `Cache-Control` header, in seconds.
    pub max_age: Option<u32>,
    /// The UNIX timestamp when the response content was created, used for the `Age` header.
    pub created_at: Option<u32>,
    /// A flag to indicate if the `Cache-Control` should be "private" (true) or "public" (false).
    pub cache_private: Option<bool>,
    /// A list of additional HTTP headers to include in the response.
    pub headers: Option<Vec<HttpHeader>>,
}

/// The failure to write a response as what it is: the client's side of
/// the connection.
///
/// pingora names the source of an error on the path of a proxied response
/// and leaves it unset on these writes, which a plugin or the proxy itself
/// answers with. Unset, a client that went away while a file was sent to
/// it counted as a failure of the proxy: a `500` in the access log and
/// the metrics, and a line at error level for every download that was
/// cancelled, where the same on a proxied response is a `499`.
fn to_client<T>(result: pingora::Result<T>) -> pingora::Result<T> {
    result.map_err(|e| e.into_down())
}

impl HttpResponse {
    /// Returns a new `HttpResponseBuilder` to start building a response.
    pub fn builder(status: StatusCode) -> HttpResponseBuilder {
        HttpResponseBuilder::new(status)
    }

    /// A convenience constructor for a 204 No Content response.
    pub fn no_content() -> Self {
        Self::builder(StatusCode::NO_CONTENT).no_store().finish()
    }

    /// A convenience constructor for a 400 Bad Request response.
    pub fn bad_request(body: impl Into<Bytes>) -> Self {
        Self::builder(StatusCode::BAD_REQUEST)
            .body(body)
            .header(HTTP_HEADER_CONTENT_TEXT.clone())
            .no_store()
            .finish()
    }

    /// A convenience constructor for a 404 Not Found response.
    pub fn not_found(body: impl Into<Bytes>) -> Self {
        Self::builder(StatusCode::NOT_FOUND)
            .body(body)
            .header(HTTP_HEADER_CONTENT_TEXT.clone())
            .no_store()
            .finish()
    }

    /// A convenience constructor for a 500 Internal Server Error response.
    pub fn unknown_error(body: impl Into<Bytes>) -> Self {
        Self::builder(StatusCode::INTERNAL_SERVER_ERROR)
            .body(body)
            .header(HTTP_HEADER_CONTENT_TEXT.clone())
            .no_store()
            .finish()
    }

    /// A convenience constructor for a 200 OK HTML response.
    pub fn html(body: impl Into<Bytes>) -> Self {
        Self::builder(StatusCode::OK)
            .body(body)
            .header(HTTP_HEADER_CONTENT_HTML.clone())
            .header(HTTP_HEADER_NO_CACHE.clone())
            .finish()
    }

    /// A convenience constructor for a 302 Temporary Redirect response.
    pub fn redirect(location: &str) -> pingora::Result<Self> {
        // Attempt to parse the location string into a valid HeaderValue.
        let value = HeaderValue::from_str(location).map_err(|e| {
            error!(error = e.to_string(), "to header value fail");
            new_internal_error(500, e)
        })?;
        // Build the redirect response.
        Ok(Self::builder(StatusCode::TEMPORARY_REDIRECT)
            .header((header::LOCATION, value))
            .header(HTTP_HEADER_NO_CACHE.clone())
            .finish())
    }

    /// A convenience constructor for a 200 OK plain text response.
    pub fn text(body: impl Into<Bytes>) -> Self {
        Self::builder(StatusCode::OK)
            .body(body)
            .header(HTTP_HEADER_CONTENT_TEXT.clone())
            .header(HTTP_HEADER_NO_CACHE.clone())
            .finish()
    }

    /// Creates a 200 OK JSON response by serializing the given value.
    pub fn try_from_json<T>(value: &T) -> pingora::Result<Self>
    where
        T: ?Sized + Serialize,
    {
        // Serialize the value to a JSON byte vector.
        let buf = serde_json::to_vec(value).map_err(|e| {
            error!(target: LOG_TARGET, error = e.to_string(), "to json fail");
            new_internal_error(400, e)
        })?;
        // Build the JSON response.
        Ok(Self::builder(StatusCode::OK)
            .body(buf)
            .header(HTTP_HEADER_CONTENT_JSON.clone())
            .finish())
    }

    /// Creates a JSON response with a specified status code by serializing the given value.
    pub fn try_from_json_status<T>(
        value: &T,
        status: StatusCode,
    ) -> pingora::Result<Self>
    where
        T: ?Sized + Serialize,
    {
        // First, create a standard 200 OK JSON response.
        let mut resp = Self::try_from_json(value)?;
        // Then, simply overwrite the status code.
        resp.status = status;
        Ok(resp)
    }

    /// Builds a `pingora::http::ResponseHeader` from the `HttpResponse`'s properties.
    pub fn new_response_header(&self) -> pingora::Result<ResponseHeader> {
        // Build the response header with the status code, sized for the
        // headers added below: Content-Length, Cache-Control, Age, plus
        // the custom ones.
        let custom = self.headers.as_ref().map_or(0, |headers| headers.len());
        let mut resp = ResponseHeader::build(self.status, Some(3 + custom))?;

        // A local helper closure to simplify adding headers and handling potential errors.
        let mut add_header =
            |name: &HeaderName, value: &HeaderValue| -> pingora::Result<()> {
                resp.insert_header(name, value)?;
                Ok(())
            };

        // Add the Content-Length header based on the body size. Not on a
        // response that has no body by definition: a 204 must not carry
        // one (RFC 9110 8.6), and on a 304 a zero would be the length of
        // the representation, which it is not.
        if !has_no_body(self.status) {
            add_header(
                &header::CONTENT_LENGTH,
                &HeaderValue::from(self.body.len()),
            )?;
        }

        // Generate and add the Cache-Control header.
        let (name, value) =
            new_cache_control_header(self.max_age, self.cache_private);
        add_header(&name, &value)?;

        // If a creation timestamp is provided, calculate and add the Age header.
        if let Some(created_at) = self.created_at {
            let secs = get_super_ts().saturating_sub(created_at);
            add_header(&header::AGE, &HeaderValue::from(secs))?;
        }

        // Add all custom headers from the `headers` field.
        if let Some(headers) = &self.headers {
            put_custom_headers(&mut resp, headers)?;
        }
        Ok(resp)
    }

    /// Sends the entire HTTP response (header and body) to the client via the session.
    pub async fn send(self, session: &mut Session) -> pingora::Result<usize> {
        // First, build the response header.
        let header = self.new_response_header()?;
        self.send_with_header(session, header).await
    }

    /// Sends the body under `header`: the one `new_response_header` builds,
    /// with whatever the caller has added to it.
    pub async fn send_with_header(
        self,
        session: &mut Session,
        header: ResponseHeader,
    ) -> pingora::Result<usize> {
        // A HEAD, a 204 or a 304 ends with its header. HTTP/1.1 drops a
        // body written after it, HTTP/2 does not: the DATA frame went out
        // and the client reset the stream with a protocol error.
        if session.req_header().method == http::Method::HEAD
            || has_no_body(self.status)
        {
            to_client(
                session.write_response_header(Box::new(header), true).await,
            )?;
            to_client(session.finish_body().await)?;
            return Ok(0);
        }
        let size = self.body.len();
        // Write the header to the session.
        to_client(
            session.write_response_header(Box::new(header), false).await,
        )?;
        // Write the body to the session. `end_stream` is true because this is the final body part.
        to_client(session.write_response_body(Some(self.body), true).await)?;
        // Finalize the response stream.
        to_client(session.finish_body().await)?;
        Ok(size)
    }
}

/// The statuses whose response has no body, whatever the request.
#[inline]
fn has_no_body(status: StatusCode) -> bool {
    status.is_informational()
        || status == StatusCode::NO_CONTENT
        || status == StatusCode::NOT_MODIFIED
}

/// The headers that are sent once per value and cannot be folded into one
/// line: a cookie's `Expires` has a comma of its own, and a challenge is a
/// list already.
fn is_repeated_header(name: &HeaderName) -> bool {
    [
        header::SET_COOKIE,
        header::WWW_AUTHENTICATE,
        header::PROXY_AUTHENTICATE,
        header::LINK,
    ]
    .contains(name)
}

/// Puts the custom headers of a response into `resp`.
///
/// A header given again replaces the earlier one, which is how a configured
/// `Content-Type` takes the place of the detected one. Not so for the few
/// headers that come once per value: each `Set-Cookie` used to replace the
/// last, and of the three an auth service answered with, the client got
/// one.
fn put_custom_headers(
    resp: &mut ResponseHeader,
    headers: &[HttpHeader],
) -> pingora::Result<()> {
    for (index, (name, value)) in headers.iter().enumerate() {
        if is_repeated_header(name)
            && headers[..index].iter().any(|(earlier, _)| earlier == name)
        {
            resp.append_header(name, value)?;
        } else {
            resp.insert_header(name, value)?;
        }
    }
    Ok(())
}

/// Represents a chunked HTTP response for streaming large bodies of data.
///
/// This is used when the response body is too large to fit in memory or is generated on-the-fly.
pub struct HttpChunkResponse<'r, R> {
    /// A pinned, mutable reference to an async reader that provides the body data.
    pub reader: Pin<&'r mut R>,
    /// The status of the response. Defaults to `200 OK`.
    pub status: StatusCode,
    /// The suggested size for each data chunk. Defaults to `DEFAULT_BUF_SIZE`.
    pub chunk_size: usize,
    /// Cache control `max-age` setting for the response.
    pub max_age: Option<u32>,
    /// Cache control privacy setting for the response.
    pub cache_private: Option<bool>,
    /// Additional headers to include in the response.
    pub headers: Option<Vec<HttpHeader>>,
    /// How fast the body is sent, where that is limited.
    pub pace: Option<crate::BodyPace>,
}

/// The default buffer size (8KB) for chunked responses.
const DEFAULT_BUF_SIZE: usize = 8 * 1024;

impl<'r, R> HttpChunkResponse<'r, R>
where
    // The reader must implement `AsyncRead` and be `Unpin`.
    R: tokio::io::AsyncRead + std::marker::Unpin,
{
    /// Creates a new `HttpChunkResponse` with a given reader and default settings.
    pub fn new(r: &'r mut R) -> Self {
        Self {
            reader: Pin::new(r),
            status: StatusCode::OK,
            chunk_size: DEFAULT_BUF_SIZE,
            max_age: None,
            headers: None,
            cache_private: None,
            pace: None,
        }
    }

    /// Builds the `ResponseHeader` for the streamed response.
    ///
    /// The body is framed by `Content-Length` when the custom headers give
    /// one, and by `Transfer-Encoding: chunked` otherwise. Both at once is
    /// not a valid message, and that is what a streamed file used to get:
    /// its length was known and sent, and `chunked` was added regardless.
    pub fn get_response_header(&self) -> pingora::Result<ResponseHeader> {
        let mut resp = ResponseHeader::build(self.status, Some(4))?;
        // Add any custom headers.
        if let Some(headers) = &self.headers {
            put_custom_headers(&mut resp, headers)?;
        }

        if !resp.headers.contains_key(header::CONTENT_LENGTH) {
            let chunked = HTTP_HEADER_TRANSFER_CHUNKED.clone();
            resp.insert_header(chunked.0, chunked.1)?;
        }

        // Add the `Cache-Control` header.
        let cache_control =
            new_cache_control_header(self.max_age, self.cache_private);
        resp.insert_header(cache_control.0, cache_control.1)?;
        Ok(resp)
    }

    /// Reads the next chunk of the body. `None` only at the end of the
    /// stream: a read that comes back short is a chunk like any other, not a
    /// sign that the reader is done - `AsyncRead` never promised to fill the
    /// buffer.
    ///
    /// The chunk is split off `buffer` rather than copied out of it; once
    /// the session has written it and dropped it, `reserve` gets the
    /// allocation back for the next read.
    async fn next_chunk(
        &mut self,
        buffer: &mut BytesMut,
    ) -> pingora::Result<Option<Bytes>> {
        buffer.reserve(self.chunk_size.max(512));
        let size = self.reader.read_buf(buffer).await.map_err(|e| {
            error!(error = e.to_string(), "read data fail");
            new_internal_error(400, e)
        })?;
        if size == 0 {
            return Ok(None);
        }
        Ok(Some(buffer.split().freeze()))
    }

    /// Streams the data from the reader to the client as a chunked response.
    ///
    /// Reads data from the `reader` in chunks and sends each chunk to the client until
    /// the reader is exhausted.
    pub async fn send(self, session: &mut Session) -> pingora::Result<usize> {
        let header = self.get_response_header()?;
        self.send_with_header(session, header).await
    }

    /// Sends the body under `header`: the one `get_response_header` builds,
    /// with whatever the caller has added to it.
    pub async fn send_with_header(
        mut self,
        session: &mut Session,
        header: ResponseHeader,
    ) -> pingora::Result<usize> {
        // `end_stream` is false because a body will follow.
        to_client(
            session.write_response_header(Box::new(header), false).await,
        )?;

        let mut sent = 0;
        let mut buffer = BytesMut::with_capacity(self.chunk_size.max(512));
        // One chunk of lookahead, so the last one goes out flagged as the end
        // of the stream.
        let mut pending = self.next_chunk(&mut buffer).await?;
        while let Some(chunk) = pending {
            let len = chunk.len();
            sent += len;
            pending = self.next_chunk(&mut buffer).await?;
            to_client(
                session
                    .write_response_body(Some(chunk), pending.is_none())
                    .await,
            )?;
            // Not after the last chunk: there is nothing left to hold
            // back.
            if pending.is_some()
                && let Some(delay) =
                    self.pace.as_mut().and_then(|pace| pace.delay(len))
            {
                tokio::time::sleep(delay).await;
            }
        }
        // Finalize the response stream.
        to_client(session.finish_body().await)?;

        Ok(sent)
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    use crate::convert_headers;
    use bytes::Bytes;
    use http::StatusCode;
    use pretty_assertions::assert_eq;
    use serde::Serialize;
    use std::io::Write;
    use tempfile::NamedTempFile;
    use tokio::fs;
    #[test]
    fn test_new_cache_control_header() {
        assert_eq!(
            r###"("cache-control", "private, max-age=3600")"###,
            format!("{:?}", new_cache_control_header(Some(3600), Some(true)))
        );
        assert_eq!(
            r###"("cache-control", "public, max-age=3600")"###,
            format!("{:?}", new_cache_control_header(Some(3600), None))
        );
        assert_eq!(
            r###"("cache-control", "private, no-cache")"###,
            format!("{:?}", new_cache_control_header(Some(0), Some(false)))
        );
        assert_eq!(
            r###"("cache-control", "private, no-cache")"###,
            format!("{:?}", new_cache_control_header(None, None))
        );
    }

    #[test]
    fn test_http_response() {
        assert_eq!(
            r###"HttpResponse { status: 204, body: b"", max_age: None, created_at: None, cache_private: None, headers: Some([("cache-control", "private, no-store")]) }"###,
            format!("{:?}", HttpResponse::no_content())
        );
        assert_eq!(
            r###"HttpResponse { status: 404, body: b"Not Found", max_age: None, created_at: None, cache_private: None, headers: Some([("content-type", "text/plain; charset=utf-8"), ("cache-control", "private, no-store")]) }"###,
            format!("{:?}", HttpResponse::not_found("Not Found"))
        );
        assert_eq!(
            r###"HttpResponse { status: 500, body: b"Unknown Error", max_age: None, created_at: None, cache_private: None, headers: Some([("content-type", "text/plain; charset=utf-8"), ("cache-control", "private, no-store")]) }"###,
            format!("{:?}", HttpResponse::unknown_error("Unknown Error"))
        );

        assert_eq!(
            r###"HttpResponse { status: 400, body: b"Bad Request", max_age: None, created_at: None, cache_private: None, headers: Some([("content-type", "text/plain; charset=utf-8"), ("cache-control", "private, no-store")]) }"###,
            format!("{:?}", HttpResponse::bad_request("Bad Request"))
        );

        assert_eq!(
            r###"HttpResponse { status: 200, body: b"<p>Pingap</p>", max_age: None, created_at: None, cache_private: None, headers: Some([("content-type", "text/html; charset=utf-8"), ("cache-control", "private, no-cache")]) }"###,
            format!("{:?}", HttpResponse::html("<p>Pingap</p>"))
        );

        assert_eq!(
            r###"HttpResponse { status: 307, body: b"", max_age: None, created_at: None, cache_private: None, headers: Some([("location", "http://example.com/"), ("cache-control", "private, no-cache")]) }"###,
            format!(
                "{:?}",
                HttpResponse::redirect("http://example.com/").unwrap()
            )
        );

        assert_eq!(
            r###"HttpResponse { status: 200, body: b"Hello World!", max_age: None, created_at: None, cache_private: None, headers: Some([("content-type", "text/plain; charset=utf-8"), ("cache-control", "private, no-cache")]) }"###,
            format!("{:?}", HttpResponse::text("Hello World!"))
        );

        #[derive(Serialize)]
        struct Data {
            message: String,
        }
        let resp = HttpResponse::try_from_json_status(
            &Data {
                message: "Hello World!".to_string(),
            },
            StatusCode::BAD_REQUEST,
        )
        .unwrap();
        assert_eq!(
            r###"HttpResponse { status: 400, body: b"{\"message\":\"Hello World!\"}", max_age: None, created_at: None, cache_private: None, headers: Some([("content-type", "application/json; charset=utf-8")]) }"###,
            format!("{resp:?}")
        );
        let resp = HttpResponse::try_from_json(&Data {
            message: "Hello World!".to_string(),
        })
        .unwrap();
        assert_eq!(
            r###"HttpResponse { status: 200, body: b"{\"message\":\"Hello World!\"}", max_age: None, created_at: None, cache_private: None, headers: Some([("content-type", "application/json; charset=utf-8")]) }"###,
            format!("{resp:?}")
        );

        let resp = HttpResponse {
            status: StatusCode::OK,
            body: Bytes::from("Hello world!"),
            max_age: Some(3600),
            created_at: Some(0),
            cache_private: Some(true),
            headers: Some(
                convert_headers(&[
                    "Contont-Type: application/json".to_string(),
                    "Content-Encoding: gzip".to_string(),
                ])
                .unwrap(),
            ),
        };

        let mut header = resp.new_response_header().unwrap();
        assert_eq!(true, !header.headers.get("Age").unwrap().is_empty());
        header.remove_header("Age").unwrap();

        assert_eq!(
            r###"ResponseHeader { base: Parts { status: 200, version: HTTP/1.1, headers: {"content-length": "12", "cache-control": "private, max-age=3600", "content-encoding": "gzip", "contont-type": "application/json"} }, header_name_map: Some({"content-length": CaseHeaderName(b"Content-Length"), "cache-control": CaseHeaderName(b"Cache-Control"), "content-encoding": CaseHeaderName(b"Content-Encoding"), "contont-type": CaseHeaderName(b"contont-type")}), reason_phrase: None }"###,
            format!("{header:?}")
        );
    }

    #[tokio::test]
    async fn test_http_chunk_response() {
        let file = include_bytes!("http_response.rs");
        let mut f = NamedTempFile::new().unwrap();
        f.write_all(file).unwrap();
        let mut f = fs::OpenOptions::new().read(true).open(f).await.unwrap();
        let mut resp = HttpChunkResponse::new(&mut f);
        resp.max_age = Some(3600);
        resp.cache_private = Some(false);
        resp.headers = Some(
            convert_headers(&["Contont-Type: text/html".to_string()]).unwrap(),
        );
        let header = resp.get_response_header().unwrap();
        assert_eq!(
            r###"ResponseHeader { base: Parts { status: 200, version: HTTP/1.1, headers: {"contont-type": "text/html", "transfer-encoding": "chunked", "cache-control": "public, max-age=3600"} }, header_name_map: Some({"contont-type": CaseHeaderName(b"contont-type"), "transfer-encoding": CaseHeaderName(b"Transfer-Encoding"), "cache-control": CaseHeaderName(b"Cache-Control")}), reason_phrase: None }"###,
            format!("{header:?}")
        );
    }

    /// A known length frames the body by itself, and the status is the one
    /// that was set.
    #[tokio::test]
    async fn test_http_chunk_response_with_content_length() {
        let mut reader = tokio_test::io::Builder::new().build();
        let mut resp = HttpChunkResponse::new(&mut reader);
        resp.status = StatusCode::PARTIAL_CONTENT;
        resp.headers = Some(
            convert_headers(&[
                "Content-Length: 10".to_string(),
                "Content-Range: bytes 0-9/100".to_string(),
            ])
            .unwrap(),
        );
        let header = resp.get_response_header().unwrap();
        assert_eq!(StatusCode::PARTIAL_CONTENT, header.status);
        assert_eq!("10", header.headers.get("content-length").unwrap());
        assert_eq!(false, header.headers.contains_key("transfer-encoding"));
    }

    /// A chunked response keeps the pace it is given: a wait after each
    /// chunk for what has gone out, and none after the last, where there
    /// is nothing left to hold back.
    #[tokio::test]
    async fn test_http_chunk_response_keeps_the_pace() {
        use tokio::io::AsyncWriteExt;
        // The waits are told on the clock of the runtime, which stands
        // still but for them.
        tokio::time::pause();
        let (mut client, server) = tokio::io::duplex(64 * 1024);
        client.write_all(b"GET / HTTP/1.1\r\n\r\n").await.unwrap();
        let mut session = Session::new_h1(Box::new(server));
        session.read_request().await.unwrap();

        let send = async |session: &mut Session, pace| {
            let mut reader = std::io::Cursor::new(vec![b'a'; 3000]);
            let mut resp = HttpChunkResponse::new(&mut reader);
            resp.chunk_size = 1000;
            resp.pace = pace;
            let start = tokio::time::Instant::now();
            let sent = resp.send(session).await.unwrap();
            (sent, start.elapsed().as_millis())
        };
        // Ten thousand bytes a second, in chunks of a thousand: a tenth
        // of a second after the first, two tenths in all after the
        // second, and the third ends it. A wait after the last as well
        // would make it six tenths.
        let (sent, waited) =
            send(&mut session, Some(crate::BodyPace::new(10_000, 0))).await;
        assert_eq!(3000, sent);
        // (A timer of the runtime ends on a whole millisecond, hence the
        // little that is allowed above three tenths.)
        assert_eq!(true, (200..=350).contains(&waited), "{waited}");
    }

    #[tokio::test]
    async fn test_http_chunk_response_short_reads() {
        // A reader that hands out less than a chunk at a time. Only an empty
        // read ends the body; a short one used to, truncating the file.
        let mut reader = tokio_test::io::Builder::new()
            .read(b"abc")
            .read(b"defgh")
            .build();
        let mut resp = HttpChunkResponse::new(&mut reader);
        resp.chunk_size = 1024;
        let mut buffer = BytesMut::new();
        assert_eq!(
            Some(Bytes::from_static(b"abc")),
            resp.next_chunk(&mut buffer).await.unwrap()
        );
        assert_eq!(
            Some(Bytes::from_static(b"defgh")),
            resp.next_chunk(&mut buffer).await.unwrap()
        );
        assert_eq!(None, resp.next_chunk(&mut buffer).await.unwrap());
    }

    /// Regression: a 204 and a 304 carried `Content-Length: 0`. A 204 must
    /// not have the header at all, and on a 304 it would be the length of
    /// the representation, which is not zero.
    #[test]
    fn test_no_content_length_without_a_body() {
        for (status, expected) in [
            (StatusCode::NO_CONTENT, None),
            (StatusCode::NOT_MODIFIED, None),
            (StatusCode::SWITCHING_PROTOCOLS, None),
            (StatusCode::OK, Some("0")),
            (StatusCode::NOT_FOUND, Some("0")),
        ] {
            let header = HttpResponse {
                status,
                ..Default::default()
            }
            .new_response_header()
            .unwrap();
            assert_eq!(
                expected,
                header
                    .headers
                    .get(header::CONTENT_LENGTH)
                    .and_then(|value| value.to_str().ok()),
                "{status}"
            );
        }
        // The length a HEAD is answered with is the caller's, and stays.
        let header = HttpResponse {
            status: StatusCode::NOT_MODIFIED,
            headers: Some(vec![(
                header::CONTENT_LENGTH,
                HeaderValue::from_static("12"),
            )]),
            ..Default::default()
        }
        .new_response_header()
        .unwrap();
        assert_eq!("12", header.headers.get(header::CONTENT_LENGTH).unwrap());
    }

    /// A response that cannot be written is the client's side failing,
    /// and says so. Regression: pingora leaves the source of such an error
    /// unset here, and unset it counted as a failure of the proxy - a
    /// `500` and a line at error level for every download a client
    /// cancelled, where the same on a proxied response is a `499`.
    #[tokio::test]
    async fn test_a_failed_write_is_the_clients() {
        let gone = async || {
            let (mut client, server) = tokio::io::duplex(1024);
            tokio::io::AsyncWriteExt::write_all(
                &mut client,
                b"GET / HTTP/1.1\r\n\r\n",
            )
            .await
            .unwrap();
            let mut session = Session::new_h1(Box::new(server));
            session.read_request().await.unwrap();
            // the client goes away with its request unanswered
            drop(client);
            session
        };
        let body = vec![b'x'; 256 * 1024];

        let mut session = gone().await;
        let error = HttpResponse {
            status: StatusCode::OK,
            body: Bytes::from(body.clone()),
            ..Default::default()
        }
        .send(&mut session)
        .await
        .expect_err("nobody is there to be written to");
        assert_eq!(&pingora::ErrorSource::Downstream, error.esource());

        let mut session = gone().await;
        let mut reader = std::io::Cursor::new(body);
        let error = HttpChunkResponse::new(&mut reader)
            .send(&mut session)
            .await
            .expect_err("nobody is there to be written to");
        assert_eq!(&pingora::ErrorSource::Downstream, error.esource());
    }

    /// Regression: the body of an answer to a HEAD, or of a 204 or a 304,
    /// was written after the header. HTTP/1.1 drops it; over HTTP/2 the
    /// DATA frame went out and the client reset the stream with a protocol
    /// error. Such a response now ends with its header.
    ///
    /// Over HTTP/1.1, which is what this runs on, the wire looks the same
    /// before and after; what tells them apart here is the number of body
    /// bytes `send` reports as written. The HTTP/2 case is only checked
    /// end to end (`curl --http2-prior-knowledge -I`).
    #[tokio::test]
    async fn test_send_writes_no_body_where_there_is_none() {
        let send = async |request: &str, status: StatusCode| {
            let (mut client, server) = tokio::io::duplex(64 * 1024);
            tokio::io::AsyncWriteExt::write_all(
                &mut client,
                request.as_bytes(),
            )
            .await
            .unwrap();
            let mut session = Session::new_h1(Box::new(server));
            session.read_request().await.unwrap();
            let sent = HttpResponse {
                status,
                body: Bytes::from_static(b"Hello world!"),
                ..Default::default()
            }
            .send(&mut session)
            .await
            .unwrap();
            drop(session);
            let mut buf = vec![];
            tokio::io::AsyncReadExt::read_to_end(&mut client, &mut buf)
                .await
                .unwrap();
            let text = String::from_utf8_lossy(&buf).into_owned();
            let (_, body) = text.split_once("\r\n\r\n").unwrap();
            (sent, body.to_string())
        };
        assert_eq!(
            (12, "Hello world!".to_string()),
            send("GET / HTTP/1.1\r\n\r\n", StatusCode::OK).await
        );
        assert_eq!(
            (0, String::new()),
            send("HEAD / HTTP/1.1\r\n\r\n", StatusCode::OK).await
        );
        assert_eq!(
            (0, String::new()),
            send("GET / HTTP/1.1\r\n\r\n", StatusCode::NO_CONTENT).await
        );
        assert_eq!(
            (0, String::new()),
            send("GET / HTTP/1.1\r\n\r\n", StatusCode::NOT_MODIFIED).await
        );
    }

    #[test]
    fn test_new_cache_control_header_logic() {
        // Test max_age > 0 case
        let (name, value) = new_cache_control_header(Some(3600), Some(true));
        assert_eq!(name, header::CACHE_CONTROL);
        assert_eq!(value.to_str().unwrap(), "private, max-age=3600");

        let (name, value) = new_cache_control_header(Some(3600), Some(false));
        assert_eq!(name, header::CACHE_CONTROL);
        assert_eq!(value.to_str().unwrap(), "public, max-age=3600");

        // Test None is equivalent to public
        let (name, value) = new_cache_control_header(Some(3600), None);
        assert_eq!(name, header::CACHE_CONTROL);
        assert_eq!(value.to_str().unwrap(), "public, max-age=3600");

        // Test max_age = 0, should return no-cache
        let (name, value) = new_cache_control_header(Some(0), Some(true));
        assert_eq!(name, header::CACHE_CONTROL);
        assert_eq!(value, HTTP_HEADER_NO_CACHE.clone().1);

        // Test max_age = None, should return no-cache
        let (name, value) = new_cache_control_header(None, Some(false));
        assert_eq!(name, header::CACHE_CONTROL);
        assert_eq!(value, HTTP_HEADER_NO_CACHE.clone().1);
    }

    #[test]
    fn test_http_response_builder_pattern() {
        // Create a custom Header
        let etag_header = (header::ETAG, HeaderValue::from_static("\"12345\""));
        let server_header =
            (header::SERVER, HeaderValue::from_static("MyTestServer"));

        let response = HttpResponse::builder(StatusCode::OK)
            .body("Test Body")
            .header(etag_header.clone())
            .headers(vec![server_header.clone()])
            .max_age(60, true) // 60 seconds, private
            .finish();

        assert_eq!(response.status, StatusCode::OK);
        assert_eq!(response.body, Bytes::from("Test Body"));
        assert_eq!(response.max_age, Some(60));
        assert_eq!(response.cache_private, Some(true));

        // Verify headers are correctly added
        let headers = response.headers.unwrap();
        assert_eq!(headers.len(), 2);
        assert!(headers.contains(&etag_header));
        assert!(headers.contains(&server_header));

        // Test no_store in the chain call
        let no_store_response = HttpResponse::builder(StatusCode::ACCEPTED)
            .no_store()
            .finish();

        assert_eq!(no_store_response.status, StatusCode::ACCEPTED);
        assert!(
            no_store_response
                .headers
                .unwrap()
                .contains(&HTTP_HEADER_NO_STORE.clone())
        );
    }

    /// Regression: every `Set-Cookie` of a response replaced the one before
    /// it. Any other header given twice is still the last one given.
    #[test]
    fn test_repeated_headers_are_kept() {
        let cookie = |value: &'static str| {
            (header::SET_COOKIE, HeaderValue::from_static(value))
        };
        let text = |value: &'static str| {
            (header::CONTENT_TYPE, HeaderValue::from_static(value))
        };
        let headers = vec![
            cookie("a=1"),
            text("text/html"),
            cookie("b=2; Expires=Wed, 21 Oct 2026 07:28:00 GMT"),
            text("text/plain"),
            cookie("c=3"),
        ];
        let values = |resp: &ResponseHeader, name: HeaderName| {
            resp.headers
                .get_all(name)
                .iter()
                .map(|value| value.to_str().unwrap().to_string())
                .collect::<Vec<_>>()
        };
        let expect_cookies =
            vec!["a=1", "b=2; Expires=Wed, 21 Oct 2026 07:28:00 GMT", "c=3"];

        let resp = HttpResponse {
            headers: Some(headers.clone()),
            ..Default::default()
        }
        .new_response_header()
        .unwrap();
        assert_eq!(expect_cookies, values(&resp, header::SET_COOKIE));
        assert_eq!(vec!["text/plain"], values(&resp, header::CONTENT_TYPE));

        let mut reader = tokio_test::io::Builder::new().build();
        let mut chunked = HttpChunkResponse::new(&mut reader);
        chunked.headers = Some(headers);
        let resp = chunked.get_response_header().unwrap();
        assert_eq!(expect_cookies, values(&resp, header::SET_COOKIE));
        assert_eq!(vec!["text/plain"], values(&resp, header::CONTENT_TYPE));
    }

    #[test]
    fn test_http_response_error_cases() {
        // Test redirect error cases (invalid location)
        // String containing \0 characters is invalid HeaderValue
        let invalid_location = "http://example.com/\0";
        let result = HttpResponse::redirect(invalid_location);
        assert!(result.is_err());
    }

    #[test]
    fn test_new_response_header_generation() {
        let resp = HttpResponse {
            status: StatusCode::OK,
            body: Bytes::from("Hello world!"),
            max_age: Some(3600),
            created_at: Some(get_super_ts().saturating_sub(10)), // 模拟10秒前创建
            cache_private: Some(true),
            headers: Some(vec![(
                header::CONTENT_ENCODING,
                HeaderValue::from_static("gzip"),
            )]),
        };

        let header = resp.new_response_header().unwrap();
        let headers_map: std::collections::HashMap<_, _> =
            header.headers.iter().collect();

        // 验证基本 headers
        assert_eq!(header.status, StatusCode::OK);
        assert_eq!(
            headers_map
                .get(&header::CONTENT_LENGTH)
                .unwrap()
                .to_str()
                .unwrap(),
            "12"
        );
        assert_eq!(
            headers_map
                .get(&header::CACHE_CONTROL)
                .unwrap()
                .to_str()
                .unwrap(),
            "private, max-age=3600"
        );
        assert_eq!(
            headers_map
                .get(&header::CONTENT_ENCODING)
                .unwrap()
                .to_str()
                .unwrap(),
            "gzip"
        );
    }
}