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

8 files changed, 74 insertions, 52 deletions

diff --git a/.bleep b/.bleep
index d1591c7..22e32a6 100644
--- a/.bleep
+++ b/.bleep
@@ -1 +1 @@
-89319e9383d6f99066dfeace750a553d45e1c167
+7dbf3a97f9e59d8ad1d8e5d199cae6ee49869b9c
\ No newline at end of file
diff --git a/pingora-core/src/protocols/http/server.rs b/pingora-core/src/protocols/http/server.rs
index 74aada6..49423e9 100644
--- a/pingora-core/src/protocols/http/server.rs
+++ b/pingora-core/src/protocols/http/server.rs
@@ -22,7 +22,6 @@ use crate::protocols::{Digest, SocketAddr, Stream};
 use bytes::Bytes;
 use http::HeaderValue;
 use http::{header::AsHeaderName, HeaderMap};
-use log::error;
 use pingora_error::Result;
 use pingora_http::{RequestHeader, ResponseHeader};
 use std::time::Duration;
@@ -294,10 +293,23 @@ impl Session {
         }
     }
 
-    /// Send error response to client
-    pub async fn respond_error(&mut self, error: u16) {
-        let resp = Self::generate_error(error);
+    /// Send error response to client using a pre-generated error message.
+    pub async fn respond_error(&mut self, error: u16) -> Result<()> {
+        self.respond_error_with_body(error, Bytes::default()).await
+    }
 
+    /// Send error response to client using a pre-generated error message and custom body.
+    pub async fn respond_error_with_body(&mut self, error: u16, body: Bytes) -> Result<()> {
+        let mut resp = Self::generate_error(error);
+        if !body.is_empty() {
+            // error responses have a default content-length of zero
+            resp.set_content_length(body.len())?
+        }
+        self.write_error_response(resp, body).await
+    }
+
+    /// Send an error response to a client with a response header and body.
+    pub async fn write_error_response(&mut self, resp: ResponseHeader, body: Bytes) -> Result<()> {
         // TODO: we shouldn't be closing downstream connections on internally generated errors
         // and possibly other upstream connect() errors (connection refused, timeout, etc)
         //
@@ -306,11 +318,12 @@ impl Session {
         // rather than a misleading the client with 'keep-alive'
         self.set_keepalive(None);
 
-        self.write_response_header(Box::new(resp))
-            .await
-            .unwrap_or_else(|e| {
-                error!("failed to send error response to downstream: {e}");
-            });
+        self.write_response_header(Box::new(resp)).await?;
+        if !body.is_empty() {
+            self.write_response_body(body, true).await?;
+            self.finish_body().await?;
+        }
+        Ok(())
     }
 
     /// Whether there is no request body
diff --git a/pingora-core/src/protocols/http/v1/server.rs b/pingora-core/src/protocols/http/v1/server.rs
index 9a4bac5..4d8a21f 100644
--- a/pingora-core/src/protocols/http/v1/server.rs
+++ b/pingora-core/src/protocols/http/v1/server.rs
@@ -18,7 +18,7 @@ use bytes::Bytes;
 use bytes::{BufMut, BytesMut};
 use http::HeaderValue;
 use http::{header, header::AsHeaderName, Method, Version};
-use log::{debug, error, warn};
+use log::{debug, warn};
 use once_cell::sync::Lazy;
 use percent_encoding::{percent_encode, AsciiSet, CONTROLS};
 use pingora_error::{Error, ErrorType::*, OrErr, Result};
@@ -30,7 +30,7 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
 use super::body::{BodyReader, BodyWriter};
 use super::common::*;
-use crate::protocols::http::{body_buffer::FixedBuffer, date, error_resp, HttpTask};
+use crate::protocols::http::{body_buffer::FixedBuffer, date, HttpTask};
 use crate::protocols::{Digest, SocketAddr, Stream};
 use crate::utils::{BufRef, KVRef};
 
@@ -895,31 +895,6 @@ impl HttpSession {
         }
     }
 
-    /// Return a error response to the client. This default error response comes with `cache-control: private, no-store`.
-    /// It has no response body.
-    pub async fn respond_error(&mut self, error_status_code: u16) {
-        let (resp, resp_tmp) = match error_status_code {
-            /* commmon error responses are pre-generated */
-            502 => (Some(&*error_resp::HTTP_502_RESPONSE), None),
-            400 => (Some(&*error_resp::HTTP_400_RESPONSE), None),
-            _ => (
-                None,
-                Some(error_resp::gen_error_response(error_status_code)),
-            ),
-        };
-
-        let resp = match resp {
-            Some(r) => r,
-            None => resp_tmp.as_ref().unwrap(),
-        };
-
-        self.write_response_header_ref(resp)
-            .await
-            .unwrap_or_else(|e| {
-                error!("failed to send error response to downstream: {}", e);
-            });
-    }
-
     /// Write a `100 Continue` response to the client.
     pub async fn write_continue_response(&mut self) -> Result<()> {
         // only send if we haven't already
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()
+        );
+    }
 }
diff --git a/pingora-proxy/examples/gateway.rs b/pingora-proxy/examples/gateway.rs
index 78f4aae..d4a5cee 100644
--- a/pingora-proxy/examples/gateway.rs
+++ b/pingora-proxy/examples/gateway.rs
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use async_trait::async_trait;
+use bytes::Bytes;
 use clap::Parser;
 use log::info;
 use prometheus::register_int_counter;
@@ -42,7 +43,9 @@ impl ProxyHttp for MyGateway {
         if session.req_header().uri.path().starts_with("/login")
             && !check_login(session.req_header())
         {
-            let _ = session.respond_error(403).await;
+            let _ = session
+                .respond_error_with_body(403, Bytes::from_static(b"no way!"))
+                .await;
             // true: early return as the response is already written
             return Ok(true);
         }
@@ -103,7 +106,7 @@ impl ProxyHttp for MyGateway {
     }
 }
 
-// RUST_LOG=INFO cargo run --example load_balancer
+// RUST_LOG=INFO cargo run --example gateway
 // curl 127.0.0.1:6191 -H "Host: one.one.one.one"
 // curl 127.0.0.1:6190/family/ -H "Host: one.one.one.one"
 // curl 127.0.0.1:6191/login/ -H "Host: one.one.one.one" -I -H "Authorization: password"
diff --git a/pingora-proxy/src/lib.rs b/pingora-proxy/src/lib.rs
index 731cdb4..fa639f2 100644
--- a/pingora-proxy/src/lib.rs
+++ b/pingora-proxy/src/lib.rs
@@ -143,9 +143,14 @@ impl<SV> HttpProxy<SV> {
             }
             Err(mut e) => {
                 e.as_down();
-                error!("Fail to proxy: {}", e);
+                error!("Fail to proxy: {e}");
                 if matches!(e.etype, InvalidHTTPHeader) {
-                    downstream_session.respond_error(400).await;
+                    downstream_session
+                        .respond_error(400)
+                        .await
+                        .unwrap_or_else(|e| {
+                            error!("failed to send error response to downstream: {e}");
+                        });
                 } // otherwise the connection must be broken, no need to send anything
                 downstream_session.shutdown().await;
                 return None;
@@ -344,16 +349,16 @@ impl Session {
         &self.downstream_session
     }
 
-    /// Write HTTP response with the given error code to the downstream
+    /// Write HTTP response with the given error code to the downstream.
     pub async fn respond_error(&mut self, error: u16) -> Result<()> {
-        let resp = HttpSession::generate_error(error);
-        self.write_response_header(Box::new(resp), true)
+        self.as_downstream_mut().respond_error(error).await
+    }
+
+    /// Write HTTP response with the given error code to the downstream with a body.
+    pub async fn respond_error_with_body(&mut self, error: u16, body: Bytes) -> Result<()> {
+        self.as_downstream_mut()
+            .respond_error_with_body(error, body)
             .await
-            .unwrap_or_else(|e| {
-                self.downstream_session.set_keepalive(None);
-                error!("failed to send error response to downstream: {e}");
-            });
-        Ok(())
     }
 
     /// Write the given HTTP response header to the downstream
diff --git a/pingora-proxy/src/proxy_cache.rs b/pingora-proxy/src/proxy_cache.rs
index 8957bc7..c720575 100644
--- a/pingora-proxy/src/proxy_cache.rs
+++ b/pingora-proxy/src/proxy_cache.rs
@@ -295,7 +295,13 @@ impl<SV> HttpProxy<SV> {
             }
             Err(e) => {
                 // TODO: more logging and error handling
-                session.as_mut().respond_error(500).await;
+                session
+                    .as_mut()
+                    .respond_error(500)
+                    .await
+                    .unwrap_or_else(|e| {
+                        error!("failed to send error response to downstream: {e}");
+                    });
                 // we have not write anything dirty to downstream, it is still reusable
                 return (true, Some(e));
             }
diff --git a/pingora-proxy/src/proxy_trait.rs b/pingora-proxy/src/proxy_trait.rs
index 4017789..6970ec0 100644
--- a/pingora-proxy/src/proxy_trait.rs
+++ b/pingora-proxy/src/proxy_trait.rs
@@ -372,7 +372,6 @@ pub trait ProxyHttp {
     where
         Self::CTX: Send + Sync,
     {
-        let server_session = session.as_mut();
         let code = match e.etype() {
             HTTPStatus(code) => *code,
             _ => {
@@ -392,7 +391,9 @@ pub trait ProxyHttp {
             }
         };
         if code > 0 {
-            server_session.respond_error(code).await
+            session.respond_error(code).await.unwrap_or_else(|e| {
+                error!("failed to send error response to downstream: {e}");
+            });
         }
         code
     }