From a5c03fa60ac0a5b814ea9082474dd17667a4571b Mon Sep 17 00:00:00 2001
From: Matthew Gumport <mgumport@cloudflare.com>
Date: Fri, 13 Dec 2024 15:07:12 -0800
Subject: add range_header_filter to proxy trait

This makes the `range_header_filter` overridable. The new default implementation
is what we were doing instead of calling a hook.
---
 .bleep                           |  2 +-
 pingora-proxy/src/lib.rs         |  1 +
 pingora-proxy/src/proxy_cache.rs |  4 ++--
 pingora-proxy/src/proxy_h1.rs    |  7 +++----
 pingora-proxy/src/proxy_h2.rs    |  3 +--
 pingora-proxy/src/proxy_trait.rs | 18 ++++++++++++++++++
 6 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/.bleep b/.bleep
index 265906b..2984400 100644
--- a/.bleep
+++ b/.bleep
@@ -1 +1 @@
-0875924524a7338c15784029df5306dd08e981dd
\ No newline at end of file
+9df2f3f6e6b919a632b08af4584a1c1a3bcee0fd
\ No newline at end of file
diff --git a/pingora-proxy/src/lib.rs b/pingora-proxy/src/lib.rs
index bf4d4ec..28608cd 100644
--- a/pingora-proxy/src/lib.rs
+++ b/pingora-proxy/src/lib.rs
@@ -78,6 +78,7 @@ pub mod subrequest;
 
 use subrequest::Ctx as SubReqCtx;
 
+pub use proxy_cache::range_filter::{range_header_filter, RangeType};
 pub use proxy_purge::PurgeStatus;
 pub use proxy_trait::ProxyHttp;
 
diff --git a/pingora-proxy/src/proxy_cache.rs b/pingora-proxy/src/proxy_cache.rs
index 454416e..b393780 100644
--- a/pingora-proxy/src/proxy_cache.rs
+++ b/pingora-proxy/src/proxy_cache.rs
@@ -279,7 +279,7 @@ impl<SV> HttpProxy<SV> {
 
         // process range header if the cache storage supports seek
         let range_type = if seekable && !session.ignore_downstream_range {
-            range_header_filter(req, &mut header)
+            self.inner.range_header_filter(req, &mut header, ctx)
         } else {
             RangeType::None
         };
@@ -826,7 +826,7 @@ fn cache_hit_header(cache: &HttpCache) -> Box<ResponseHeader> {
 }
 
 // https://datatracker.ietf.org/doc/html/rfc7233#section-3
-pub(crate) mod range_filter {
+pub mod range_filter {
     use super::*;
     use http::header::*;
     use std::ops::Range;
diff --git a/pingora-proxy/src/proxy_h1.rs b/pingora-proxy/src/proxy_h1.rs
index d323562..580b070 100644
--- a/pingora-proxy/src/proxy_h1.rs
+++ b/pingora-proxy/src/proxy_h1.rs
@@ -492,10 +492,9 @@ impl<SV> HttpProxy<SV> {
                         ctx,
                     );
                     if !session.ignore_downstream_range {
-                        let range_type = proxy_cache::range_filter::range_header_filter(
-                            session.req_header(),
-                            &mut header,
-                        );
+                        let range_type =
+                            self.inner
+                                .range_header_filter(session.req_header(), &mut header, ctx);
                         range_body_filter.set(range_type);
                     }
                 }
diff --git a/pingora-proxy/src/proxy_h2.rs b/pingora-proxy/src/proxy_h2.rs
index a27c559..ca10ef5 100644
--- a/pingora-proxy/src/proxy_h2.rs
+++ b/pingora-proxy/src/proxy_h2.rs
@@ -445,8 +445,7 @@ impl<SV> HttpProxy<SV> {
                         ctx,
                     );
                     if !session.ignore_downstream_range {
-                        let range_type =
-                            proxy_cache::range_filter::range_header_filter(req, &mut header);
+                        let range_type = self.inner.range_header_filter(req, &mut header, ctx);
                         range_body_filter.set(range_type);
                     }
                 }
diff --git a/pingora-proxy/src/proxy_trait.rs b/pingora-proxy/src/proxy_trait.rs
index 00e54f4..e67a9f0 100644
--- a/pingora-proxy/src/proxy_trait.rs
+++ b/pingora-proxy/src/proxy_trait.rs
@@ -18,6 +18,7 @@ use pingora_cache::{
     CacheKey, CacheMeta, ForcedInvalidationKind,
     RespCacheable::{self, *},
 };
+use proxy_cache::range_filter::{self};
 use std::time::Duration;
 
 /// The interface to control the HTTP proxy
@@ -216,6 +217,23 @@ pub trait ProxyHttp {
         )
     }
 
+    /// This filter is called when cache is enabled to determine what byte range to return (in both
+    /// cache hit and miss cases) from the response body. It is only used when caching is enabled,
+    /// otherwise the upstream is responsible for any filtering. It allows users to define the range
+    /// this request is for via its return type `range_filter::RangeType`.
+    ///
+    /// It also allow users to modify the response header accordingly.
+    ///
+    /// The default implementation can handle a single-range as per [RFC7232].
+    fn range_header_filter(
+        &self,
+        req: &RequestHeader,
+        resp: &mut ResponseHeader,
+        _ctx: &mut Self::CTX,
+    ) -> range_filter::RangeType {
+        proxy_cache::range_filter::range_header_filter(req, resp)
+    }
+
     /// Modify the request before it is sent to the upstream
     ///
     /// Unlike [Self::request_filter()], this filter allows to change the request headers to send
-- 
cgit v1.2.3