Add support for downstream h2 trailers and add an upstream h2 response trailer filter

4 files changed, 35 insertions, 12 deletions

diff --git a/pingora-proxy/src/lib.rs b/pingora-proxy/src/lib.rs
index 930087d..453e9dd 100644
--- a/pingora-proxy/src/lib.rs
+++ b/pingora-proxy/src/lib.rs
@@ -222,7 +222,12 @@ impl<SV> HttpProxy<SV> {
         }
     }
 
-    fn upstream_filter(&self, session: &mut Session, task: &mut HttpTask, ctx: &mut SV::CTX)
+    fn upstream_filter(
+        &self,
+        session: &mut Session,
+        task: &mut HttpTask,
+        ctx: &mut SV::CTX,
+    ) -> Result<()>
     where
         SV: ProxyHttp,
     {
@@ -233,10 +238,14 @@ impl<SV> HttpProxy<SV> {
             HttpTask::Body(data, eos) => self
                 .inner
                 .upstream_response_body_filter(session, data, *eos, ctx),
+            HttpTask::Trailer(Some(trailers)) => self
+                .inner
+                .upstream_response_trailer_filter(session, trailers, ctx)?,
             _ => {
-                // TODO: add other upstream filter traits
+                // task does not support a filter
             }
         }
+        Ok(())
     }
 
     async fn finish(
diff --git a/pingora-proxy/src/proxy_h1.rs b/pingora-proxy/src/proxy_h1.rs
index 9856921..c702af3 100644
--- a/pingora-proxy/src/proxy_h1.rs
+++ b/pingora-proxy/src/proxy_h1.rs
@@ -431,7 +431,7 @@ impl<SV> HttpProxy<SV> {
     {
         // skip caching if already served from cache
         if !from_cache {
-            self.upstream_filter(session, &mut task, ctx);
+            self.upstream_filter(session, &mut task, ctx)?;
 
             // cache the original response before any downstream transformation
             // requests that bypassed cache still need to run filters to see if the response has become cacheable
@@ -515,7 +515,7 @@ impl<SV> HttpProxy<SV> {
                 }
                 Ok(HttpTask::Body(data, end))
             }
-            HttpTask::Trailer(h) => Ok(HttpTask::Trailer(h)), // no h1 trailer filter yet
+            HttpTask::Trailer(h) => Ok(HttpTask::Trailer(h)), // TODO: support trailers for h1
             HttpTask::Done => Ok(task),
             HttpTask::Failed(_) => Ok(task), // Do nothing just pass the error down
         }
diff --git a/pingora-proxy/src/proxy_h2.rs b/pingora-proxy/src/proxy_h2.rs
index 43b82ef..8f0dffa 100644
--- a/pingora-proxy/src/proxy_h2.rs
+++ b/pingora-proxy/src/proxy_h2.rs
@@ -381,7 +381,7 @@ impl<SV> HttpProxy<SV> {
         SV::CTX: Send + Sync,
     {
         if !from_cache {
-            self.upstream_filter(session, &mut task, ctx);
+            self.upstream_filter(session, &mut task, ctx)?;
 
             // cache the original response before any downstream transformation
             // requests that bypassed cache still need to run filters to see if the response has become cacheable
@@ -457,18 +457,18 @@ impl<SV> HttpProxy<SV> {
                     .inner
                     .response_body_filter(session, &mut data, eos, ctx)?
                 {
-                    trace!("delaying response for {:?}", duration);
+                    trace!("delaying response for {duration:?}");
                     time::sleep(duration).await;
                 }
                 Ok(HttpTask::Body(data, eos))
             }
-            HttpTask::Trailer(header_map) => {
-                let trailer_buffer = match header_map {
-                    Some(mut trailer_map) => {
+            HttpTask::Trailer(mut trailers) => {
+                let trailer_buffer = match trailers.as_mut() {
+                    Some(trailers) => {
                         debug!("Parsing response trailers..");
                         match self
                             .inner
-                            .response_trailer_filter(session, &mut trailer_map, ctx)
+                            .response_trailer_filter(session, trailers, ctx)
                             .await
                         {
                             Ok(buf) => buf,
@@ -490,7 +490,7 @@ impl<SV> HttpProxy<SV> {
                     // https://http2.github.io/http2-spec/#malformed
                     Ok(HttpTask::Body(Some(buffer), true))
                 } else {
-                    Ok(HttpTask::Done)
+                    Ok(HttpTask::Trailer(trailers))
                 }
             }
             HttpTask::Done => Ok(task),
diff --git a/pingora-proxy/src/proxy_trait.rs b/pingora-proxy/src/proxy_trait.rs
index 0ee76b1..888854c 100644
--- a/pingora-proxy/src/proxy_trait.rs
+++ b/pingora-proxy/src/proxy_trait.rs
@@ -221,6 +221,16 @@ pub trait ProxyHttp {
     ) {
     }
 
+    /// Similar to [Self::upstream_response_filter()] but for response trailers
+    fn upstream_response_trailer_filter(
+        &self,
+        _session: &mut Session,
+        _upstream_trailers: &mut header::HeaderMap,
+        _ctx: &mut Self::CTX,
+    ) -> Result<()> {
+        Ok(())
+    }
+
     /// Similar to [Self::response_filter()] but for response body chunks
     fn response_body_filter(
         &self,
@@ -235,7 +245,11 @@ pub trait ProxyHttp {
         Ok(None)
     }
 
-    /// When a trailer is received.
+    /// Similar to [Self::response_filter()] but for response trailers.
+    /// Note, returning an Ok(Some(Bytes)) will result in the downstream response
+    /// trailers being written to the response body.
+    ///
+    /// TODO: make this interface more intuitive
     async fn response_trailer_filter(
         &self,
         _session: &mut Session,