Handle UpgradedBody in more filters

drcaramelsyrup · drcaramelsyrup · commit aca00b609524 · 2026-02-02T22:35:59.000-08:00
UpgradedBody was incorrectly no longer invoked for upstream response
body filters or for downstream modules. Additionally allow supporting
the rare case where UpgradedBody may be cached.
diff --git a/.bleep b/.bleep
@@ -1 +1 @@
-a83d44c1042cc8071f1361441e606e50a22c68f4
+71f26703aeb326cc03ccf2d200a1784c915ffb49
diff --git a/pingora-proxy/src/lib.rs b/pingora-proxy/src/lib.rs
@@ -357,7 +357,7 @@ where
                     .await?;
                 None
             }
-            HttpTask::Body(data, eos) => self
+            HttpTask::Body(data, eos) | HttpTask::UpgradedBody(data, eos) => self
                 .inner
                 .upstream_response_body_filter(session, data, *eos, ctx)?,
             HttpTask::Trailer(Some(trailers)) => {
@@ -563,6 +563,7 @@ impl Session {
     }
 
     pub async fn write_response_tasks(&mut self, mut tasks: Vec<HttpTask>) -> Result<bool> {
+        let mut seen_upgraded = self.was_upgraded();
         for task in tasks.iter_mut() {
             match task {
                 HttpTask::Header(resp, end) => {
@@ -574,6 +575,11 @@ impl Session {
                     self.downstream_modules_ctx
                         .response_body_filter(data, *end)?;
                 }
+                HttpTask::UpgradedBody(data, end) => {
+                    seen_upgraded = true;
+                    self.downstream_modules_ctx
+                        .response_body_filter(data, *end)?;
+                }
                 HttpTask::Trailer(trailers) => {
                     if let Some(buf) = self
                         .downstream_modules_ctx
@@ -584,6 +590,7 @@ impl Session {
                         //
                         // Note, this will not work if end of stream has already
                         // been seen or we've written content-length bytes.
+                        // (Trailers should never come after upgraded body)
                         *task = HttpTask::Body(Some(buf), true);
                     }
                 }
@@ -597,7 +604,11 @@ impl Session {
                     // Note, this will not work if end of stream has already
                     // been seen or we've written content-length bytes.
                     if let Some(buf) = self.downstream_modules_ctx.response_done_filter()? {
-                        *task = HttpTask::Body(Some(buf), true);
+                        if seen_upgraded {
+                            *task = HttpTask::UpgradedBody(Some(buf), true);
+                        } else {
+                            *task = HttpTask::Body(Some(buf), true);
+                        }
                     }
                 }
                 _ => { /* Failed */ }
diff --git a/pingora-proxy/src/proxy_cache.rs b/pingora-proxy/src/proxy_cache.rs
@@ -661,52 +661,49 @@ where
                     }
                 }
             }
-            HttpTask::Body(data, end_stream) => match data {
-                Some(d) => {
-                    if session.cache.enabled() {
-                        // TODO: do this async
-                        // fail if writing the body would exceed the max_file_size_bytes
-                        let body_size_allowed =
-                            session.cache.track_body_bytes_for_max_file_size(d.len());
-                        if !body_size_allowed {
-                            debug!("chunked response exceeded max cache size, remembering that it is uncacheable");
-                            session
-                                .cache
-                                .response_became_uncacheable(NoCacheReason::ResponseTooLarge);
-
-                            return Error::e_explain(
-                                ERR_RESPONSE_TOO_LARGE,
-                                format!(
-                                    "writing data of size {} bytes would exceed max file size of {} bytes",
-                                    d.len(),
-                                    session.cache.max_file_size_bytes().expect("max file size bytes must be set to exceed size")
-                                ),
-                            );
-                        }
+            HttpTask::Body(data, end_stream) | HttpTask::UpgradedBody(data, end_stream) => {
+                // It is not normally advisable to cache upgraded responses
+                // e.g. they are essentially close-delimited, so they are easily truncated
+                // but the framework still allows for it
+                match data {
+                    Some(d) => {
+                        if session.cache.enabled() {
+                            // TODO: do this async
+                            // fail if writing the body would exceed the max_file_size_bytes
+                            let body_size_allowed =
+                                session.cache.track_body_bytes_for_max_file_size(d.len());
+                            if !body_size_allowed {
+                                debug!("chunked response exceeded max cache size, remembering that it is uncacheable");
+                                session
+                                    .cache
+                                    .response_became_uncacheable(NoCacheReason::ResponseTooLarge);
+
+                                return Error::e_explain(
+                                    ERR_RESPONSE_TOO_LARGE,
+                                    format!(
+                                        "writing data of size {} bytes would exceed max file size of {} bytes",
+                                        d.len(),
+                                        session.cache.max_file_size_bytes().expect("max file size bytes must be set to exceed size")
+                                    ),
+                                );
+                            }
 
-                        // this will panic if more data is sent after we see end_stream
-                        // but should be impossible in real world
-                        let miss_handler = session.cache.miss_handler().unwrap();
+                            // this will panic if more data is sent after we see end_stream
+                            // but should be impossible in real world
+                            let miss_handler = session.cache.miss_handler().unwrap();
 
-                        miss_handler.write_body(d.clone(), *end_stream).await?;
-                        if *end_stream {
-                            session.cache.finish_miss_handler().await?;
+                            miss_handler.write_body(d.clone(), *end_stream).await?;
+                            if *end_stream {
+                                session.cache.finish_miss_handler().await?;
+                            }
                         }
                     }
-                }
-                None => {
-                    if session.cache.enabled() && *end_stream {
-                        session.cache.finish_miss_handler().await?;
+                    None => {
+                        if session.cache.enabled() && *end_stream {
+                            session.cache.finish_miss_handler().await?;
+                        }
                     }
                 }
-            },
-            HttpTask::UpgradedBody(..) => {
-                // caching upgraded bodies isn't supported with and doesn't make sense with the HttpCache
-                // (caller of cache http task will disable cache in the session)
-                return Error::e_explain(
-                    InternalError,
-                    "Unexpected UpgradedBody task while caching",
-                );
             }
             HttpTask::Trailer(_) => {} // h1 trailer is not supported yet
             HttpTask::Done => {
@@ -2285,7 +2282,16 @@ impl ServeFromCache {
         &mut self,
         cache: &mut HttpCache,
         range: &mut RangeBodyFilter,
+        upgraded: bool,
     ) -> Result<HttpTask> {
+        fn body_task(data: Bytes, upgraded: bool) -> HttpTask {
+            if upgraded {
+                HttpTask::UpgradedBody(Some(data), false)
+            } else {
+                HttpTask::Body(Some(data), false)
+            }
+        }
+
         if !cache.enabled() {
             // Cache is disabled due to internal error
             // TODO: if nothing is sent to eyeball yet, figure out a way to recovery by
@@ -2317,7 +2323,7 @@ impl ServeFromCache {
                 }
                 loop {
                     if let Some(b) = cache.hit_handler().read_body().await? {
-                        return Ok(HttpTask::Body(Some(b), false)); // false for now
+                        return Ok(body_task(b, upgraded));
                     }
                     // EOF from hit handler for body requested
                     // if multipart, then seek again
@@ -2336,7 +2342,7 @@ impl ServeFromCache {
                 // safety: caller of enable_miss() call it only if the async_body_reader exist
                 loop {
                     if let Some(b) = cache.miss_body_reader().unwrap().read_body().await? {
-                        return Ok(HttpTask::Body(Some(b), false)); // false for now
+                        return Ok(body_task(b, upgraded));
                     } else {
                         // EOF from hit handler for body requested
                         // if multipart, then seek again
diff --git a/pingora-proxy/src/proxy_custom.rs b/pingora-proxy/src/proxy_custom.rs
@@ -317,6 +317,7 @@ where
             // partial read support, this check will also be false if cache is disabled.
             let support_cache_partial_read =
                 session.cache.support_streaming_partial_write() == Some(true);
+            let upgraded = session.was_upgraded();
 
             tokio::select! {
                 body = session.downstream_session.read_body_or_idle(downstream_state.is_done()), if downstream_state.can_poll() => {
@@ -423,7 +424,7 @@ where
                     }
                 }
 
-                task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter),
+                task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter, upgraded),
                     if !response_state.cached_done() && !downstream_state.is_errored() && serve_from_cache.is_on() => {
                     let task = self.custom_response_filter(session, task?, ctx,
                         &mut serve_from_cache,
diff --git a/pingora-proxy/src/proxy_h1.rs b/pingora-proxy/src/proxy_h1.rs
@@ -370,6 +370,7 @@ where
             // partial read support, this check will also be false if cache is disabled.
             let support_cache_partial_read =
                 session.cache.support_streaming_partial_write() == Some(true);
+            let upgraded = session.was_upgraded();
 
             tokio::select! {
                 // only try to send to pipe if there is capacity to avoid deadlock
@@ -500,7 +501,7 @@ where
                     }
                 },
 
-                task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter),
+                task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter, upgraded),
                     if !response_state.cached_done() && !downstream_state.is_errored() && serve_from_cache.is_on() => {
 
                     let task = self.h1_response_filter(session, task?, ctx,
@@ -720,7 +721,13 @@ where
                 Ok(HttpTask::Body(data, end))
             }
             HttpTask::UpgradedBody(mut data, end) => {
-                // range / caching doesn't apply to upgraded body
+                if track_max_cache_size {
+                    session
+                        .cache
+                        .track_body_bytes_for_max_file_size(data.as_ref().map_or(0, |d| d.len()));
+                }
+
+                // range doesn't apply to upgraded body
                 if let Some(duration) = self
                     .inner
                     .response_body_filter(session, &mut data, end, ctx)?
diff --git a/pingora-proxy/src/proxy_h2.rs b/pingora-proxy/src/proxy_h2.rs
@@ -340,6 +340,7 @@ where
             // partial read support, this check will also be false if cache is disabled.
             let support_cache_partial_read =
                 session.cache.support_streaming_partial_write() == Some(true);
+            let upgraded = session.was_upgraded();
 
             // Similar logic in h1 need to reserve capacity first to avoid deadlock
             // But we don't need to do the same because the h2 client_body pipe is unbounded (never block)
@@ -447,7 +448,7 @@ where
                     }
                 }
 
-                task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter),
+                task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter, upgraded),
                     if !response_state.cached_done() && !downstream_state.is_errored() && serve_from_cache.is_on() => {
                     let task = self.h2_response_filter(session, task?, ctx,
                         &mut serve_from_cache,
diff --git a/pingora-proxy/tests/test_upstream.rs b/pingora-proxy/tests/test_upstream.rs
@@ -1138,6 +1138,55 @@ mod test_cache {
         assert_eq!(res.text().await.unwrap(), "");
     }
 
+    #[tokio::test]
+    async fn test_cache_websocket_101() {
+        // Test the unlikely scenario in which users may want to cache WS
+        init();
+
+        // First request - should be a miss
+        let mut stream = TcpStream::connect("127.0.0.1:6148").await.unwrap();
+        let req = concat!(
+            "GET /unique/test_cache_websocket_101/upgrade HTTP/1.1\r\n",
+            "Host: 127.0.0.1\r\n",
+            "Upgrade: websocket\r\n",
+            "Connection: Upgrade\r\n",
+            "X-Cache-Websocket: 1\r\n",
+            "\r\n"
+        );
+        stream.write_all(req.as_bytes()).await.unwrap();
+        stream.flush().await.unwrap();
+
+        let expected_payload = b"hello\n";
+        let fut = read_response(&mut stream, expected_payload.len());
+        let (resp_header, resp_body) = timeout(Duration::from_secs(5), fut).await.unwrap();
+
+        assert_eq!(resp_header.status, 101);
+        assert_eq!(resp_header.headers["Upgrade"], "websocket");
+        assert_eq!(resp_header.headers["x-cache-status"], "miss");
+        assert_eq!(resp_body, expected_payload);
+
+        // Second request - should be a cache hit
+        let mut stream = TcpStream::connect("127.0.0.1:6148").await.unwrap();
+        let req = concat!(
+            "GET /unique/test_cache_websocket_101/upgrade HTTP/1.1\r\n",
+            "Host: 127.0.0.1\r\n",
+            "Upgrade: websocket\r\n",
+            "Connection: Upgrade\r\n",
+            "X-Cache-Websocket: 1\r\n",
+            "\r\n"
+        );
+        stream.write_all(req.as_bytes()).await.unwrap();
+        stream.flush().await.unwrap();
+
+        let fut = read_response(&mut stream, expected_payload.len());
+        let (resp_header, resp_body) = timeout(Duration::from_secs(5), fut).await.unwrap();
+
+        assert_eq!(resp_header.status, 101);
+        assert_eq!(resp_header.headers["Upgrade"], "websocket");
+        assert_eq!(resp_header.headers["x-cache-status"], "hit");
+        assert_eq!(resp_body, expected_payload);
+    }
+
     #[tokio::test]
     async fn test_1xx_caching() {
         // 1xx shouldn't interfere with HTTP caching
diff --git a/pingora-proxy/tests/utils/server_utils.rs b/pingora-proxy/tests/utils/server_utils.rs
@@ -47,7 +47,7 @@ use pingora_proxy::{FailToProxy, ProxyHttp, Session};
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 use std::thread;
-use std::time::Duration;
+use std::time::{Duration, SystemTime};
 
 pub struct ExampleProxyHttps {}
 
@@ -574,10 +574,26 @@ impl ProxyHttp for ExampleProxyCache {
 
     fn response_cache_filter(
         &self,
-        _session: &Session,
+        session: &Session,
         resp: &ResponseHeader,
         _ctx: &mut Self::CTX,
     ) -> Result<RespCacheable> {
+        // Allow testing the unlikely case of caching a 101 response
+        if resp.status == 101
+            && session
+                .req_header()
+                .headers
+                .contains_key("x-cache-websocket")
+        {
+            return Ok(RespCacheable::Cacheable(CacheMeta::new(
+                SystemTime::now() + Duration::from_secs(5),
+                SystemTime::now(),
+                0,
+                0,
+                resp.clone(),
+            )));
+        }
+
         let cc = CacheControl::from_resp_headers(resp);
         Ok(resp_cacheable(
             cc.as_ref(),

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-a83d44c1042cc8071f1361441e606e50a22c68f4`
	`1`	`+71f26703aeb326cc03ccf2d200a1784c915ffb49`
Original file line number	Diff line number	Diff line change
`@@ -317,6 +317,7 @@ where`
`317`	`317`	`// partial read support, this check will also be false if cache is disabled.`
`318`	`318`	`let support_cache_partial_read =`
`319`	`319`	`session.cache.support_streaming_partial_write() == Some(true);`
	`320`	`+ let upgraded = session.was_upgraded();`
`320`	`321`
`321`	`322`	`tokio::select! {`
`322`	`323`	`body = session.downstream_session.read_body_or_idle(downstream_state.is_done()), if downstream_state.can_poll() => {`
`@@ -423,7 +424,7 @@ where`
`423`	`424`	`}`
`424`	`425`	`}`
`425`	`426`
`426`		`- task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter),`
	`427`	`+ task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter, upgraded),`
`427`	`428`	`if !response_state.cached_done() && !downstream_state.is_errored() && serve_from_cache.is_on() => {`
`428`	`429`	`let task = self.custom_response_filter(session, task?, ctx,`
`429`	`430`	`&mut serve_from_cache,`
Original file line number	Diff line number	Diff line change
`@@ -340,6 +340,7 @@ where`
`340`	`340`	`// partial read support, this check will also be false if cache is disabled.`
`341`	`341`	`let support_cache_partial_read =`
`342`	`342`	`session.cache.support_streaming_partial_write() == Some(true);`
	`343`	`+ let upgraded = session.was_upgraded();`
`343`	`344`
`344`	`345`	`// Similar logic in h1 need to reserve capacity first to avoid deadlock`
`345`	`346`	`// But we don't need to do the same because the h2 client_body pipe is unbounded (never block)`
`@@ -447,7 +448,7 @@ where`
`447`	`448`	`}`
`448`	`449`	`}`
`449`	`450`
`450`		`- task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter),`
	`451`	`+ task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter, upgraded),`
`451`	`452`	`if !response_state.cached_done() && !downstream_state.is_errored() && serve_from_cache.is_on() => {`
`452`	`453`	`let task = self.h2_response_filter(session, task?, ctx,`
`453`	`454`	`&mut serve_from_cache,`