Improve support for sending custom response headers and bodies for error messages

1 files changed, 19 insertions, 0 deletions

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()
+        );
+    }
 }