Skip to main content

pingap_core/
http_response.rs

1// Copyright 2024-2025 Tree xie.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use super::{
16    HTTP_HEADER_CONTENT_HTML, HTTP_HEADER_CONTENT_JSON,
17    HTTP_HEADER_CONTENT_TEXT, HTTP_HEADER_NO_CACHE, HTTP_HEADER_NO_STORE,
18    HTTP_HEADER_TRANSFER_CHUNKED, HttpHeader, LOG_TARGET, get_super_ts,
19    new_internal_error,
20};
21use bytes::{Bytes, BytesMut};
22use http::StatusCode;
23use http::header;
24use http::{HeaderName, HeaderValue};
25use pingora::http::ResponseHeader;
26use pingora::proxy::Session;
27use serde::Serialize;
28use std::pin::Pin;
29use tokio::io::AsyncReadExt;
30use tracing::error;
31
32/// A helper function to generate a `Cache-Control` header.
33///
34/// It determines the appropriate `Cache-Control` value based on the provided
35/// `max_age` and `cache_private` settings. This is a performance-sensitive
36/// function, so it's optimized to avoid allocations where possible.
37///
38/// # Arguments
39/// * `max_age`: An `Option<u32>` specifying the max-age in seconds. If `Some(0)` or `None`,
40///   a "no-cache" header is returned.
41/// * `cache_private`: An `Option<bool>` indicating if the cache is private. Defaults to public.
42///
43/// # Returns
44/// An `HttpHeader` tuple `(HeaderName, HeaderValue)` for the `Cache-Control` header.
45fn new_cache_control_header(
46    max_age: Option<u32>,
47    cache_private: Option<bool>,
48) -> HttpHeader {
49    // Determine the max_age value. If it's 0 or not provided, return the static no-cache header immediately.
50    let max_age = match max_age {
51        Some(0) | None => return HTTP_HEADER_NO_CACHE.clone(),
52        Some(age) => age,
53    };
54
55    // Determine the cache visibility ("private" or "public").
56    let category: &[u8] = if cache_private.unwrap_or_default() {
57        b"private"
58    } else {
59        b"public"
60    };
61
62    // Use a pre-allocated buffer (`BytesMut`) and the `itoa` crate to efficiently build the
63    // header value without creating an intermediate `String` via `format!`.
64    // The capacity is estimated to prevent reallocations.
65    let mut buf = BytesMut::with_capacity(category.len() + 9 + 10); // e.g., "public, max-age=" + up to 10 digits for a u32
66    buf.extend_from_slice(category);
67    buf.extend_from_slice(b", max-age=");
68    buf.extend_from_slice(itoa::Buffer::new().format(max_age).as_bytes());
69
70    // Try to create a `HeaderValue` from the constructed bytes.
71    if let Ok(value) = HeaderValue::from_bytes(&buf) {
72        return (header::CACHE_CONTROL, value);
73    }
74    // If creation fails for any reason, fall back to the safe "no-cache" header.
75    HTTP_HEADER_NO_CACHE.clone()
76}
77
78/// A builder for creating `HttpResponse` instances fluently using the builder pattern.
79///
80/// This allows for chaining method calls to configure a response, improving readability.
81#[derive(Default, Debug)]
82pub struct HttpResponseBuilder {
83    /// The `HttpResponse` instance being built.
84    response: HttpResponse,
85}
86
87impl HttpResponseBuilder {
88    /// Creates a new builder with a given status code.
89    pub fn new(status: StatusCode) -> Self {
90        Self {
91            response: HttpResponse {
92                status,
93                ..Default::default()
94            },
95        }
96    }
97
98    /// Sets the response body.
99    ///
100    /// This method is generic over `impl Into<Bytes>`, allowing various types like
101    /// `Vec<u8>`, `String`, or `&'static str` to be passed as the body.
102    pub fn body(mut self, body: impl Into<Bytes>) -> Self {
103        self.response.body = body.into();
104        self
105    }
106
107    /// Adds a single HTTP header to the response.
108    ///
109    /// If the headers vector doesn't exist yet, it will be created.
110    pub fn header(mut self, header: HttpHeader) -> Self {
111        self.response
112            .headers
113            .get_or_insert_with(Vec::new)
114            .push(header);
115        self
116    }
117
118    /// Appends multiple HTTP headers to the response.
119    pub fn headers(mut self, headers: Vec<HttpHeader>) -> Self {
120        self.response
121            .headers
122            .get_or_insert_with(Vec::new)
123            .extend(headers);
124        self
125    }
126
127    /// Sets the `Cache-Control` max-age and privacy directive.
128    pub fn max_age(mut self, seconds: u32, is_private: bool) -> Self {
129        self.response.max_age = Some(seconds);
130        self.response.cache_private = Some(is_private);
131        self
132    }
133
134    /// A convenience method to add a "no-store" `Cache-Control` header.
135    pub fn no_store(self) -> Self {
136        self.header(HTTP_HEADER_NO_STORE.clone())
137    }
138
139    /// Consumes the builder and returns the final `HttpResponse` instance.
140    pub fn finish(self) -> HttpResponse {
141        self.response
142    }
143}
144
145/// Represents a complete HTTP response, including status, headers, and a body.
146/// This struct is used for responses where the entire body is known upfront.
147#[derive(Default, Clone, Debug)]
148pub struct HttpResponse {
149    /// The HTTP status code (e.g., 200 OK, 404 Not Found).
150    pub status: StatusCode,
151    /// The response body, stored as a `Bytes` object for efficiency.
152    pub body: Bytes,
153    /// The `max-age` directive for the `Cache-Control` header, in seconds.
154    pub max_age: Option<u32>,
155    /// The UNIX timestamp when the response content was created, used for the `Age` header.
156    pub created_at: Option<u32>,
157    /// A flag to indicate if the `Cache-Control` should be "private" (true) or "public" (false).
158    pub cache_private: Option<bool>,
159    /// A list of additional HTTP headers to include in the response.
160    pub headers: Option<Vec<HttpHeader>>,
161}
162
163/// The failure to write a response as what it is: the client's side of
164/// the connection.
165///
166/// pingora names the source of an error on the path of a proxied response
167/// and leaves it unset on these writes, which a plugin or the proxy itself
168/// answers with. Unset, a client that went away while a file was sent to
169/// it counted as a failure of the proxy: a `500` in the access log and
170/// the metrics, and a line at error level for every download that was
171/// cancelled, where the same on a proxied response is a `499`.
172fn to_client<T>(result: pingora::Result<T>) -> pingora::Result<T> {
173    result.map_err(|e| e.into_down())
174}
175
176impl HttpResponse {
177    /// Returns a new `HttpResponseBuilder` to start building a response.
178    pub fn builder(status: StatusCode) -> HttpResponseBuilder {
179        HttpResponseBuilder::new(status)
180    }
181
182    /// A convenience constructor for a 204 No Content response.
183    pub fn no_content() -> Self {
184        Self::builder(StatusCode::NO_CONTENT).no_store().finish()
185    }
186
187    /// A convenience constructor for a 400 Bad Request response.
188    pub fn bad_request(body: impl Into<Bytes>) -> Self {
189        Self::builder(StatusCode::BAD_REQUEST)
190            .body(body)
191            .header(HTTP_HEADER_CONTENT_TEXT.clone())
192            .no_store()
193            .finish()
194    }
195
196    /// A convenience constructor for a 404 Not Found response.
197    pub fn not_found(body: impl Into<Bytes>) -> Self {
198        Self::builder(StatusCode::NOT_FOUND)
199            .body(body)
200            .header(HTTP_HEADER_CONTENT_TEXT.clone())
201            .no_store()
202            .finish()
203    }
204
205    /// A convenience constructor for a 500 Internal Server Error response.
206    pub fn unknown_error(body: impl Into<Bytes>) -> Self {
207        Self::builder(StatusCode::INTERNAL_SERVER_ERROR)
208            .body(body)
209            .header(HTTP_HEADER_CONTENT_TEXT.clone())
210            .no_store()
211            .finish()
212    }
213
214    /// A convenience constructor for a 200 OK HTML response.
215    pub fn html(body: impl Into<Bytes>) -> Self {
216        Self::builder(StatusCode::OK)
217            .body(body)
218            .header(HTTP_HEADER_CONTENT_HTML.clone())
219            .header(HTTP_HEADER_NO_CACHE.clone())
220            .finish()
221    }
222
223    /// A convenience constructor for a 302 Temporary Redirect response.
224    pub fn redirect(location: &str) -> pingora::Result<Self> {
225        // Attempt to parse the location string into a valid HeaderValue.
226        let value = HeaderValue::from_str(location).map_err(|e| {
227            error!(error = e.to_string(), "to header value fail");
228            new_internal_error(500, e)
229        })?;
230        // Build the redirect response.
231        Ok(Self::builder(StatusCode::TEMPORARY_REDIRECT)
232            .header((header::LOCATION, value))
233            .header(HTTP_HEADER_NO_CACHE.clone())
234            .finish())
235    }
236
237    /// A convenience constructor for a 200 OK plain text response.
238    pub fn text(body: impl Into<Bytes>) -> Self {
239        Self::builder(StatusCode::OK)
240            .body(body)
241            .header(HTTP_HEADER_CONTENT_TEXT.clone())
242            .header(HTTP_HEADER_NO_CACHE.clone())
243            .finish()
244    }
245
246    /// Creates a 200 OK JSON response by serializing the given value.
247    pub fn try_from_json<T>(value: &T) -> pingora::Result<Self>
248    where
249        T: ?Sized + Serialize,
250    {
251        // Serialize the value to a JSON byte vector.
252        let buf = serde_json::to_vec(value).map_err(|e| {
253            error!(target: LOG_TARGET, error = e.to_string(), "to json fail");
254            new_internal_error(400, e)
255        })?;
256        // Build the JSON response.
257        Ok(Self::builder(StatusCode::OK)
258            .body(buf)
259            .header(HTTP_HEADER_CONTENT_JSON.clone())
260            .finish())
261    }
262
263    /// Creates a JSON response with a specified status code by serializing the given value.
264    pub fn try_from_json_status<T>(
265        value: &T,
266        status: StatusCode,
267    ) -> pingora::Result<Self>
268    where
269        T: ?Sized + Serialize,
270    {
271        // First, create a standard 200 OK JSON response.
272        let mut resp = Self::try_from_json(value)?;
273        // Then, simply overwrite the status code.
274        resp.status = status;
275        Ok(resp)
276    }
277
278    /// Builds a `pingora::http::ResponseHeader` from the `HttpResponse`'s properties.
279    pub fn new_response_header(&self) -> pingora::Result<ResponseHeader> {
280        // Build the response header with the status code, sized for the
281        // headers added below: Content-Length, Cache-Control, Age, plus
282        // the custom ones.
283        let custom = self.headers.as_ref().map_or(0, |headers| headers.len());
284        let mut resp = ResponseHeader::build(self.status, Some(3 + custom))?;
285
286        // A local helper closure to simplify adding headers and handling potential errors.
287        let mut add_header =
288            |name: &HeaderName, value: &HeaderValue| -> pingora::Result<()> {
289                resp.insert_header(name, value)?;
290                Ok(())
291            };
292
293        // Add the Content-Length header based on the body size. Not on a
294        // response that has no body by definition: a 204 must not carry
295        // one (RFC 9110 8.6), and on a 304 a zero would be the length of
296        // the representation, which it is not.
297        if !has_no_body(self.status) {
298            add_header(
299                &header::CONTENT_LENGTH,
300                &HeaderValue::from(self.body.len()),
301            )?;
302        }
303
304        // Generate and add the Cache-Control header.
305        let (name, value) =
306            new_cache_control_header(self.max_age, self.cache_private);
307        add_header(&name, &value)?;
308
309        // If a creation timestamp is provided, calculate and add the Age header.
310        if let Some(created_at) = self.created_at {
311            let secs = get_super_ts().saturating_sub(created_at);
312            add_header(&header::AGE, &HeaderValue::from(secs))?;
313        }
314
315        // Add all custom headers from the `headers` field.
316        if let Some(headers) = &self.headers {
317            put_custom_headers(&mut resp, headers)?;
318        }
319        Ok(resp)
320    }
321
322    /// Sends the entire HTTP response (header and body) to the client via the session.
323    pub async fn send(self, session: &mut Session) -> pingora::Result<usize> {
324        // First, build the response header.
325        let header = self.new_response_header()?;
326        self.send_with_header(session, header).await
327    }
328
329    /// Sends the body under `header`: the one `new_response_header` builds,
330    /// with whatever the caller has added to it.
331    pub async fn send_with_header(
332        self,
333        session: &mut Session,
334        header: ResponseHeader,
335    ) -> pingora::Result<usize> {
336        // A HEAD, a 204 or a 304 ends with its header. HTTP/1.1 drops a
337        // body written after it, HTTP/2 does not: the DATA frame went out
338        // and the client reset the stream with a protocol error.
339        if session.req_header().method == http::Method::HEAD
340            || has_no_body(self.status)
341        {
342            to_client(
343                session.write_response_header(Box::new(header), true).await,
344            )?;
345            to_client(session.finish_body().await)?;
346            return Ok(0);
347        }
348        let size = self.body.len();
349        // Write the header to the session.
350        to_client(
351            session.write_response_header(Box::new(header), false).await,
352        )?;
353        // Write the body to the session. `end_stream` is true because this is the final body part.
354        to_client(session.write_response_body(Some(self.body), true).await)?;
355        // Finalize the response stream.
356        to_client(session.finish_body().await)?;
357        Ok(size)
358    }
359}
360
361/// The statuses whose response has no body, whatever the request.
362#[inline]
363fn has_no_body(status: StatusCode) -> bool {
364    status.is_informational()
365        || status == StatusCode::NO_CONTENT
366        || status == StatusCode::NOT_MODIFIED
367}
368
369/// The headers that are sent once per value and cannot be folded into one
370/// line: a cookie's `Expires` has a comma of its own, and a challenge is a
371/// list already.
372fn is_repeated_header(name: &HeaderName) -> bool {
373    [
374        header::SET_COOKIE,
375        header::WWW_AUTHENTICATE,
376        header::PROXY_AUTHENTICATE,
377        header::LINK,
378    ]
379    .contains(name)
380}
381
382/// Puts the custom headers of a response into `resp`.
383///
384/// A header given again replaces the earlier one, which is how a configured
385/// `Content-Type` takes the place of the detected one. Not so for the few
386/// headers that come once per value: each `Set-Cookie` used to replace the
387/// last, and of the three an auth service answered with, the client got
388/// one.
389fn put_custom_headers(
390    resp: &mut ResponseHeader,
391    headers: &[HttpHeader],
392) -> pingora::Result<()> {
393    for (index, (name, value)) in headers.iter().enumerate() {
394        if is_repeated_header(name)
395            && headers[..index].iter().any(|(earlier, _)| earlier == name)
396        {
397            resp.append_header(name, value)?;
398        } else {
399            resp.insert_header(name, value)?;
400        }
401    }
402    Ok(())
403}
404
405/// Represents a chunked HTTP response for streaming large bodies of data.
406///
407/// This is used when the response body is too large to fit in memory or is generated on-the-fly.
408pub struct HttpChunkResponse<'r, R> {
409    /// A pinned, mutable reference to an async reader that provides the body data.
410    pub reader: Pin<&'r mut R>,
411    /// The status of the response. Defaults to `200 OK`.
412    pub status: StatusCode,
413    /// The suggested size for each data chunk. Defaults to `DEFAULT_BUF_SIZE`.
414    pub chunk_size: usize,
415    /// Cache control `max-age` setting for the response.
416    pub max_age: Option<u32>,
417    /// Cache control privacy setting for the response.
418    pub cache_private: Option<bool>,
419    /// Additional headers to include in the response.
420    pub headers: Option<Vec<HttpHeader>>,
421    /// How fast the body is sent, where that is limited.
422    pub pace: Option<crate::BodyPace>,
423}
424
425/// The default buffer size (8KB) for chunked responses.
426const DEFAULT_BUF_SIZE: usize = 8 * 1024;
427
428impl<'r, R> HttpChunkResponse<'r, R>
429where
430    // The reader must implement `AsyncRead` and be `Unpin`.
431    R: tokio::io::AsyncRead + std::marker::Unpin,
432{
433    /// Creates a new `HttpChunkResponse` with a given reader and default settings.
434    pub fn new(r: &'r mut R) -> Self {
435        Self {
436            reader: Pin::new(r),
437            status: StatusCode::OK,
438            chunk_size: DEFAULT_BUF_SIZE,
439            max_age: None,
440            headers: None,
441            cache_private: None,
442            pace: None,
443        }
444    }
445
446    /// Builds the `ResponseHeader` for the streamed response.
447    ///
448    /// The body is framed by `Content-Length` when the custom headers give
449    /// one, and by `Transfer-Encoding: chunked` otherwise. Both at once is
450    /// not a valid message, and that is what a streamed file used to get:
451    /// its length was known and sent, and `chunked` was added regardless.
452    pub fn get_response_header(&self) -> pingora::Result<ResponseHeader> {
453        let mut resp = ResponseHeader::build(self.status, Some(4))?;
454        // Add any custom headers.
455        if let Some(headers) = &self.headers {
456            put_custom_headers(&mut resp, headers)?;
457        }
458
459        if !resp.headers.contains_key(header::CONTENT_LENGTH) {
460            let chunked = HTTP_HEADER_TRANSFER_CHUNKED.clone();
461            resp.insert_header(chunked.0, chunked.1)?;
462        }
463
464        // Add the `Cache-Control` header.
465        let cache_control =
466            new_cache_control_header(self.max_age, self.cache_private);
467        resp.insert_header(cache_control.0, cache_control.1)?;
468        Ok(resp)
469    }
470
471    /// Reads the next chunk of the body. `None` only at the end of the
472    /// stream: a read that comes back short is a chunk like any other, not a
473    /// sign that the reader is done - `AsyncRead` never promised to fill the
474    /// buffer.
475    ///
476    /// The chunk is split off `buffer` rather than copied out of it; once
477    /// the session has written it and dropped it, `reserve` gets the
478    /// allocation back for the next read.
479    async fn next_chunk(
480        &mut self,
481        buffer: &mut BytesMut,
482    ) -> pingora::Result<Option<Bytes>> {
483        buffer.reserve(self.chunk_size.max(512));
484        let size = self.reader.read_buf(buffer).await.map_err(|e| {
485            error!(error = e.to_string(), "read data fail");
486            new_internal_error(400, e)
487        })?;
488        if size == 0 {
489            return Ok(None);
490        }
491        Ok(Some(buffer.split().freeze()))
492    }
493
494    /// Streams the data from the reader to the client as a chunked response.
495    ///
496    /// Reads data from the `reader` in chunks and sends each chunk to the client until
497    /// the reader is exhausted.
498    pub async fn send(self, session: &mut Session) -> pingora::Result<usize> {
499        let header = self.get_response_header()?;
500        self.send_with_header(session, header).await
501    }
502
503    /// Sends the body under `header`: the one `get_response_header` builds,
504    /// with whatever the caller has added to it.
505    pub async fn send_with_header(
506        mut self,
507        session: &mut Session,
508        header: ResponseHeader,
509    ) -> pingora::Result<usize> {
510        // `end_stream` is false because a body will follow.
511        to_client(
512            session.write_response_header(Box::new(header), false).await,
513        )?;
514
515        let mut sent = 0;
516        let mut buffer = BytesMut::with_capacity(self.chunk_size.max(512));
517        // One chunk of lookahead, so the last one goes out flagged as the end
518        // of the stream.
519        let mut pending = self.next_chunk(&mut buffer).await?;
520        while let Some(chunk) = pending {
521            let len = chunk.len();
522            sent += len;
523            pending = self.next_chunk(&mut buffer).await?;
524            to_client(
525                session
526                    .write_response_body(Some(chunk), pending.is_none())
527                    .await,
528            )?;
529            // Not after the last chunk: there is nothing left to hold
530            // back.
531            if pending.is_some()
532                && let Some(delay) =
533                    self.pace.as_mut().and_then(|pace| pace.delay(len))
534            {
535                tokio::time::sleep(delay).await;
536            }
537        }
538        // Finalize the response stream.
539        to_client(session.finish_body().await)?;
540
541        Ok(sent)
542    }
543}
544#[cfg(test)]
545mod tests {
546    use super::*;
547    use crate::convert_headers;
548    use bytes::Bytes;
549    use http::StatusCode;
550    use pretty_assertions::assert_eq;
551    use serde::Serialize;
552    use std::io::Write;
553    use tempfile::NamedTempFile;
554    use tokio::fs;
555    #[test]
556    fn test_new_cache_control_header() {
557        assert_eq!(
558            r###"("cache-control", "private, max-age=3600")"###,
559            format!("{:?}", new_cache_control_header(Some(3600), Some(true)))
560        );
561        assert_eq!(
562            r###"("cache-control", "public, max-age=3600")"###,
563            format!("{:?}", new_cache_control_header(Some(3600), None))
564        );
565        assert_eq!(
566            r###"("cache-control", "private, no-cache")"###,
567            format!("{:?}", new_cache_control_header(Some(0), Some(false)))
568        );
569        assert_eq!(
570            r###"("cache-control", "private, no-cache")"###,
571            format!("{:?}", new_cache_control_header(None, None))
572        );
573    }
574
575    #[test]
576    fn test_http_response() {
577        assert_eq!(
578            r###"HttpResponse { status: 204, body: b"", max_age: None, created_at: None, cache_private: None, headers: Some([("cache-control", "private, no-store")]) }"###,
579            format!("{:?}", HttpResponse::no_content())
580        );
581        assert_eq!(
582            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")]) }"###,
583            format!("{:?}", HttpResponse::not_found("Not Found"))
584        );
585        assert_eq!(
586            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")]) }"###,
587            format!("{:?}", HttpResponse::unknown_error("Unknown Error"))
588        );
589
590        assert_eq!(
591            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")]) }"###,
592            format!("{:?}", HttpResponse::bad_request("Bad Request"))
593        );
594
595        assert_eq!(
596            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")]) }"###,
597            format!("{:?}", HttpResponse::html("<p>Pingap</p>"))
598        );
599
600        assert_eq!(
601            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")]) }"###,
602            format!(
603                "{:?}",
604                HttpResponse::redirect("http://example.com/").unwrap()
605            )
606        );
607
608        assert_eq!(
609            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")]) }"###,
610            format!("{:?}", HttpResponse::text("Hello World!"))
611        );
612
613        #[derive(Serialize)]
614        struct Data {
615            message: String,
616        }
617        let resp = HttpResponse::try_from_json_status(
618            &Data {
619                message: "Hello World!".to_string(),
620            },
621            StatusCode::BAD_REQUEST,
622        )
623        .unwrap();
624        assert_eq!(
625            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")]) }"###,
626            format!("{resp:?}")
627        );
628        let resp = HttpResponse::try_from_json(&Data {
629            message: "Hello World!".to_string(),
630        })
631        .unwrap();
632        assert_eq!(
633            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")]) }"###,
634            format!("{resp:?}")
635        );
636
637        let resp = HttpResponse {
638            status: StatusCode::OK,
639            body: Bytes::from("Hello world!"),
640            max_age: Some(3600),
641            created_at: Some(0),
642            cache_private: Some(true),
643            headers: Some(
644                convert_headers(&[
645                    "Contont-Type: application/json".to_string(),
646                    "Content-Encoding: gzip".to_string(),
647                ])
648                .unwrap(),
649            ),
650        };
651
652        let mut header = resp.new_response_header().unwrap();
653        assert_eq!(true, !header.headers.get("Age").unwrap().is_empty());
654        header.remove_header("Age").unwrap();
655
656        assert_eq!(
657            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 }"###,
658            format!("{header:?}")
659        );
660    }
661
662    #[tokio::test]
663    async fn test_http_chunk_response() {
664        let file = include_bytes!("http_response.rs");
665        let mut f = NamedTempFile::new().unwrap();
666        f.write_all(file).unwrap();
667        let mut f = fs::OpenOptions::new().read(true).open(f).await.unwrap();
668        let mut resp = HttpChunkResponse::new(&mut f);
669        resp.max_age = Some(3600);
670        resp.cache_private = Some(false);
671        resp.headers = Some(
672            convert_headers(&["Contont-Type: text/html".to_string()]).unwrap(),
673        );
674        let header = resp.get_response_header().unwrap();
675        assert_eq!(
676            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 }"###,
677            format!("{header:?}")
678        );
679    }
680
681    /// A known length frames the body by itself, and the status is the one
682    /// that was set.
683    #[tokio::test]
684    async fn test_http_chunk_response_with_content_length() {
685        let mut reader = tokio_test::io::Builder::new().build();
686        let mut resp = HttpChunkResponse::new(&mut reader);
687        resp.status = StatusCode::PARTIAL_CONTENT;
688        resp.headers = Some(
689            convert_headers(&[
690                "Content-Length: 10".to_string(),
691                "Content-Range: bytes 0-9/100".to_string(),
692            ])
693            .unwrap(),
694        );
695        let header = resp.get_response_header().unwrap();
696        assert_eq!(StatusCode::PARTIAL_CONTENT, header.status);
697        assert_eq!("10", header.headers.get("content-length").unwrap());
698        assert_eq!(false, header.headers.contains_key("transfer-encoding"));
699    }
700
701    /// A chunked response keeps the pace it is given: a wait after each
702    /// chunk for what has gone out, and none after the last, where there
703    /// is nothing left to hold back.
704    #[tokio::test]
705    async fn test_http_chunk_response_keeps_the_pace() {
706        use tokio::io::AsyncWriteExt;
707        // The waits are told on the clock of the runtime, which stands
708        // still but for them.
709        tokio::time::pause();
710        let (mut client, server) = tokio::io::duplex(64 * 1024);
711        client.write_all(b"GET / HTTP/1.1\r\n\r\n").await.unwrap();
712        let mut session = Session::new_h1(Box::new(server));
713        session.read_request().await.unwrap();
714
715        let send = async |session: &mut Session, pace| {
716            let mut reader = std::io::Cursor::new(vec![b'a'; 3000]);
717            let mut resp = HttpChunkResponse::new(&mut reader);
718            resp.chunk_size = 1000;
719            resp.pace = pace;
720            let start = tokio::time::Instant::now();
721            let sent = resp.send(session).await.unwrap();
722            (sent, start.elapsed().as_millis())
723        };
724        // Ten thousand bytes a second, in chunks of a thousand: a tenth
725        // of a second after the first, two tenths in all after the
726        // second, and the third ends it. A wait after the last as well
727        // would make it six tenths.
728        let (sent, waited) =
729            send(&mut session, Some(crate::BodyPace::new(10_000, 0))).await;
730        assert_eq!(3000, sent);
731        // (A timer of the runtime ends on a whole millisecond, hence the
732        // little that is allowed above three tenths.)
733        assert_eq!(true, (200..=350).contains(&waited), "{waited}");
734    }
735
736    #[tokio::test]
737    async fn test_http_chunk_response_short_reads() {
738        // A reader that hands out less than a chunk at a time. Only an empty
739        // read ends the body; a short one used to, truncating the file.
740        let mut reader = tokio_test::io::Builder::new()
741            .read(b"abc")
742            .read(b"defgh")
743            .build();
744        let mut resp = HttpChunkResponse::new(&mut reader);
745        resp.chunk_size = 1024;
746        let mut buffer = BytesMut::new();
747        assert_eq!(
748            Some(Bytes::from_static(b"abc")),
749            resp.next_chunk(&mut buffer).await.unwrap()
750        );
751        assert_eq!(
752            Some(Bytes::from_static(b"defgh")),
753            resp.next_chunk(&mut buffer).await.unwrap()
754        );
755        assert_eq!(None, resp.next_chunk(&mut buffer).await.unwrap());
756    }
757
758    /// Regression: a 204 and a 304 carried `Content-Length: 0`. A 204 must
759    /// not have the header at all, and on a 304 it would be the length of
760    /// the representation, which is not zero.
761    #[test]
762    fn test_no_content_length_without_a_body() {
763        for (status, expected) in [
764            (StatusCode::NO_CONTENT, None),
765            (StatusCode::NOT_MODIFIED, None),
766            (StatusCode::SWITCHING_PROTOCOLS, None),
767            (StatusCode::OK, Some("0")),
768            (StatusCode::NOT_FOUND, Some("0")),
769        ] {
770            let header = HttpResponse {
771                status,
772                ..Default::default()
773            }
774            .new_response_header()
775            .unwrap();
776            assert_eq!(
777                expected,
778                header
779                    .headers
780                    .get(header::CONTENT_LENGTH)
781                    .and_then(|value| value.to_str().ok()),
782                "{status}"
783            );
784        }
785        // The length a HEAD is answered with is the caller's, and stays.
786        let header = HttpResponse {
787            status: StatusCode::NOT_MODIFIED,
788            headers: Some(vec![(
789                header::CONTENT_LENGTH,
790                HeaderValue::from_static("12"),
791            )]),
792            ..Default::default()
793        }
794        .new_response_header()
795        .unwrap();
796        assert_eq!("12", header.headers.get(header::CONTENT_LENGTH).unwrap());
797    }
798
799    /// A response that cannot be written is the client's side failing,
800    /// and says so. Regression: pingora leaves the source of such an error
801    /// unset here, and unset it counted as a failure of the proxy - a
802    /// `500` and a line at error level for every download a client
803    /// cancelled, where the same on a proxied response is a `499`.
804    #[tokio::test]
805    async fn test_a_failed_write_is_the_clients() {
806        let gone = async || {
807            let (mut client, server) = tokio::io::duplex(1024);
808            tokio::io::AsyncWriteExt::write_all(
809                &mut client,
810                b"GET / HTTP/1.1\r\n\r\n",
811            )
812            .await
813            .unwrap();
814            let mut session = Session::new_h1(Box::new(server));
815            session.read_request().await.unwrap();
816            // the client goes away with its request unanswered
817            drop(client);
818            session
819        };
820        let body = vec![b'x'; 256 * 1024];
821
822        let mut session = gone().await;
823        let error = HttpResponse {
824            status: StatusCode::OK,
825            body: Bytes::from(body.clone()),
826            ..Default::default()
827        }
828        .send(&mut session)
829        .await
830        .expect_err("nobody is there to be written to");
831        assert_eq!(&pingora::ErrorSource::Downstream, error.esource());
832
833        let mut session = gone().await;
834        let mut reader = std::io::Cursor::new(body);
835        let error = HttpChunkResponse::new(&mut reader)
836            .send(&mut session)
837            .await
838            .expect_err("nobody is there to be written to");
839        assert_eq!(&pingora::ErrorSource::Downstream, error.esource());
840    }
841
842    /// Regression: the body of an answer to a HEAD, or of a 204 or a 304,
843    /// was written after the header. HTTP/1.1 drops it; over HTTP/2 the
844    /// DATA frame went out and the client reset the stream with a protocol
845    /// error. Such a response now ends with its header.
846    ///
847    /// Over HTTP/1.1, which is what this runs on, the wire looks the same
848    /// before and after; what tells them apart here is the number of body
849    /// bytes `send` reports as written. The HTTP/2 case is only checked
850    /// end to end (`curl --http2-prior-knowledge -I`).
851    #[tokio::test]
852    async fn test_send_writes_no_body_where_there_is_none() {
853        let send = async |request: &str, status: StatusCode| {
854            let (mut client, server) = tokio::io::duplex(64 * 1024);
855            tokio::io::AsyncWriteExt::write_all(
856                &mut client,
857                request.as_bytes(),
858            )
859            .await
860            .unwrap();
861            let mut session = Session::new_h1(Box::new(server));
862            session.read_request().await.unwrap();
863            let sent = HttpResponse {
864                status,
865                body: Bytes::from_static(b"Hello world!"),
866                ..Default::default()
867            }
868            .send(&mut session)
869            .await
870            .unwrap();
871            drop(session);
872            let mut buf = vec![];
873            tokio::io::AsyncReadExt::read_to_end(&mut client, &mut buf)
874                .await
875                .unwrap();
876            let text = String::from_utf8_lossy(&buf).into_owned();
877            let (_, body) = text.split_once("\r\n\r\n").unwrap();
878            (sent, body.to_string())
879        };
880        assert_eq!(
881            (12, "Hello world!".to_string()),
882            send("GET / HTTP/1.1\r\n\r\n", StatusCode::OK).await
883        );
884        assert_eq!(
885            (0, String::new()),
886            send("HEAD / HTTP/1.1\r\n\r\n", StatusCode::OK).await
887        );
888        assert_eq!(
889            (0, String::new()),
890            send("GET / HTTP/1.1\r\n\r\n", StatusCode::NO_CONTENT).await
891        );
892        assert_eq!(
893            (0, String::new()),
894            send("GET / HTTP/1.1\r\n\r\n", StatusCode::NOT_MODIFIED).await
895        );
896    }
897
898    #[test]
899    fn test_new_cache_control_header_logic() {
900        // Test max_age > 0 case
901        let (name, value) = new_cache_control_header(Some(3600), Some(true));
902        assert_eq!(name, header::CACHE_CONTROL);
903        assert_eq!(value.to_str().unwrap(), "private, max-age=3600");
904
905        let (name, value) = new_cache_control_header(Some(3600), Some(false));
906        assert_eq!(name, header::CACHE_CONTROL);
907        assert_eq!(value.to_str().unwrap(), "public, max-age=3600");
908
909        // Test None is equivalent to public
910        let (name, value) = new_cache_control_header(Some(3600), None);
911        assert_eq!(name, header::CACHE_CONTROL);
912        assert_eq!(value.to_str().unwrap(), "public, max-age=3600");
913
914        // Test max_age = 0, should return no-cache
915        let (name, value) = new_cache_control_header(Some(0), Some(true));
916        assert_eq!(name, header::CACHE_CONTROL);
917        assert_eq!(value, HTTP_HEADER_NO_CACHE.clone().1);
918
919        // Test max_age = None, should return no-cache
920        let (name, value) = new_cache_control_header(None, Some(false));
921        assert_eq!(name, header::CACHE_CONTROL);
922        assert_eq!(value, HTTP_HEADER_NO_CACHE.clone().1);
923    }
924
925    #[test]
926    fn test_http_response_builder_pattern() {
927        // Create a custom Header
928        let etag_header = (header::ETAG, HeaderValue::from_static("\"12345\""));
929        let server_header =
930            (header::SERVER, HeaderValue::from_static("MyTestServer"));
931
932        let response = HttpResponse::builder(StatusCode::OK)
933            .body("Test Body")
934            .header(etag_header.clone())
935            .headers(vec![server_header.clone()])
936            .max_age(60, true) // 60 seconds, private
937            .finish();
938
939        assert_eq!(response.status, StatusCode::OK);
940        assert_eq!(response.body, Bytes::from("Test Body"));
941        assert_eq!(response.max_age, Some(60));
942        assert_eq!(response.cache_private, Some(true));
943
944        // Verify headers are correctly added
945        let headers = response.headers.unwrap();
946        assert_eq!(headers.len(), 2);
947        assert!(headers.contains(&etag_header));
948        assert!(headers.contains(&server_header));
949
950        // Test no_store in the chain call
951        let no_store_response = HttpResponse::builder(StatusCode::ACCEPTED)
952            .no_store()
953            .finish();
954
955        assert_eq!(no_store_response.status, StatusCode::ACCEPTED);
956        assert!(
957            no_store_response
958                .headers
959                .unwrap()
960                .contains(&HTTP_HEADER_NO_STORE.clone())
961        );
962    }
963
964    /// Regression: every `Set-Cookie` of a response replaced the one before
965    /// it. Any other header given twice is still the last one given.
966    #[test]
967    fn test_repeated_headers_are_kept() {
968        let cookie = |value: &'static str| {
969            (header::SET_COOKIE, HeaderValue::from_static(value))
970        };
971        let text = |value: &'static str| {
972            (header::CONTENT_TYPE, HeaderValue::from_static(value))
973        };
974        let headers = vec![
975            cookie("a=1"),
976            text("text/html"),
977            cookie("b=2; Expires=Wed, 21 Oct 2026 07:28:00 GMT"),
978            text("text/plain"),
979            cookie("c=3"),
980        ];
981        let values = |resp: &ResponseHeader, name: HeaderName| {
982            resp.headers
983                .get_all(name)
984                .iter()
985                .map(|value| value.to_str().unwrap().to_string())
986                .collect::<Vec<_>>()
987        };
988        let expect_cookies =
989            vec!["a=1", "b=2; Expires=Wed, 21 Oct 2026 07:28:00 GMT", "c=3"];
990
991        let resp = HttpResponse {
992            headers: Some(headers.clone()),
993            ..Default::default()
994        }
995        .new_response_header()
996        .unwrap();
997        assert_eq!(expect_cookies, values(&resp, header::SET_COOKIE));
998        assert_eq!(vec!["text/plain"], values(&resp, header::CONTENT_TYPE));
999
1000        let mut reader = tokio_test::io::Builder::new().build();
1001        let mut chunked = HttpChunkResponse::new(&mut reader);
1002        chunked.headers = Some(headers);
1003        let resp = chunked.get_response_header().unwrap();
1004        assert_eq!(expect_cookies, values(&resp, header::SET_COOKIE));
1005        assert_eq!(vec!["text/plain"], values(&resp, header::CONTENT_TYPE));
1006    }
1007
1008    #[test]
1009    fn test_http_response_error_cases() {
1010        // Test redirect error cases (invalid location)
1011        // String containing \0 characters is invalid HeaderValue
1012        let invalid_location = "http://example.com/\0";
1013        let result = HttpResponse::redirect(invalid_location);
1014        assert!(result.is_err());
1015    }
1016
1017    #[test]
1018    fn test_new_response_header_generation() {
1019        let resp = HttpResponse {
1020            status: StatusCode::OK,
1021            body: Bytes::from("Hello world!"),
1022            max_age: Some(3600),
1023            created_at: Some(get_super_ts().saturating_sub(10)), // 模拟10秒前创建
1024            cache_private: Some(true),
1025            headers: Some(vec![(
1026                header::CONTENT_ENCODING,
1027                HeaderValue::from_static("gzip"),
1028            )]),
1029        };
1030
1031        let header = resp.new_response_header().unwrap();
1032        let headers_map: std::collections::HashMap<_, _> =
1033            header.headers.iter().collect();
1034
1035        // 验证基本 headers
1036        assert_eq!(header.status, StatusCode::OK);
1037        assert_eq!(
1038            headers_map
1039                .get(&header::CONTENT_LENGTH)
1040                .unwrap()
1041                .to_str()
1042                .unwrap(),
1043            "12"
1044        );
1045        assert_eq!(
1046            headers_map
1047                .get(&header::CACHE_CONTROL)
1048                .unwrap()
1049                .to_str()
1050                .unwrap(),
1051            "private, max-age=3600"
1052        );
1053        assert_eq!(
1054            headers_map
1055                .get(&header::CONTENT_ENCODING)
1056                .unwrap()
1057                .to_str()
1058                .unwrap(),
1059            "gzip"
1060        );
1061    }
1062}