From a8a6e77eef2c0f4d2a45f00c5b0e316dd373f2f2 Mon Sep 17 00:00:00 2001
From: Andrew Hauck <ahauck@cloudflare.com>
Date: Fri, 8 Nov 2024 14:34:14 -0800
Subject: Improve support for sending custom response headers and bodies for
 error messages

---
 pingora-http/src/lib.rs | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

(limited to 'pingora-http')

diff --git a/pingora-http/src/lib.rs b/pingora-http/src/lib.rs
index d57998d..d5cf8a8 100644
--- a/pingora-http/src/lib.rs
+++ b/pingora-http/src/lib.rs
@@ -492,6 +492,11 @@ impl ResponseHeader {
     pub fn as_owned_parts(&self) -> RespParts {
         clone_resp_parts(&self.base)
     }
+
+    /// Helper function to set the HTTP content length on the response header.
+    pub fn set_content_length(&mut self, len: usize) -> Result<()> {
+        self.insert_header(http::header::CONTENT_LENGTH, len)
+    }
 }
 
 fn clone_req_parts(me: &ReqParts) -> ReqParts {
@@ -754,4 +759,18 @@ mod tests {
         // Some(false)
         assert!(!req.send_end_stream().unwrap());
     }
+
+    #[test]
+    fn set_test_set_content_length() {
+        let mut resp = ResponseHeader::new(None);
+        resp.set_content_length(10).unwrap();
+
+        assert_eq!(
+            b"10",
+            resp.headers
+                .get(http::header::CONTENT_LENGTH)
+                .map(|d| d.as_bytes())
+                .unwrap()
+        );
+    }
 }
-- 
cgit v1.2.3