feat: support cache control and age http header for response

src/cache/http_header.rs

@@ -1,5 +1,7 @@
 use crate::utils::split_to_two_trim;
+use http::header;
 use http::{HeaderName, HeaderValue};
+use once_cell::sync::Lazy;
 use snafu::{ResultExt, Snafu};
 use std::str::FromStr;
 
@@ -32,3 +34,17 @@ pub fn convert_headers(header_values: &[String]) -> Result<Vec<HttpHeader>> {
     }
     Ok(arr)
 }
+
+pub static HTTP_HEADER_NO_STORE: Lazy<HttpHeader> = Lazy::new(|| {
+    (
+        header::CACHE_CONTROL,
+        HeaderValue::from_str("private, no-store").unwrap(),
+    )
+});
+
+pub static HTTP_HEADER_CONTENT_JSON: Lazy<HttpHeader> = Lazy::new(|| {
+    (
+        header::CONTENT_TYPE,
+        HeaderValue::from_str("application/json; charset=utf-8").unwrap(),
+    )
+});

src/cache/http_response.rs

@@ -1,26 +1,64 @@
-use super::HttpHeader;
+use super::{HttpHeader, HTTP_HEADER_NO_STORE};
 use bytes::Bytes;
+use http::header;
 use http::StatusCode;
 use pingora::{http::ResponseHeader, proxy::Session};
+use std::time::{SystemTime, UNIX_EPOCH};
 
 #[derive(Default)]
 pub struct HttpResponse {
     pub status: StatusCode,
     pub body: Bytes,
     pub max_age: Option<u32>,
     pub created_at: Option<u64>,
+    pub private: Option<bool>,
     pub headers: Option<Vec<HttpHeader>>,
 }
 
 impl HttpResponse {
     pub async fn send(&self, session: &mut Session) -> pingora::Result<usize> {
-        let mut resp = ResponseHeader::build(self.status, Some(4))?;
+        let fix_size = 3;
+        let size = self
+            .headers
+            .as_ref()
+            .map_or_else(|| fix_size, |headers| headers.len() + fix_size);
+        let mut resp = ResponseHeader::build(self.status, Some(size))?;
         resp.insert_header(http::header::CONTENT_LENGTH, self.body.len().to_string())?;
+
+        // set cache control
+        if let Some(max_age) = self.max_age {
+            let category = if self.private.unwrap_or_default() {
+                "private"
+            } else {
+                "public"
+            };
+            if let Ok(value) =
+                header::HeaderValue::from_str(&format!("{category}, max-age={max_age}"))
+            {
+                resp.insert_header(header::CACHE_CONTROL, value)?;
+            }
+        } else {
+            let h = HTTP_HEADER_NO_STORE.clone();
+            resp.insert_header(h.0, h.1)?;
+        }
+
+        if let Some(created_at) = self.created_at {
+            let secs = SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap_or_default()
+                .as_secs()
+                - created_at;
+            if let Ok(value) = header::HeaderValue::from_str(&secs.to_string()) {
+                resp.insert_header(header::AGE, value)?;
+            }
+        }
+
         if let Some(headers) = &self.headers {
             for (name, value) in headers {
                 resp.insert_header(name.to_owned(), value)?;
             }
         }
+
         let buf = self.body.clone();
         let size = buf.len();
         session.write_response_header(Box::new(resp)).await?;

src/proxy/server.rs

@@ -1,7 +1,7 @@
 use super::logger::Parser;
 use super::state::State;
 use super::{Location, Upstream};
-use crate::cache::{convert_headers, HttpResponse};
+use crate::cache;
 use crate::config::{LocationConf, PingapConf, UpstreamConf};
 use crate::utils;
 use async_trait::async_trait;
@@ -272,23 +272,18 @@ impl Server {
         }
         Ok(ServerServices { lb, bg_services })
     }
-    fn get_stats_response(&self) -> HttpResponse {
+    fn get_stats_response(&self) -> cache::HttpResponse {
         let buf = serde_json::to_vec(&ServerStats {
             accepted: self.accepted.load(Ordering::Relaxed),
             processing: self.processing.load(Ordering::Relaxed),
             hostname: HOST_NAME.to_string(),
         })
         .unwrap_or_default();
-        let headers = convert_headers(&[
-            "Content-Type: application/json; charset=utf-8".to_string(),
-            "Cache-Control: private, no-store".to_string(),
-        ])
-        .unwrap_or_default();
 
-        HttpResponse {
+        cache::HttpResponse {
             status: StatusCode::OK,
             body: Bytes::from(buf),
-            headers: Some(headers),
+            headers: Some(vec![cache::HTTP_HEADER_CONTENT_JSON.clone()]),
             ..Default::default()
         }
     }