Skip to main content

pingora_proxy/
proxy_cache.rs

1// Copyright 2026 Cloudflare, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use super::*;
16use http::header::{CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE, TRANSFER_ENCODING};
17use http::{Method, StatusCode};
18use pingora_cache::key::CacheHashKey;
19use pingora_cache::lock::LockWaitOutcome;
20use pingora_cache::max_file_size::ERR_RESPONSE_TOO_LARGE;
21use pingora_cache::{ForcedFreshness, HitHandler, HitStatus, RespCacheable::*};
22use pingora_core::protocols::http::conditional_filter::to_304;
23use pingora_core::protocols::http::v1::common::header_value_content_length;
24use pingora_core::ErrorType;
25use range_filter::RangeBodyFilter;
26use std::time::SystemTime;
27
28const DEFAULT_MAX_CACHE_LOCK_RETRIES: usize = 2;
29
30impl<SV, C> HttpProxy<SV, C>
31where
32    C: custom::Connector,
33{
34    // return bool: server_session can be reused, and error if any
35    pub(crate) async fn proxy_cache(
36        self: &Arc<Self>,
37        session: &mut Session,
38        ctx: &mut SV::CTX,
39    ) -> Option<(bool, Option<Box<Error>>)>
40    // None: continue to proxy, Some: return
41    where
42        SV: ProxyHttp + Send + Sync + 'static,
43        SV::CTX: Send + Sync,
44    {
45        // Cache logic request phase
46        if let Err(e) = self.inner.request_cache_filter(session, ctx) {
47            // TODO: handle this error
48            warn!(
49                "Fail to request_cache_filter: {e}, {}",
50                self.inner.request_summary(session, ctx)
51            );
52        }
53
54        // cache key logic, should this be part of request_cache_filter?
55        if session.cache.enabled() {
56            match self.inner.cache_key_callback(session, ctx) {
57                Ok(key) => {
58                    session.cache.set_cache_key(key);
59                }
60                Err(e) => {
61                    // TODO: handle this error
62                    session.cache.disable(NoCacheReason::StorageError);
63                    warn!(
64                        "Fail to cache_key_callback: {e}, {}",
65                        self.inner.request_summary(session, ctx)
66                    );
67                }
68            }
69        }
70
71        // cache purge logic: PURGE short-circuits rest of request
72        if self.inner.is_purge(session, ctx) {
73            return self.proxy_purge(session, ctx).await;
74        }
75
76        // bypass cache lookup if we predict to be uncacheable
77        if session.cache.enabled() && !session.cache.cacheable_prediction() {
78            session.cache.bypass();
79        }
80
81        if !session.cache.enabled() {
82            return None;
83        }
84
85        // cache lookup logic
86        let mut cache_lock_retries = 0;
87        loop {
88            match session.cache.cache_lookup().await {
89                Ok(res) => {
90                    let mut hit_status_opt = None;
91                    if let Some((mut meta, mut handler)) = res {
92                        // Vary logic
93                        // Because this branch can be called multiple times in a loop, and we only
94                        // need to update the vary once, check if variance is already set to
95                        // prevent unnecessary vary lookups.
96                        let cache_key = session.cache.cache_key();
97                        if let Some(variance) = cache_key.variance_bin() {
98                            // We've looked up a secondary slot.
99                            // Adhoc double check that the variance found is the variance we want.
100                            if Some(variance) != meta.variance() {
101                                warn!("Cache variance mismatch, {variance:?}, {cache_key:?}");
102                                session.cache.disable(NoCacheReason::InternalError);
103                                break None;
104                            }
105                        } else {
106                            // Basic cache key; either variance is off, or this is the primary slot.
107                            let req_header = session.req_header();
108                            let variance = self.inner.cache_vary_filter(&meta, ctx, req_header);
109                            if let Some(variance) = variance {
110                                // Variance is on. This is the primary slot.
111                                if !session.cache.cache_vary_lookup(variance, &meta) {
112                                    // This wasn't the desired variant. Updated cache key variance, cause another
113                                    // lookup to get the desired variant, which would be in a secondary slot.
114                                    continue;
115                                }
116                            } // else: vary is not in use
117                        }
118
119                        // Either no variance, or the current handler targets the correct variant.
120
121                        // hit
122                        // TODO: maybe round and/or cache now()
123                        let now = SystemTime::now();
124                        let is_fresh = meta.is_fresh(now);
125                        // check if we should force expire or force miss
126                        let hit_status = match self
127                            .inner
128                            .cache_hit_filter(session, &meta, &mut handler, is_fresh, ctx)
129                            .await
130                        {
131                            Err(e) => {
132                                error!(
133                                    "Failed to filter cache hit: {e}, {}",
134                                    self.inner.request_summary(session, ctx)
135                                );
136                                // this return value will cause us to fetch from upstream
137                                HitStatus::FailedHitFilter
138                            }
139                            Ok(None) => {
140                                if is_fresh {
141                                    HitStatus::Fresh
142                                } else {
143                                    HitStatus::Expired
144                                }
145                            }
146                            Ok(Some(ForcedFreshness::ForceExpired)) => {
147                                // this variant exists to take data out of service, so the
148                                // stale body must not be served while it revalidates
149                                meta.disable_serve_stale();
150                                HitStatus::ForceExpired
151                            }
152                            Ok(Some(ForcedFreshness::ForceExpiredServeStale { expired_at })) => {
153                                // this variant keeps the serve stale windows, so they run
154                                // from when the asset went out of service, not its deadline
155                                if let Some(expired_at) = expired_at {
156                                    meta.expire_at(expired_at);
157                                }
158                                HitStatus::ForceExpiredServeStale
159                            }
160                            Ok(Some(ForcedFreshness::ForceMiss)) => HitStatus::ForceMiss,
161                            Ok(Some(ForcedFreshness::ForceFresh)) => HitStatus::ForceFresh,
162                        };
163
164                        hit_status_opt = Some(hit_status);
165
166                        // init cache for hit / stale
167                        session.cache.cache_found(meta, handler, hit_status);
168                    }
169
170                    if hit_status_opt.is_none_or(HitStatus::is_treated_as_miss) {
171                        // cache miss
172                        if !session.cache.enabled() {
173                            // An admission policy may have disabled caching during cache_lookup().
174                            break None;
175                        } else if session.cache.is_cache_locked() {
176                            // Another request is filling the cache; try waiting til that's done and retry.
177                            let outcome = session.cache.cache_lock_wait().await;
178                            if self.handle_lock_wait_outcome(session, ctx, outcome) {
179                                if self.cache_lock_retry_limit_exceeded(
180                                    session,
181                                    ctx,
182                                    &mut cache_lock_retries,
183                                ) {
184                                    break None;
185                                }
186                                continue;
187                            } else {
188                                break None;
189                            }
190                        } else {
191                            self.inner.cache_miss(session, ctx);
192                            break None;
193                        }
194                    }
195
196                    // Safe because an empty hit status would have broken out
197                    // in the block above
198                    let hit_status = hit_status_opt.expect("None case handled as miss");
199
200                    if !hit_status.is_fresh() {
201                        // expired or force expired asset
202                        if session.cache.is_cache_locked() {
203                            // first if this is the sub request for the background cache update
204                            if let Some(write_lock) = session
205                                .subrequest_ctx
206                                .as_mut()
207                                .and_then(|ctx| ctx.take_write_lock())
208                            {
209                                // Put the write lock in the request
210                                session.cache.set_write_lock(write_lock);
211                                session.cache.tag_as_subrequest();
212                                // and then let it go to upstream
213                                break None;
214                            }
215                            let will_serve_stale = session.cache.can_serve_stale_updating()
216                                && self.inner.should_serve_stale(session, ctx, None);
217                            if !will_serve_stale {
218                                let outcome = session.cache.cache_lock_wait().await;
219                                if self.handle_lock_wait_outcome(session, ctx, outcome) {
220                                    if self.cache_lock_retry_limit_exceeded(
221                                        session,
222                                        ctx,
223                                        &mut cache_lock_retries,
224                                    ) {
225                                        break None;
226                                    }
227                                    continue;
228                                } else {
229                                    break None;
230                                }
231                            }
232                            // else continue to serve stale
233                            session.cache.set_stale_updating();
234                        } else if session.cache.is_cache_lock_writer() {
235                            // stale while revalidate logic for the writer
236                            let will_serve_stale = session.cache.can_serve_stale_updating()
237                                && self.inner.should_serve_stale(session, ctx, None);
238                            if will_serve_stale {
239                                // create a background thread to do the actual update
240                                // the subrequest handle is only None by this phase in unit tests
241                                // that don't go through process_new_http
242                                let (permit, cache_lock) = session.cache.take_write_lock();
243                                SubrequestSpawner::new(self.clone()).spawn_background_subrequest(
244                                    session.as_ref(),
245                                    subrequest::Ctx::builder()
246                                        .cache_write_lock(
247                                            cache_lock,
248                                            session.cache.cache_key().clone(),
249                                            permit,
250                                        )
251                                        .build(),
252                                );
253                                // continue to serve stale for this request
254                                session.cache.set_stale_updating();
255                            } else {
256                                // return to fetch from upstream
257                                break None;
258                            }
259                        } else {
260                            // return to fetch from upstream
261                            break None;
262                        }
263                    }
264
265                    let (reuse, err) = self.proxy_cache_hit(session, ctx).await;
266                    if let Some(e) = err.as_ref() {
267                        error!(
268                            "Fail to serve cache: {e}, {}",
269                            self.inner.request_summary(session, ctx)
270                        );
271                    }
272                    // responses is served from cache, exit
273                    break Some((reuse, err));
274                }
275                Err(e) => {
276                    // Allow cache miss to fill cache even if cache lookup errors
277                    // this is mostly to support backward incompatible metadata update
278                    // TODO: check error types
279                    // session.cache.disable();
280                    self.inner.cache_miss(session, ctx);
281                    warn!(
282                        "Fail to cache lookup: {e}, {}",
283                        self.inner.request_summary(session, ctx)
284                    );
285                    break None;
286                }
287            }
288        }
289    }
290
291    // return bool: server_session can be reused, and error if any
292    pub(crate) async fn proxy_cache_hit(
293        &self,
294        session: &mut Session,
295        ctx: &mut SV::CTX,
296    ) -> (bool, Option<Box<Error>>)
297    where
298        SV: ProxyHttp + Send + Sync,
299        SV::CTX: Send + Sync,
300    {
301        use range_filter::*;
302
303        let seekable = session.cache.hit_handler().can_seek();
304        let mut header = cache_hit_header(&session.cache);
305
306        let req = session.req_header();
307
308        let not_modified = match self.inner.cache_not_modified_filter(session, &header, ctx) {
309            Ok(not_modified) => not_modified,
310            Err(e) => {
311                // fail open if cache_not_modified_filter errors,
312                // just return the whole original response
313                warn!(
314                    "Failed to run cache not modified filter: {e}, {}",
315                    self.inner.request_summary(session, ctx)
316                );
317                false
318            }
319        };
320        if not_modified {
321            to_304(&mut header);
322        }
323        let header_only = not_modified || req.method == http::method::Method::HEAD;
324
325        // process range header if the cache storage supports seek
326        let range_type = if seekable && !session.ignore_downstream_range {
327            self.inner.range_header_filter(session, &mut header, ctx)
328        } else {
329            RangeType::None
330        };
331
332        // return a 416 with an empty body for simplicity
333        let header_only = header_only || matches!(range_type, RangeType::Invalid);
334        debug!("header: {header:?}");
335
336        // TODO: use ProxyUseCache to replace the logic below
337        match self.inner.response_filter(session, &mut header, ctx).await {
338            Ok(_) => {
339                if let Err(e) = session
340                    .downstream_modules_ctx
341                    .response_header_filter(&mut header, header_only)
342                    .await
343                {
344                    error!(
345                        "Failed to run downstream modules response header filter in hit: {e}, {}",
346                        self.inner.request_summary(session, ctx)
347                    );
348                    session
349                        .as_mut()
350                        .respond_error(500)
351                        .await
352                        .unwrap_or_else(|e| {
353                            error!("failed to send error response to downstream: {e}");
354                        });
355                    // we have not write anything dirty to downstream, it is still reusable
356                    return (true, Some(e));
357                }
358
359                if let Err(e) = session
360                    .as_mut()
361                    .write_response_header(header)
362                    .await
363                    .map_err(|e| e.into_down())
364                {
365                    // downstream connection is bad already
366                    return (false, Some(e));
367                }
368            }
369            Err(e) => {
370                error!(
371                    "Failed to run response filter in hit: {e}, {}",
372                    self.inner.request_summary(session, ctx)
373                );
374                session
375                    .as_mut()
376                    .respond_error(500)
377                    .await
378                    .unwrap_or_else(|e| {
379                        error!("failed to send error response to downstream: {e}");
380                    });
381                // we have not write anything dirty to downstream, it is still reusable
382                return (true, Some(e));
383            }
384        }
385        debug!("finished sending cached header to downstream");
386
387        // If the function returns an Err, there was an issue seeking from the hit handler.
388        //
389        // Returning false means that no seeking or state change was done, either because the
390        // hit handler doesn't support the seek or because multipart doesn't apply.
391        fn seek_multipart(
392            hit_handler: &mut HitHandler,
393            range_filter: &mut RangeBodyFilter,
394        ) -> Result<bool> {
395            if !range_filter.is_multipart_range() || !hit_handler.can_seek_multipart() {
396                return Ok(false);
397            }
398            let r = range_filter.next_cache_multipart_range()?;
399            hit_handler.seek_multipart(r.start, Some(r.end))?;
400            // we still need RangeBodyFilter's help to transform the byte
401            // range into a multipart response.
402            range_filter.set_current_cursor(r.start);
403            Ok(true)
404        }
405
406        if !header_only {
407            let mut maybe_range_filter = match &range_type {
408                RangeType::Single(r) => {
409                    if session.cache.hit_handler().can_seek() {
410                        if let Err(e) = session.cache.hit_handler().seek(r.start, Some(r.end)) {
411                            return (false, Some(e));
412                        }
413                        None
414                    } else {
415                        Some(RangeBodyFilter::new_range(range_type.clone()))
416                    }
417                }
418                RangeType::Multi(_) => {
419                    let mut range_filter = RangeBodyFilter::new_range(range_type.clone());
420                    if let Err(e) = seek_multipart(session.cache.hit_handler(), &mut range_filter) {
421                        return (false, Some(e));
422                    }
423                    Some(range_filter)
424                }
425                RangeType::Invalid => unreachable!(),
426                RangeType::None => None,
427            };
428            loop {
429                match session.cache.hit_handler().read_body().await {
430                    Ok(raw_body) => {
431                        let end = raw_body.is_none();
432
433                        if end {
434                            if let Some(range_filter) = maybe_range_filter.as_mut() {
435                                if range_filter.should_cache_seek_again() {
436                                    let e = match seek_multipart(
437                                        session.cache.hit_handler(),
438                                        range_filter,
439                                    ) {
440                                        Ok(true) => {
441                                            // called seek(), read again
442                                            continue;
443                                        }
444                                        Ok(false) => {
445                                            // body reader can no longer seek multipart,
446                                            // but cache wants to continue seeking
447                                            // the body will just end in this case if we pass the
448                                            // None through
449                                            // (TODO: how might hit handlers want to recover from
450                                            // this situation)?
451                                            Error::explain(
452                                                InternalError,
453                                                "hit handler cannot seek for multipart again",
454                                            )
455                                            // the body will just end in this case.
456                                        }
457                                        Err(e) => e,
458                                    };
459                                    return (false, Some(e));
460                                }
461                            }
462                        }
463
464                        let mut body = if let Some(range_filter) = maybe_range_filter.as_mut() {
465                            range_filter.filter_body(raw_body)
466                        } else {
467                            raw_body
468                        };
469
470                        match self
471                            .inner
472                            .response_body_filter(session, &mut body, end, ctx)
473                        {
474                            Ok(Some(duration)) => {
475                                trace!("delaying response for {duration:?}");
476                                time::sleep(duration).await;
477                            }
478                            Ok(None) => { /* continue */ }
479                            Err(e) => {
480                                // body is being sent, don't treat downstream as reusable
481                                return (false, Some(e));
482                            }
483                        }
484
485                        if let Err(e) = session
486                            .downstream_modules_ctx
487                            .response_body_filter(&mut body, end)
488                        {
489                            // body is being sent, don't treat downstream as reusable
490                            return (false, Some(e));
491                        }
492
493                        if !end && body.as_ref().is_none_or(|b| b.is_empty()) {
494                            // Don't write empty body which will end session,
495                            // still more hit handler bytes to read
496                            continue;
497                        }
498
499                        // write to downstream
500                        let b = body.unwrap_or_default();
501                        if let Err(e) = session
502                            .as_mut()
503                            .write_response_body(b, end)
504                            .await
505                            .map_err(|e| e.into_down())
506                        {
507                            return (false, Some(e));
508                        }
509                        if end {
510                            break;
511                        }
512                    }
513                    Err(e) => return (false, Some(e)),
514                }
515            }
516        }
517
518        // No enabled() guard: no concurrent upstream can disable cache here.
519        if let Err(e) = session.cache.finish_hit_handler().await {
520            warn!("Error during finish_hit_handler: {}", e);
521        }
522
523        match session.as_mut().finish_body().await {
524            Ok(_) => {
525                debug!("finished sending cached body to downstream");
526                (true, None)
527            }
528            Err(e) => (false, Some(e)),
529        }
530    }
531
532    /* Downstream revalidation, only needed when cache is on because otherwise origin
533     * will handle it */
534    pub(crate) fn downstream_response_conditional_filter(
535        &self,
536        use_cache: &mut ServeFromCache,
537        session: &Session,
538        resp: &mut ResponseHeader,
539        ctx: &mut SV::CTX,
540    ) where
541        SV: ProxyHttp,
542    {
543        // TODO: range
544        let req = session.req_header();
545
546        let not_modified = match self.inner.cache_not_modified_filter(session, resp, ctx) {
547            Ok(not_modified) => not_modified,
548            Err(e) => {
549                // fail open if cache_not_modified_filter errors,
550                // just return the whole original response
551                warn!(
552                    "Failed to run cache not modified filter: {e}, {}",
553                    self.inner.request_summary(session, ctx)
554                );
555                false
556            }
557        };
558
559        if not_modified {
560            to_304(resp);
561        }
562        let header_only = not_modified || req.method == http::method::Method::HEAD;
563        if header_only && use_cache.is_on() {
564            // tell cache to stop serving downstream after yielding header
565            // (misses will continue to allow admitting upstream into cache)
566            use_cache.enable_header_only();
567        }
568    }
569
570    // TODO: cache upstream header filter to add/remove headers
571
572    async fn finish_miss_handler_best_effort(&self, session: &mut Session, ctx: &SV::CTX)
573    where
574        SV: ProxyHttp,
575    {
576        if let Err(e) = session.cache.finish_miss_handler().await {
577            warn!(
578                "Failed to finish cache miss admission: {e}, {}",
579                self.inner.request_summary(session, ctx)
580            );
581            session.cache.disable(NoCacheReason::StorageError);
582        }
583    }
584
585    pub(crate) async fn cache_http_task(
586        &self,
587        session: &mut Session,
588        task: &HttpTask,
589        ctx: &mut SV::CTX,
590        serve_from_cache: &mut ServeFromCache,
591    ) -> Result<()>
592    where
593        SV: ProxyHttp + Send + Sync,
594        SV::CTX: Send + Sync,
595    {
596        if !session.cache.enabled() && !session.cache.bypassing() {
597            return Ok(());
598        }
599
600        match task {
601            HttpTask::Header(header, end_stream) => {
602                // decide if cacheable and create cache meta
603                // for now, skip 1xxs (should not affect response cache decisions)
604                // However 101 is an exception because it is the final response header
605                if header.status.is_informational()
606                    && header.status != StatusCode::SWITCHING_PROTOCOLS
607                {
608                    return Ok(());
609                }
610                match self.inner.response_cache_filter(session, header, ctx)? {
611                    Cacheable(meta) => {
612                        let mut fill_cache = true;
613                        if session.cache.bypassing() {
614                            // Only hold this request back if the predictor bypassed it over size.
615                            // Re-enabling the cache without a known content length would fail the
616                            // request mid-body if the response exceeds the maximum file size
617                            // again, so wait for the body to finish and let the response filters
618                            // re-admit the key. Every other bypass reason says nothing about size
619                            // and must not be reported as PredictedResponseTooLarge.
620                            let bypassed_over_size =
621                                match session.cache.predicted_uncacheable_reason() {
622                                    Some(reason) => reason == NoCacheReason::ResponseTooLarge,
623                                    // unknown reason, stay conservative
624                                    None => true,
625                                };
626                            if bypassed_over_size
627                                && session.cache.max_file_size_bytes().is_some()
628                                && !meta.headers().contains_key(header::CONTENT_LENGTH)
629                            {
630                                session
631                                    .cache
632                                    .disable(NoCacheReason::PredictedResponseTooLarge);
633                                return Ok(());
634                            }
635
636                            session.cache.response_became_cacheable();
637
638                            if session.req_header().method == Method::GET
639                                && meta.response_header().status == StatusCode::OK
640                            {
641                                self.inner.cache_miss(session, ctx);
642                                if !session.cache.enabled() {
643                                    fill_cache = false;
644                                }
645                            } else {
646                                // we've allowed caching on the next request,
647                                // but do not cache _this_ request if bypassed and not 200
648                                // (We didn't run upstream request cache filters to strip range or condition headers,
649                                // so this could be an uncacheable response e.g. 206 or 304 or HEAD.
650                                // Exclude all non-200/GET for simplicity, may expand allowable codes in the future.)
651                                fill_cache = false;
652                                session.cache.disable(NoCacheReason::Deferred);
653                            }
654                        }
655
656                        // If the Content-Length is known, and a maximum asset size has been configured
657                        // on the cache, validate that the response does not exceed the maximum asset size.
658                        if session.cache.enabled() {
659                            if let Some(max_file_size) = session.cache.max_file_size_bytes() {
660                                let content_length_hdr = meta.headers().get(header::CONTENT_LENGTH);
661                                if let Some(content_length) =
662                                    header_value_content_length(content_length_hdr)
663                                {
664                                    if content_length > max_file_size {
665                                        fill_cache = false;
666                                        session.cache.response_became_uncacheable(
667                                            NoCacheReason::ResponseTooLarge,
668                                        );
669                                        session.cache.disable(NoCacheReason::ResponseTooLarge);
670                                        // too large to cache, disable ranging
671                                        session.ignore_downstream_range = true;
672                                    }
673                                }
674                                // if the content-length header is not specified, the miss handler
675                                // will count the response size on the fly, aborting the request
676                                // mid-transfer if the max file size is exceeded
677                            }
678                        }
679                        if fill_cache {
680                            let req_header = session.req_header();
681                            // Update the variance in the meta via the same callback,
682                            // cache_vary_filter(), used in cache lookup for consistency.
683                            // Future cache lookups need a matching variance in the meta
684                            // with the cache key to pick up the correct variance
685                            let variance = self.inner.cache_vary_filter(&meta, ctx, req_header);
686                            session.cache.set_cache_meta(meta);
687                            session.cache.update_variance(variance);
688                            // this sends the meta and header
689                            session.cache.set_miss_handler().await?;
690                            if session.cache.miss_body_reader().is_some() {
691                                serve_from_cache.enable_miss();
692                            }
693                            if *end_stream {
694                                session
695                                    .cache
696                                    .miss_handler()
697                                    .unwrap() // safe, it is set above
698                                    .write_body(Bytes::new(), true)
699                                    .await?;
700                                self.finish_miss_handler_best_effort(session, ctx).await;
701                            }
702                        }
703                    }
704                    Uncacheable(reason) => {
705                        if !session.cache.bypassing() {
706                            // mark as uncacheable, so we bypass cache next time
707                            session.cache.response_became_uncacheable(reason);
708                        }
709                        session.cache.disable(reason);
710                    }
711                }
712            }
713            HttpTask::Body(data, end_stream) | HttpTask::UpgradedBody(data, end_stream) => {
714                // It is not normally advisable to cache upgraded responses
715                // e.g. they are essentially close-delimited, so they are easily truncated
716                // but the framework still allows for it
717                match data {
718                    Some(d) => {
719                        if session.cache.enabled() {
720                            // TODO: do this async
721                            // fail if writing the body would exceed the max_file_size_bytes
722                            let body_size_allowed =
723                                session.cache.track_body_bytes_for_max_file_size(d.len());
724                            if !body_size_allowed {
725                                debug!("chunked response exceeded max cache size, remembering that it is uncacheable");
726                                session
727                                    .cache
728                                    .response_became_uncacheable(NoCacheReason::ResponseTooLarge);
729
730                                return Error::e_explain(
731                                    ERR_RESPONSE_TOO_LARGE,
732                                    format!(
733                                        "writing data of size {} bytes would exceed max file size of {} bytes",
734                                        d.len(),
735                                        session.cache.max_file_size_bytes().expect("max file size bytes must be set to exceed size")
736                                    ),
737                                );
738                            }
739
740                            // this will panic if more data is sent after we see end_stream
741                            // but should be impossible in real world
742                            let miss_handler = session.cache.miss_handler().unwrap();
743
744                            miss_handler.write_body(d.clone(), *end_stream).await?;
745                            if *end_stream {
746                                self.finish_miss_handler_best_effort(session, ctx).await;
747                            }
748                        }
749                    }
750                    None => {
751                        if session.cache.enabled() && *end_stream {
752                            self.finish_miss_handler_best_effort(session, ctx).await;
753                        }
754                    }
755                }
756            }
757            HttpTask::Trailer(_) => {} // h1 trailer is not supported yet
758            HttpTask::Done => {
759                if session.cache.enabled() {
760                    self.finish_miss_handler_best_effort(session, ctx).await;
761                }
762            }
763            HttpTask::Failed(_) => {
764                // TODO: handle this failure: delete the temp files?
765            }
766        }
767        Ok(())
768    }
769
770    // Decide if local cache can be used according to upstream http header
771    // 1. when upstream returns 304, the local cache is refreshed and served fresh
772    // 2. when upstream returns certain HTTP error status, the local cache is served stale
773    // Return true if local cache should be used, false otherwise
774    pub(crate) async fn revalidate_or_stale(
775        &self,
776        session: &mut Session,
777        task: &mut HttpTask,
778        ctx: &mut SV::CTX,
779    ) -> bool
780    where
781        SV: ProxyHttp + Send + Sync,
782        SV::CTX: Send + Sync,
783    {
784        if !session.cache.enabled() {
785            return false;
786        }
787
788        match task {
789            HttpTask::Header(resp, _eos) => {
790                if resp.status == StatusCode::NOT_MODIFIED {
791                    if session.cache.maybe_cache_meta().is_some() {
792                        // run upstream response filters on upstream 304 first
793                        if let Err(err) = self
794                            .inner
795                            .upstream_response_filter(session, resp, ctx)
796                            .await
797                        {
798                            error!("upstream response filter error on 304: {err:?}");
799                            session.cache.revalidate_uncacheable(
800                                *resp.clone(),
801                                NoCacheReason::InternalError,
802                            );
803                            // always serve from cache after receiving the 304
804                            return true;
805                        }
806                        // 304 doesn't contain all the headers, merge 304 into cached 200 header
807                        // in order for response_cache_filter to run correctly
808                        let merged_header = session.cache.revalidate_merge_header(resp);
809                        match self
810                            .inner
811                            .response_cache_filter(session, &merged_header, ctx)
812                        {
813                            Ok(Cacheable(mut meta)) => {
814                                // For simplicity, ignore changes to variance over 304 for now.
815                                // Note this means upstream can only update variance via 2xx
816                                // (expired response).
817                                //
818                                // TODO: if we choose to respect changing Vary / variance over 304,
819                                // then there are a few cases to consider. See `update_variance` in
820                                // the `pingora-cache` module.
821                                let old_meta = session.cache.maybe_cache_meta().unwrap(); // safe, checked above
822                                if let Some(old_variance) = old_meta.variance() {
823                                    meta.set_variance(old_variance);
824                                }
825                                if let Err(e) = session.cache.revalidate_cache_meta(meta).await {
826                                    // Fail open: we can continue use the revalidated response even
827                                    // if the meta failed to write to storage
828                                    warn!("revalidate_cache_meta failed {e:?}");
829                                }
830                            }
831                            Ok(Uncacheable(reason)) => {
832                                // This response was once cacheable, and upstream tells us it has not changed
833                                // but now we decided it is uncacheable!
834                                // RFC 9111: still allowed to reuse stored response this time because
835                                // it was "successfully validated"
836                                // https://www.rfc-editor.org/rfc/rfc9111#constructing.responses.from.caches
837                                // Serve the response, but do not update cache
838
839                                // We also want to avoid poisoning downstream's cache with an unsolicited 304
840                                // if we did not receive a conditional request from downstream
841                                // (downstream may have a different cacheability assessment and could cache the 304)
842
843                                //TODO: log more
844                                debug!("Uncacheable {reason:?} 304 received");
845                                session.cache.response_became_uncacheable(reason);
846                                session.cache.revalidate_uncacheable(merged_header, reason);
847                            }
848                            Err(e) => {
849                                // Error during revalidation, similarly to the reasons above
850                                // (avoid poisoning downstream cache with passthrough 304),
851                                // allow serving the stored response without updating cache
852                                warn!("Error {e:?} response_cache_filter during revalidation");
853                                session.cache.revalidate_uncacheable(
854                                    merged_header,
855                                    NoCacheReason::InternalError,
856                                );
857                                // Assume the next 304 may succeed, so don't mark uncacheable
858                            }
859                        }
860                        // always serve from cache after receiving the 304
861                        true
862                    } else {
863                        //TODO: log more
864                        warn!("304 received without cached asset, disable caching");
865                        let reason = NoCacheReason::Custom("304 on miss");
866                        session.cache.response_became_uncacheable(reason);
867                        session.cache.disable(reason);
868                        false
869                    }
870                } else if resp.status.is_server_error() {
871                    // stale if error logic, 5xx only for now
872
873                    // this is response header filter, response_written should always be None?
874                    if !session.cache.can_serve_stale_error()
875                        || session.response_written().is_some()
876                    {
877                        return false;
878                    }
879
880                    // create an error to encode the http status code
881                    let http_status_error = Error::create(
882                        ErrorType::HTTPStatus(resp.status.as_u16()),
883                        ErrorSource::Upstream,
884                        None,
885                        None,
886                    );
887                    if self
888                        .inner
889                        .should_serve_stale(session, ctx, Some(&http_status_error))
890                    {
891                        // no more need to keep the write lock
892                        session
893                            .cache
894                            .release_write_lock(NoCacheReason::UpstreamError);
895                        true
896                    } else {
897                        false
898                    }
899                } else {
900                    false // not 304, not stale if error status code
901                }
902            }
903            _ => false, // not header
904        }
905    }
906
907    // None: no staled asset is used, Some(_): staled asset is sent to downstream
908    // bool: can the downstream connection be reused
909    pub(crate) async fn handle_stale_if_error(
910        &self,
911        session: &mut Session,
912        ctx: &mut SV::CTX,
913        error: &Error,
914    ) -> Option<(bool, Option<Box<Error>>)>
915    where
916        SV: ProxyHttp + Send + Sync,
917        SV::CTX: Send + Sync,
918    {
919        // the caller might already checked this as an optimization
920        if !session.cache.can_serve_stale_error() {
921            return None;
922        }
923
924        // the error happen halfway through a regular response to downstream
925        // can't resend the response
926        if session.response_written().is_some() {
927            return None;
928        }
929
930        // check error types
931        if !self.inner.should_serve_stale(session, ctx, Some(error)) {
932            return None;
933        }
934
935        // log the original error
936        warn!(
937            "Fail to proxy: {}, serving stale, {}",
938            error,
939            self.inner.request_summary(session, ctx)
940        );
941
942        // no more need to hang onto the cache lock
943        session
944            .cache
945            .release_write_lock(NoCacheReason::UpstreamError);
946
947        Some(self.proxy_cache_hit(session, ctx).await)
948    }
949
950    // helper function to check when to continue to retry lock (true) or give up (false)
951    fn handle_lock_wait_outcome(
952        &self,
953        session: &mut Session,
954        ctx: &SV::CTX,
955        outcome: LockWaitOutcome,
956    ) -> bool
957    where
958        SV: ProxyHttp,
959    {
960        debug!("cache unlocked {outcome:?}");
961        match outcome {
962            // should lookup the cached asset again
963            LockWaitOutcome::Done => true,
964            // should compete to be a new writer
965            LockWaitOutcome::TransientError => true,
966            // the writer found no lock was needed; every reader goes upstream
967            LockWaitOutcome::GiveUp => {
968                session.cache.disable(NoCacheReason::CacheLockGiveUp);
969                false
970            }
971            // this reader alone stopped waiting, over a fill it cannot use; the
972            // cause travels with the outcome, so there is nothing to look up
973            LockWaitOutcome::Abandoned { reason, .. } => {
974                session.cache.disable(reason);
975                false
976            }
977            // treat this the same as TransientError
978            LockWaitOutcome::Dangling => {
979                // software bug, but request can recover from this
980                warn!(
981                    "Dangling cache lock, {}",
982                    self.inner.request_summary(session, ctx)
983                );
984                true
985            }
986            // If this reader has spent too long waiting on locks, let the request
987            // through while disabling cache (to avoid amplifying disk writes).
988            LockWaitOutcome::WaitTimeout => {
989                warn!(
990                    "Cache lock timeout, {}",
991                    self.inner.request_summary(session, ctx)
992                );
993                session.cache.disable(NoCacheReason::CacheLockTimeout);
994                // not cacheable, just go to the origin.
995                false
996            }
997            // When a singular cache lock has been held for too long,
998            // we should allow requests to recompete for the lock
999            // to protect upstreams from load.
1000            LockWaitOutcome::AgeTimeout => true,
1001        }
1002    }
1003
1004    fn cache_lock_retry_limit_exceeded(
1005        &self,
1006        session: &mut Session,
1007        ctx: &SV::CTX,
1008        cache_lock_retries: &mut usize,
1009    ) -> bool
1010    where
1011        SV: ProxyHttp,
1012    {
1013        *cache_lock_retries += 1;
1014        let max_retries = session
1015            .cache
1016            .cache_lock_max_retries()
1017            .unwrap_or(DEFAULT_MAX_CACHE_LOCK_RETRIES);
1018        if *cache_lock_retries <= max_retries {
1019            return false;
1020        }
1021
1022        warn!(
1023            "Cache lock retry limit exceeded, {}",
1024            self.inner.request_summary(session, ctx)
1025        );
1026        session.cache.disable(NoCacheReason::CacheLockRetryLimit);
1027        true
1028    }
1029}
1030
1031fn cache_hit_header(cache: &HttpCache) -> Box<ResponseHeader> {
1032    let mut header = Box::new(cache.cache_meta().response_header_copy());
1033    // convert cache response
1034
1035    // these status codes / method cannot have body, so no need to add chunked encoding
1036    let no_body = matches!(header.status.as_u16(), 204 | 304);
1037
1038    // https://www.rfc-editor.org/rfc/rfc9111#section-4:
1039    // When a stored response is used to satisfy a request without validation, a cache
1040    // MUST generate an Age header field
1041    if !cache.upstream_used() {
1042        let age = cache.cache_meta().age().as_secs();
1043        header.insert_header(http::header::AGE, age).unwrap();
1044    }
1045    log::debug!("cache header: {header:?} {:?}", cache.phase());
1046
1047    // currently storage cache is always considered an h1 upstream
1048    // (header-serde serializes as h1.0 or h1.1)
1049    // set this header to be h1.1
1050    header.set_version(Version::HTTP_11);
1051
1052    /* Add chunked header to tell downstream to use chunked encoding
1053     * during the absent of content-length in h2 */
1054    if !no_body
1055        && !header.status.is_informational()
1056        && header.headers.get(http::header::CONTENT_LENGTH).is_none()
1057    {
1058        header
1059            .insert_header(http::header::TRANSFER_ENCODING, "chunked")
1060            .unwrap();
1061    }
1062    header
1063}
1064
1065// https://datatracker.ietf.org/doc/html/rfc7233#section-3
1066pub mod range_filter {
1067    use super::*;
1068    use bytes::BytesMut;
1069    use http::header::*;
1070    use std::ops::Range;
1071
1072    // parse bytes into usize, ignores specific error
1073    fn parse_number(input: &[u8]) -> Option<usize> {
1074        str::from_utf8(input).ok()?.parse().ok()
1075    }
1076
1077    fn parse_range_header(
1078        range: &[u8],
1079        content_length: usize,
1080        max_multipart_ranges: Option<usize>,
1081    ) -> RangeType {
1082        use regex::Regex;
1083
1084        // Match individual range parts, (e.g. "0-100", "-5", "1-")
1085        static RE_SINGLE_RANGE_PART: Lazy<Regex> =
1086            Lazy::new(|| Regex::new(r"(?i)^\s*(?P<start>\d*)-(?P<end>\d*)\s*$").unwrap());
1087
1088        // Convert bytes to UTF-8 string
1089        let range_str = match str::from_utf8(range) {
1090            Ok(s) => s,
1091            Err(_) => return RangeType::None,
1092        };
1093
1094        // Split into "bytes=" and the actual range(s)
1095        let mut parts = range_str.splitn(2, "=");
1096
1097        // Check if it starts with "bytes="
1098        let prefix = parts.next();
1099        if !prefix.is_some_and(|s| s.eq_ignore_ascii_case("bytes")) {
1100            return RangeType::None;
1101        }
1102
1103        let Some(ranges_str) = parts.next() else {
1104            // No ranges provided
1105            return RangeType::None;
1106        };
1107
1108        // "bytes=" with an empty (or whitespace-only) range-set is syntactically a
1109        // range request with zero satisfiable range-specs, so return 416.
1110        if ranges_str.trim().is_empty() {
1111            return RangeType::Invalid;
1112        }
1113
1114        // Get the actual range string (e.g."100-200,300-400")
1115        let mut range_count = 0;
1116        for _ in ranges_str.split(',') {
1117            range_count += 1;
1118            if let Some(max_ranges) = max_multipart_ranges {
1119                if range_count >= max_ranges {
1120                    // If we get more than max configured ranges, return None for now to save parsing time
1121                    return RangeType::None;
1122                }
1123            }
1124        }
1125        let mut ranges: Vec<Range<usize>> = Vec::with_capacity(range_count);
1126
1127        // Process each range
1128        let mut last_range_end = 0;
1129        for part in ranges_str.split(',') {
1130            let captured = match RE_SINGLE_RANGE_PART.captures(part) {
1131                Some(c) => c,
1132                None => {
1133                    return RangeType::None;
1134                }
1135            };
1136
1137            let maybe_start = captured
1138                .name("start")
1139                .and_then(|s| s.as_str().parse::<usize>().ok());
1140            let end = captured
1141                .name("end")
1142                .and_then(|s| s.as_str().parse::<usize>().ok());
1143
1144            let range = if let Some(start) = maybe_start {
1145                if start >= content_length {
1146                    // Skip the invalid range
1147                    continue;
1148                }
1149                // open-ended range should end at the last byte
1150                // over sized end is allowed but ignored
1151                // range end is inclusive
1152                let end = std::cmp::min(end.unwrap_or(content_length - 1), content_length - 1) + 1;
1153                if end <= start {
1154                    // Skip the invalid range
1155                    continue;
1156                }
1157                start..end
1158            } else {
1159                // start is empty, this changes the meaning of the value of `end`
1160                // Now it means to read the last `end` bytes
1161                if let Some(end) = end {
1162                    if content_length >= end {
1163                        (content_length - end)..content_length
1164                    } else {
1165                        // over sized end is allowed but ignored
1166                        0..content_length
1167                    }
1168                } else {
1169                    // No start or end, skip the invalid range
1170                    continue;
1171                }
1172            };
1173            // For now we stick to non-overlapping, ascending ranges for simplicity
1174            // and parity with nginx
1175            if range.start < last_range_end {
1176                return RangeType::None;
1177            }
1178            last_range_end = range.end;
1179            ranges.push(range);
1180        }
1181
1182        // Note for future: we can technically coalesce multiple ranges for multipart
1183        //
1184        // https://www.rfc-editor.org/rfc/rfc9110#section-17.15
1185        // "Servers ought to ignore, coalesce, or reject egregious range
1186        // requests, such as requests for more than two overlapping ranges or
1187        // for many small ranges in a single set, particularly when the ranges
1188        // are requested out of order for no apparent reason. Multipart range
1189        // requests are not designed to support random access."
1190
1191        if ranges.is_empty() {
1192            // We got some ranges, processed them but none were valid
1193            RangeType::Invalid
1194        } else if ranges.len() == 1 {
1195            RangeType::Single(ranges[0].clone()) // Only 1 index
1196        } else {
1197            RangeType::Multi(MultiRangeInfo::new(ranges))
1198        }
1199    }
1200    #[test]
1201    fn test_parse_range() {
1202        assert_eq!(
1203            parse_range_header(b"bytes=0-1", 10, None),
1204            RangeType::new_single(0, 2)
1205        );
1206        assert_eq!(
1207            parse_range_header(b"bYTes=0-9", 10, None),
1208            RangeType::new_single(0, 10)
1209        );
1210        assert_eq!(
1211            parse_range_header(b"bytes=0-12", 10, None),
1212            RangeType::new_single(0, 10)
1213        );
1214        assert_eq!(
1215            parse_range_header(b"bytes=0-", 10, None),
1216            RangeType::new_single(0, 10)
1217        );
1218        assert_eq!(
1219            parse_range_header(b"bytes=2-1", 10, None),
1220            RangeType::Invalid
1221        );
1222        assert_eq!(
1223            parse_range_header(b"bytes=10-11", 10, None),
1224            RangeType::Invalid
1225        );
1226        assert_eq!(
1227            parse_range_header(b"bytes=-2", 10, None),
1228            RangeType::new_single(8, 10)
1229        );
1230        assert_eq!(
1231            parse_range_header(b"bytes=-12", 10, None),
1232            RangeType::new_single(0, 10)
1233        );
1234        assert_eq!(parse_range_header(b"bytes=-", 10, None), RangeType::Invalid);
1235        assert_eq!(parse_range_header(b"bytes=", 10, None), RangeType::Invalid);
1236        assert_eq!(
1237            parse_range_header(b"bytes=  ", 10, None),
1238            RangeType::Invalid
1239        );
1240    }
1241
1242    // Add some tests for multi-range too
1243    #[test]
1244    fn test_parse_range_header_multi() {
1245        assert_eq!(
1246            parse_range_header(b"bytes=0-1,4-5", 10, None)
1247                .get_multirange_info()
1248                .expect("Should have multipart info for Multipart range request")
1249                .ranges,
1250            (vec![Range { start: 0, end: 2 }, Range { start: 4, end: 6 }])
1251        );
1252        // Last range is invalid because the content-length is too small
1253        assert_eq!(
1254            parse_range_header(b"bytEs=0-99,200-299,400-499", 320, None)
1255                .get_multirange_info()
1256                .expect("Should have multipart info for Multipart range request")
1257                .ranges,
1258            (vec![
1259                Range { start: 0, end: 100 },
1260                Range {
1261                    start: 200,
1262                    end: 300
1263                }
1264            ])
1265        );
1266        // Same as above but appropriate content length
1267        assert_eq!(
1268            parse_range_header(b"bytEs=0-99,200-299,400-499", 500, None)
1269                .get_multirange_info()
1270                .expect("Should have multipart info for Multipart range request")
1271                .ranges,
1272            vec![
1273                Range { start: 0, end: 100 },
1274                Range {
1275                    start: 200,
1276                    end: 300
1277                },
1278                Range {
1279                    start: 400,
1280                    end: 500
1281                },
1282            ]
1283        );
1284        // Looks like a range request but it is continuous, we decline to range
1285        assert_eq!(
1286            parse_range_header(b"bytes=0-,-2", 10, None),
1287            RangeType::None,
1288        );
1289        // Should not have multirange info set
1290        assert!(parse_range_header(b"bytes=0-,-2", 10, None)
1291            .get_multirange_info()
1292            .is_none());
1293        // Overlapping ranges, these ranges are currently declined
1294        assert_eq!(
1295            parse_range_header(b"bytes=0-3,2-5", 10, None),
1296            RangeType::None,
1297        );
1298        assert!(parse_range_header(b"bytes=0-3,2-5", 10, None)
1299            .get_multirange_info()
1300            .is_none());
1301
1302        // Content length is 2, so only range is 0-2.
1303        assert_eq!(
1304            parse_range_header(b"bytes=0-5,10-", 2, None),
1305            RangeType::new_single(0, 2)
1306        );
1307        assert!(parse_range_header(b"bytes=0-5,10-", 2, None)
1308            .get_multirange_info()
1309            .is_none());
1310
1311        // We should ignore the last incorrect range and return the other acceptable ranges
1312        assert_eq!(
1313            parse_range_header(b"bytes=0-5, 10-20, 30-18", 200, None)
1314                .get_multirange_info()
1315                .expect("Should have multipart info for Multipart range request")
1316                .ranges,
1317            vec![Range { start: 0, end: 6 }, Range { start: 10, end: 21 },]
1318        );
1319        // All invalid ranges
1320        assert_eq!(
1321            parse_range_header(b"bytes=5-0, 20-15, 30-25", 200, None),
1322            RangeType::Invalid
1323        );
1324
1325        // Helper function to generate a large number of ranges for the next test
1326        fn generate_range_header(count: usize) -> Vec<u8> {
1327            let mut s = String::from("bytes=");
1328            for i in 0..count {
1329                let start = i * 4;
1330                let end = start + 1;
1331                if i > 0 {
1332                    s.push(',');
1333                }
1334                s.push_str(&start.to_string());
1335                s.push('-');
1336                s.push_str(&end.to_string());
1337            }
1338            s.into_bytes()
1339        }
1340
1341        // Test 200 range limit for parsing.
1342        let ranges = generate_range_header(201);
1343        assert_eq!(
1344            parse_range_header(&ranges, 1000, Some(200)),
1345            RangeType::None
1346        )
1347    }
1348
1349    // For Multipart Requests, we need to know the boundary, content length and type across
1350    // the headers and the body. So let us store this information as part of the range
1351    #[derive(Debug, Eq, PartialEq, Clone)]
1352    pub struct MultiRangeInfo {
1353        pub ranges: Vec<Range<usize>>,
1354        pub boundary: String,
1355        total_length: usize,
1356        pub content_type: Option<String>,
1357    }
1358
1359    impl MultiRangeInfo {
1360        // Create a new MultiRangeInfo, when we just have the ranges
1361        pub fn new(ranges: Vec<Range<usize>>) -> Self {
1362            Self {
1363                ranges,
1364                // Directly create boundary string on initialization
1365                boundary: Self::generate_boundary(),
1366                total_length: 0,
1367                content_type: None,
1368            }
1369        }
1370        pub fn set_content_type(&mut self, content_type: String) {
1371            self.content_type = Some(content_type)
1372        }
1373        pub fn set_total_length(&mut self, total_length: usize) {
1374            self.total_length = total_length;
1375        }
1376        // Per [RFC 9110](https://www.rfc-editor.org/rfc/rfc9110.html#multipart.byteranges),
1377        // we need generate a boundary string for each body part.
1378        // Per [RFC 2046](https://www.rfc-editor.org/rfc/rfc2046#section-5.1.1), the boundary should be no longer than 70 characters
1379        // and it must not match the body content.
1380        fn generate_boundary() -> String {
1381            use rand::Rng;
1382            let mut rng: rand::prelude::ThreadRng = rand::thread_rng();
1383            format!("{:016x}", rng.gen::<u64>())
1384        }
1385        pub fn calculate_multipart_length(&self) -> usize {
1386            let mut total_length = 0;
1387            let content_type = self.content_type.as_ref();
1388            for range in self.ranges.clone() {
1389                // Each part should have
1390                // \r\n--boundary\r\n                         --> 4 + boundary.len() (16) + 2 = 20
1391                // Content-Type: original-content-type\r\n    --> 14 + content_type.len() + 2
1392                // Content-Range: bytes start-end/total\r\n   --> Variable +2
1393                // \r\n                                       --> 2
1394                // [data]                                     --> data.len()
1395                total_length += 4 + self.boundary.len() + 2;
1396                total_length += content_type.map_or(0, |ct| 14 + ct.len() + 2);
1397                total_length += format!(
1398                    "Content-Range: bytes {}-{}/{}",
1399                    range.start,
1400                    range.end - 1,
1401                    self.total_length
1402                )
1403                .len()
1404                    + 2;
1405                total_length += 2;
1406                total_length += range.end - range.start;
1407            }
1408            // Final boundary: "\r\n--<boundary>--\r\n"
1409            total_length += 4 + self.boundary.len() + 4;
1410            total_length
1411        }
1412    }
1413    #[derive(Debug, Eq, PartialEq, Clone)]
1414    pub enum RangeType {
1415        None,
1416        Single(Range<usize>),
1417        Multi(MultiRangeInfo),
1418        Invalid,
1419    }
1420
1421    impl RangeType {
1422        // Helper functions for tests
1423        #[allow(dead_code)]
1424        fn new_single(start: usize, end: usize) -> Self {
1425            RangeType::Single(Range { start, end })
1426        }
1427        #[allow(dead_code)]
1428        pub fn new_multi(ranges: Vec<Range<usize>>) -> Self {
1429            RangeType::Multi(MultiRangeInfo::new(ranges))
1430        }
1431        #[allow(dead_code)]
1432        fn get_multirange_info(&self) -> Option<&MultiRangeInfo> {
1433            match self {
1434                RangeType::Multi(multi_range_info) => Some(multi_range_info),
1435                _ => None,
1436            }
1437        }
1438        #[allow(dead_code)]
1439        fn update_multirange_info(&mut self, content_length: usize, content_type: Option<String>) {
1440            if let RangeType::Multi(multipart_range_info) = self {
1441                multipart_range_info.content_type = content_type;
1442                multipart_range_info.set_total_length(content_length);
1443            }
1444        }
1445    }
1446
1447    // Handles both single-range and multipart-range requests
1448    pub fn range_header_filter(
1449        req: &RequestHeader,
1450        resp: &mut ResponseHeader,
1451        max_multipart_ranges: Option<usize>,
1452    ) -> RangeType {
1453        // The Range header field is evaluated after evaluating the precondition
1454        // header fields defined in [RFC7232], and only if the result in absence
1455        // of the Range header field would be a 200 (OK) response
1456        if resp.status != StatusCode::OK {
1457            return RangeType::None;
1458        }
1459
1460        // Content-Length is not required by RFC but it is what nginx does and easier to implement
1461        // with this header present.
1462        let Some(content_length_bytes) = resp.headers.get(CONTENT_LENGTH) else {
1463            return RangeType::None;
1464        };
1465        // bail on invalid content length
1466        let Some(content_length) = parse_number(content_length_bytes.as_bytes()) else {
1467            return RangeType::None;
1468        };
1469
1470        // At this point the response is allowed to be served as ranges
1471        // TODO: we can also check Accept-Range header from resp. Nginx gives uses the option
1472        // see proxy_force_ranges
1473
1474        fn request_range_type(
1475            req: &RequestHeader,
1476            resp: &ResponseHeader,
1477            content_length: usize,
1478            max_multipart_ranges: Option<usize>,
1479        ) -> RangeType {
1480            // "A server MUST ignore a Range header field received with a request method other than GET."
1481            if req.method != http::Method::GET && req.method != http::Method::HEAD {
1482                return RangeType::None;
1483            }
1484
1485            let Some(range_header) = req.headers.get(RANGE) else {
1486                return RangeType::None;
1487            };
1488
1489            // if-range wants to understand if the Last-Modified / ETag value matches exactly for use
1490            // with resumable downloads.
1491            // https://datatracker.ietf.org/doc/html/rfc9110#name-if-range
1492            // Note that the RFC wants strong validation, and suggests that
1493            // "A valid entity-tag can be distinguished from a valid HTTP-date
1494            // by examining the first three characters for a DQUOTE,"
1495            // but this current etag matching behavior most closely mirrors nginx.
1496            if let Some(if_range) = req.headers.get(IF_RANGE) {
1497                let ir = if_range.as_bytes();
1498                let matches = if ir.len() >= 2 && ir.last() == Some(&b'"') {
1499                    resp.headers.get(ETAG).is_some_and(|etag| etag == if_range)
1500                } else if let Some(last_modified) = resp.headers.get(LAST_MODIFIED) {
1501                    last_modified == if_range
1502                } else {
1503                    false
1504                };
1505                if !matches {
1506                    return RangeType::None;
1507                }
1508            }
1509
1510            parse_range_header(
1511                range_header.as_bytes(),
1512                content_length,
1513                max_multipart_ranges,
1514            )
1515        }
1516
1517        let mut range_type = request_range_type(req, resp, content_length, max_multipart_ranges);
1518
1519        match &mut range_type {
1520            RangeType::None => {
1521                // At this point, the response is _eligible_ to be served in ranges
1522                // in the future, so add Accept-Ranges, mirroring nginx behavior
1523                resp.insert_header(&ACCEPT_RANGES, "bytes").unwrap();
1524            }
1525            RangeType::Single(r) => {
1526                // 206 response
1527                resp.set_status(StatusCode::PARTIAL_CONTENT).unwrap();
1528                resp.remove_header(&ACCEPT_RANGES);
1529                resp.insert_header(&CONTENT_LENGTH, r.end - r.start)
1530                    .unwrap();
1531                resp.insert_header(
1532                    &CONTENT_RANGE,
1533                    format!("bytes {}-{}/{content_length}", r.start, r.end - 1), // range end is inclusive
1534                )
1535                .unwrap()
1536            }
1537
1538            RangeType::Multi(multi_range_info) => {
1539                let content_type = resp
1540                    .headers
1541                    .get(CONTENT_TYPE)
1542                    .and_then(|v| v.to_str().ok())
1543                    .unwrap_or("application/octet-stream");
1544                // Update multipart info
1545                multi_range_info.set_total_length(content_length);
1546                multi_range_info.set_content_type(content_type.to_string());
1547
1548                let total_length = multi_range_info.calculate_multipart_length();
1549
1550                resp.set_status(StatusCode::PARTIAL_CONTENT).unwrap();
1551                resp.remove_header(&ACCEPT_RANGES);
1552                resp.insert_header(CONTENT_LENGTH, total_length).unwrap();
1553                resp.insert_header(
1554                    CONTENT_TYPE,
1555                    format!(
1556                        "multipart/byteranges; boundary={}",
1557                        multi_range_info.boundary
1558                    ), // RFC 2046
1559                )
1560                .unwrap();
1561                resp.remove_header(&CONTENT_RANGE);
1562            }
1563            RangeType::Invalid => {
1564                // 416 response
1565                resp.set_status(StatusCode::RANGE_NOT_SATISFIABLE).unwrap();
1566                // empty body for simplicity
1567                resp.insert_header(&CONTENT_LENGTH, HeaderValue::from_static("0"))
1568                    .unwrap();
1569                resp.remove_header(&ACCEPT_RANGES);
1570                resp.remove_header(&CONTENT_TYPE);
1571                resp.remove_header(&CONTENT_ENCODING);
1572                resp.remove_header(&TRANSFER_ENCODING);
1573                resp.insert_header(&CONTENT_RANGE, format!("bytes */{content_length}"))
1574                    .unwrap()
1575            }
1576        }
1577
1578        range_type
1579    }
1580
1581    #[test]
1582    fn test_range_filter_single() {
1583        fn gen_req() -> RequestHeader {
1584            RequestHeader::build(http::Method::GET, b"/", Some(1)).unwrap()
1585        }
1586        fn gen_resp() -> ResponseHeader {
1587            let mut resp = ResponseHeader::build(200, Some(1)).unwrap();
1588            resp.append_header("Content-Length", "10").unwrap();
1589            resp
1590        }
1591
1592        // no range
1593        let req = gen_req();
1594        let mut resp = gen_resp();
1595        assert_eq!(RangeType::None, range_header_filter(&req, &mut resp, None));
1596        assert_eq!(resp.status.as_u16(), 200);
1597        assert_eq!(
1598            resp.headers.get("accept-ranges").unwrap().as_bytes(),
1599            b"bytes"
1600        );
1601
1602        // no range, try HEAD
1603        let mut req = gen_req();
1604        req.set_method(Method::HEAD);
1605        let mut resp = gen_resp();
1606        assert_eq!(RangeType::None, range_header_filter(&req, &mut resp, None));
1607        assert_eq!(resp.status.as_u16(), 200);
1608        assert_eq!(
1609            resp.headers.get("accept-ranges").unwrap().as_bytes(),
1610            b"bytes"
1611        );
1612
1613        // regular range
1614        let mut req = gen_req();
1615        req.insert_header("Range", "bytes=0-1").unwrap();
1616        let mut resp = gen_resp();
1617        assert_eq!(
1618            RangeType::new_single(0, 2),
1619            range_header_filter(&req, &mut resp, None)
1620        );
1621        assert_eq!(resp.status.as_u16(), 206);
1622        assert_eq!(resp.headers.get("content-length").unwrap().as_bytes(), b"2");
1623        assert_eq!(
1624            resp.headers.get("content-range").unwrap().as_bytes(),
1625            b"bytes 0-1/10"
1626        );
1627        assert!(resp.headers.get("accept-ranges").is_none());
1628
1629        // regular range, accept-ranges included
1630        let mut req = gen_req();
1631        req.insert_header("Range", "bytes=0-1").unwrap();
1632        let mut resp = gen_resp();
1633        resp.insert_header("Accept-Ranges", "bytes").unwrap();
1634        assert_eq!(
1635            RangeType::new_single(0, 2),
1636            range_header_filter(&req, &mut resp, None)
1637        );
1638        assert_eq!(resp.status.as_u16(), 206);
1639        assert_eq!(resp.headers.get("content-length").unwrap().as_bytes(), b"2");
1640        assert_eq!(
1641            resp.headers.get("content-range").unwrap().as_bytes(),
1642            b"bytes 0-1/10"
1643        );
1644        // accept-ranges stripped
1645        assert!(resp.headers.get("accept-ranges").is_none());
1646
1647        // bad range
1648        let mut req = gen_req();
1649        req.insert_header("Range", "bytes=1-0").unwrap();
1650        let mut resp = gen_resp();
1651        resp.insert_header("Accept-Ranges", "bytes").unwrap();
1652        resp.insert_header("Content-Encoding", "gzip").unwrap();
1653        resp.insert_header("Transfer-Encoding", "chunked").unwrap();
1654        assert_eq!(
1655            RangeType::Invalid,
1656            range_header_filter(&req, &mut resp, None)
1657        );
1658        assert_eq!(resp.status.as_u16(), 416);
1659        assert_eq!(resp.headers.get("content-length").unwrap().as_bytes(), b"0");
1660        assert_eq!(
1661            resp.headers.get("content-range").unwrap().as_bytes(),
1662            b"bytes */10"
1663        );
1664        assert!(resp.headers.get("accept-ranges").is_none());
1665        assert!(resp.headers.get("content-encoding").is_none());
1666        assert!(resp.headers.get("transfer-encoding").is_none());
1667    }
1668
1669    // Multipart Tests
1670    #[test]
1671    fn test_range_filter_multipart() {
1672        fn gen_req() -> RequestHeader {
1673            let mut req: RequestHeader =
1674                RequestHeader::build(http::Method::GET, b"/", Some(1)).unwrap();
1675            req.append_header("Range", "bytes=0-1,3-4,6-7").unwrap();
1676            req
1677        }
1678        fn gen_req_overlap_range() -> RequestHeader {
1679            let mut req: RequestHeader =
1680                RequestHeader::build(http::Method::GET, b"/", Some(1)).unwrap();
1681            req.append_header("Range", "bytes=0-3,2-5,7-8").unwrap();
1682            req
1683        }
1684        fn gen_resp() -> ResponseHeader {
1685            let mut resp = ResponseHeader::build(200, Some(1)).unwrap();
1686            resp.append_header("Content-Length", "10").unwrap();
1687            resp
1688        }
1689
1690        // valid multipart range
1691        let req = gen_req();
1692        let mut resp = gen_resp();
1693        let result = range_header_filter(&req, &mut resp, None);
1694        let mut boundary_str = String::new();
1695
1696        assert!(matches!(result, RangeType::Multi(_)));
1697        if let RangeType::Multi(multi_part_info) = result {
1698            assert_eq!(multi_part_info.ranges.len(), 3);
1699            assert_eq!(multi_part_info.ranges[0], Range { start: 0, end: 2 });
1700            assert_eq!(multi_part_info.ranges[1], Range { start: 3, end: 5 });
1701            assert_eq!(multi_part_info.ranges[2], Range { start: 6, end: 8 });
1702            // Verify that multipart info has been set
1703            assert!(multi_part_info.content_type.is_some());
1704            assert_eq!(multi_part_info.total_length, 10);
1705            assert!(!multi_part_info.boundary.is_empty());
1706            boundary_str = multi_part_info.boundary;
1707        }
1708        assert_eq!(resp.status.as_u16(), 206);
1709        // Verify that boundary is the same in header and in multipartinfo
1710        assert_eq!(
1711            resp.headers.get("content-type").unwrap().to_str().unwrap(),
1712            format!("multipart/byteranges; boundary={boundary_str}")
1713        );
1714        assert!(resp.headers.get("content_length").is_none());
1715        assert!(resp.headers.get("accept-ranges").is_none());
1716
1717        // overlapping range, multipart range is declined
1718        let req = gen_req_overlap_range();
1719        let mut resp = gen_resp();
1720        let result = range_header_filter(&req, &mut resp, None);
1721
1722        assert!(matches!(result, RangeType::None));
1723        assert_eq!(resp.status.as_u16(), 200);
1724        assert!(resp.headers.get("content-type").is_none());
1725        assert_eq!(
1726            resp.headers.get("accept-ranges").unwrap().as_bytes(),
1727            b"bytes"
1728        );
1729
1730        // bad multipart range
1731        let mut req = gen_req();
1732        req.insert_header("Range", "bytes=1-0, 12-9, 50-40")
1733            .unwrap();
1734        let mut resp = gen_resp();
1735        resp.insert_header("Content-Encoding", "br").unwrap();
1736        resp.insert_header("Transfer-Encoding", "chunked").unwrap();
1737        let result = range_header_filter(&req, &mut resp, None);
1738        assert!(matches!(result, RangeType::Invalid));
1739        assert_eq!(resp.status.as_u16(), 416);
1740        assert!(resp.headers.get("accept-ranges").is_none());
1741        assert!(resp.headers.get("content-encoding").is_none());
1742        assert!(resp.headers.get("transfer-encoding").is_none());
1743    }
1744
1745    #[test]
1746    fn test_if_range() {
1747        const DATE: &str = "Fri, 07 Jul 2023 22:03:29 GMT";
1748        const ETAG: &str = "\"1234\"";
1749
1750        fn gen_req() -> RequestHeader {
1751            let mut req = RequestHeader::build(http::Method::GET, b"/", Some(1)).unwrap();
1752            req.append_header("Range", "bytes=0-1").unwrap();
1753            req
1754        }
1755        fn get_multipart_req() -> RequestHeader {
1756            let mut req = RequestHeader::build(http::Method::GET, b"/", Some(1)).unwrap();
1757            _ = req.append_header("Range", "bytes=0-1,3-4,6-7");
1758            req
1759        }
1760        fn gen_resp() -> ResponseHeader {
1761            let mut resp = ResponseHeader::build(200, Some(1)).unwrap();
1762            resp.append_header("Content-Length", "10").unwrap();
1763            resp.append_header("Last-Modified", DATE).unwrap();
1764            resp.append_header("ETag", ETAG).unwrap();
1765            resp
1766        }
1767
1768        // matching Last-Modified date
1769        let mut req = gen_req();
1770        req.insert_header("If-Range", DATE).unwrap();
1771        let mut resp = gen_resp();
1772        assert_eq!(
1773            RangeType::new_single(0, 2),
1774            range_header_filter(&req, &mut resp, None)
1775        );
1776
1777        // non-matching date
1778        let mut req = gen_req();
1779        req.insert_header("If-Range", "Fri, 07 Jul 2023 22:03:25 GMT")
1780            .unwrap();
1781        let mut resp = gen_resp();
1782        assert_eq!(RangeType::None, range_header_filter(&req, &mut resp, None));
1783        assert_eq!(resp.status.as_u16(), 200);
1784        assert_eq!(
1785            resp.headers.get("accept-ranges").unwrap().as_bytes(),
1786            b"bytes"
1787        );
1788
1789        // match ETag
1790        let mut req = gen_req();
1791        req.insert_header("If-Range", ETAG).unwrap();
1792        let mut resp = gen_resp();
1793        assert_eq!(
1794            RangeType::new_single(0, 2),
1795            range_header_filter(&req, &mut resp, None)
1796        );
1797        assert_eq!(resp.status.as_u16(), 206);
1798        assert!(resp.headers.get("accept-ranges").is_none());
1799
1800        // non-matching ETags do not result in range
1801        let mut req = gen_req();
1802        req.insert_header("If-Range", "\"4567\"").unwrap();
1803        let mut resp = gen_resp();
1804        assert_eq!(RangeType::None, range_header_filter(&req, &mut resp, None));
1805        assert_eq!(resp.status.as_u16(), 200);
1806        assert_eq!(
1807            resp.headers.get("accept-ranges").unwrap().as_bytes(),
1808            b"bytes"
1809        );
1810
1811        let mut req = gen_req();
1812        req.insert_header("If-Range", "1234").unwrap();
1813        let mut resp = gen_resp();
1814        assert_eq!(RangeType::None, range_header_filter(&req, &mut resp, None));
1815        assert_eq!(resp.status.as_u16(), 200);
1816        assert_eq!(
1817            resp.headers.get("accept-ranges").unwrap().as_bytes(),
1818            b"bytes"
1819        );
1820
1821        // multipart range with If-Range
1822        let mut req = get_multipart_req();
1823        req.insert_header("If-Range", DATE).unwrap();
1824        let mut resp = gen_resp();
1825        let result = range_header_filter(&req, &mut resp, None);
1826        assert!(matches!(result, RangeType::Multi(_)));
1827        assert_eq!(resp.status.as_u16(), 206);
1828        assert!(resp.headers.get("accept-ranges").is_none());
1829
1830        // multipart with matching ETag
1831        let req = get_multipart_req();
1832        let mut resp = gen_resp();
1833        assert!(matches!(
1834            range_header_filter(&req, &mut resp, None),
1835            RangeType::Multi(_)
1836        ));
1837
1838        // multipart with non-matching If-Range
1839        let mut req = get_multipart_req();
1840        req.insert_header("If-Range", "\"wrong\"").unwrap();
1841        let mut resp = gen_resp();
1842        assert_eq!(RangeType::None, range_header_filter(&req, &mut resp, None));
1843        assert_eq!(resp.status.as_u16(), 200);
1844        assert_eq!(
1845            resp.headers.get("accept-ranges").unwrap().as_bytes(),
1846            b"bytes"
1847        );
1848    }
1849
1850    pub struct RangeBodyFilter {
1851        pub range: RangeType,
1852        current: usize,
1853        multipart_idx: Option<usize>,
1854        cache_multipart_idx: Option<usize>,
1855    }
1856
1857    impl Default for RangeBodyFilter {
1858        fn default() -> Self {
1859            Self::new()
1860        }
1861    }
1862
1863    impl RangeBodyFilter {
1864        pub fn new() -> Self {
1865            RangeBodyFilter {
1866                range: RangeType::None,
1867                current: 0,
1868                multipart_idx: None,
1869                cache_multipart_idx: None,
1870            }
1871        }
1872
1873        pub fn new_range(range: RangeType) -> Self {
1874            RangeBodyFilter {
1875                multipart_idx: matches!(range, RangeType::Multi(_)).then_some(0),
1876                range,
1877                ..Default::default()
1878            }
1879        }
1880
1881        pub fn is_multipart_range(&self) -> bool {
1882            matches!(self.range, RangeType::Multi(_))
1883        }
1884
1885        /// Whether we should expect the cache body reader to seek again
1886        /// for a different range.
1887        pub fn should_cache_seek_again(&self) -> bool {
1888            match &self.range {
1889                RangeType::Multi(multipart_info) => self
1890                    .cache_multipart_idx
1891                    .is_some_and(|idx| idx != multipart_info.ranges.len() - 1),
1892                _ => false,
1893            }
1894        }
1895
1896        /// Returns the next multipart range to seek for the cache body reader.
1897        ///
1898        /// The body filter and seekable cache reader must advance through each
1899        /// part together. If they diverge, the response body cannot be served
1900        /// correctly; report an internal error rather than panic.
1901        pub fn next_cache_multipart_range(&mut self) -> Result<Range<usize>> {
1902            let RangeType::Multi(multipart_info) = &self.range else {
1903                return Error::e_explain(
1904                    InternalError,
1905                    "tried to advance cache multipart range on a non-multipart response",
1906                );
1907            };
1908
1909            let cache_multipart_idx = self.cache_multipart_idx.map_or(0, |idx| idx + 1);
1910            let Some(multipart_idx) = self.multipart_idx else {
1911                return Error::e_explain(
1912                    InternalError,
1913                    "multipart response is missing body filter progress state",
1914                );
1915            };
1916            if multipart_idx != cache_multipart_idx {
1917                return Error::e_explain(
1918                    InternalError,
1919                    format!(
1920                        "cache multipart progress mismatch: body_filter_idx={multipart_idx}, cache_reader_idx={cache_multipart_idx}, ranges={}",
1921                        multipart_info.ranges.len(),
1922                    ),
1923                );
1924            }
1925
1926            let Some(range) = multipart_info.ranges.get(cache_multipart_idx).cloned() else {
1927                return Error::e_explain(
1928                    InternalError,
1929                    "cache multipart reader advanced past the final requested range",
1930                );
1931            };
1932            self.cache_multipart_idx = Some(cache_multipart_idx);
1933            Ok(range)
1934        }
1935
1936        pub fn set_current_cursor(&mut self, current: usize) {
1937            self.current = current;
1938        }
1939
1940        pub fn set(&mut self, range: RangeType) {
1941            self.multipart_idx = matches!(range, RangeType::Multi(_)).then_some(0);
1942            self.range = range;
1943        }
1944
1945        // Emit final boundary footer for multipart requests
1946        pub fn finalize(&self, boundary: &String) -> Option<Bytes> {
1947            if let RangeType::Multi(_) = self.range {
1948                Some(Bytes::from(format!("\r\n--{boundary}--\r\n")))
1949            } else {
1950                None
1951            }
1952        }
1953
1954        pub fn filter_body(&mut self, data: Option<Bytes>) -> Option<Bytes> {
1955            match &self.range {
1956                RangeType::None => data,
1957                RangeType::Invalid => None,
1958                RangeType::Single(r) => {
1959                    let current = self.current;
1960                    self.current += data.as_ref().map_or(0, |d| d.len());
1961                    data.and_then(|d| Self::filter_range_data(r.start, r.end, current, d))
1962                }
1963
1964                RangeType::Multi(_) => {
1965                    let data = data?;
1966                    let current = self.current;
1967                    let data_len = data.len();
1968                    self.current += data_len;
1969                    self.filter_multi_range_body(data, current, data_len)
1970                }
1971            }
1972        }
1973
1974        fn filter_range_data(
1975            start: usize,
1976            end: usize,
1977            current: usize,
1978            data: Bytes,
1979        ) -> Option<Bytes> {
1980            if current + data.len() < start || current >= end {
1981                // if the current data is out side the desired range, just drop the data
1982                None
1983            } else if current >= start && current + data.len() <= end {
1984                // all data is within the slice
1985                Some(data)
1986            } else {
1987                // data:  current........current+data.len()
1988                // range: start...........end
1989                let slice_start = start.saturating_sub(current);
1990                let slice_end = std::cmp::min(data.len(), end - current);
1991                Some(data.slice(slice_start..slice_end))
1992            }
1993        }
1994
1995        // Returns the multipart header for a given range
1996        fn build_multipart_header(
1997            &self,
1998            range: &Range<usize>,
1999            boundary: &str,
2000            total_length: &usize,
2001            content_type: Option<&str>,
2002        ) -> Bytes {
2003            Bytes::from(format!(
2004                "\r\n--{}\r\n{}Content-Range: bytes {}-{}/{}\r\n\r\n",
2005                boundary,
2006                content_type.map_or(String::new(), |ct| format!("Content-Type: {ct}\r\n")),
2007                range.start,
2008                range.end - 1,
2009                total_length
2010            ))
2011        }
2012
2013        // Return true if chunk includes the start of the given range
2014        fn current_chunk_includes_range_start(
2015            &self,
2016            range: &Range<usize>,
2017            current: usize,
2018            data_len: usize,
2019        ) -> bool {
2020            range.start >= current && range.start < current + data_len
2021        }
2022
2023        // Return true if chunk includes the end of the given range
2024        fn current_chunk_includes_range_end(
2025            &self,
2026            range: &Range<usize>,
2027            current: usize,
2028            data_len: usize,
2029        ) -> bool {
2030            range.end > current && range.end <= current + data_len
2031        }
2032
2033        fn filter_multi_range_body(
2034            &mut self,
2035            data: Bytes,
2036            current: usize,
2037            data_len: usize,
2038        ) -> Option<Bytes> {
2039            let mut result = BytesMut::new();
2040
2041            let RangeType::Multi(multi_part_info) = &self.range else {
2042                return None;
2043            };
2044
2045            let multipart_idx = self.multipart_idx.expect("must be set on multirange");
2046            let final_range = multi_part_info.ranges.last()?;
2047
2048            let (_, remaining_ranges) = multi_part_info.ranges.as_slice().split_at(multipart_idx);
2049            // NOTE: current invariant is that the multipart info ranges are disjoint ascending
2050            // this code is invalid if this invariant is not upheld
2051            for range in remaining_ranges {
2052                if let Some(sliced) =
2053                    Self::filter_range_data(range.start, range.end, current, data.clone())
2054                {
2055                    if self.current_chunk_includes_range_start(range, current, data_len) {
2056                        result.extend_from_slice(&self.build_multipart_header(
2057                            range,
2058                            multi_part_info.boundary.as_ref(),
2059                            &multi_part_info.total_length,
2060                            multi_part_info.content_type.as_deref(),
2061                        ));
2062                    }
2063                    // Emit the actual data bytes
2064                    result.extend_from_slice(&sliced);
2065                    if self.current_chunk_includes_range_end(range, current, data_len) {
2066                        // If this was the last range, we should emit the final footer too
2067                        if range == final_range {
2068                            if let Some(final_chunk) = self.finalize(&multi_part_info.boundary) {
2069                                result.extend_from_slice(&final_chunk);
2070                            }
2071                        }
2072                        // done with this range
2073                        self.multipart_idx = Some(self.multipart_idx.expect("must be set") + 1);
2074                    }
2075                } else {
2076                    // no part of the data was within this range,
2077                    // so lower bound of this range (and remaining ranges) must be
2078                    // > current + data_len
2079                    break;
2080                }
2081            }
2082            if result.is_empty() {
2083                None
2084            } else {
2085                Some(result.freeze())
2086            }
2087        }
2088    }
2089
2090    #[test]
2091    fn test_range_body_filter_single() {
2092        let mut body_filter = RangeBodyFilter::new_range(RangeType::None);
2093        assert_eq!(body_filter.filter_body(Some("123".into())).unwrap(), "123");
2094
2095        let mut body_filter = RangeBodyFilter::new_range(RangeType::Invalid);
2096        assert!(body_filter.filter_body(Some("123".into())).is_none());
2097
2098        let mut body_filter = RangeBodyFilter::new_range(RangeType::new_single(0, 1));
2099        assert_eq!(body_filter.filter_body(Some("012".into())).unwrap(), "0");
2100        assert!(body_filter.filter_body(Some("345".into())).is_none());
2101
2102        let mut body_filter = RangeBodyFilter::new_range(RangeType::new_single(4, 6));
2103        assert!(body_filter.filter_body(Some("012".into())).is_none());
2104        assert_eq!(body_filter.filter_body(Some("345".into())).unwrap(), "45");
2105        assert!(body_filter.filter_body(Some("678".into())).is_none());
2106
2107        let mut body_filter = RangeBodyFilter::new_range(RangeType::new_single(1, 7));
2108        assert_eq!(body_filter.filter_body(Some("012".into())).unwrap(), "12");
2109        assert_eq!(body_filter.filter_body(Some("345".into())).unwrap(), "345");
2110        assert_eq!(body_filter.filter_body(Some("678".into())).unwrap(), "6");
2111    }
2112
2113    #[test]
2114    fn test_range_body_filter_multipart() {
2115        // Test #1 - Test multipart ranges from 1 chunk
2116        let data = Bytes::from("0123456789");
2117        let ranges = vec![0..3, 6..9];
2118        let content_length = data.len();
2119        let mut body_filter = RangeBodyFilter::new();
2120        body_filter.set(RangeType::new_multi(ranges.clone()));
2121
2122        body_filter
2123            .range
2124            .update_multirange_info(content_length, None);
2125
2126        let multi_range_info = body_filter
2127            .range
2128            .get_multirange_info()
2129            .cloned()
2130            .expect("Multipart Ranges should have MultiPartInfo struct");
2131
2132        // Pass the whole body in one chunk
2133        let output = body_filter.filter_body(Some(data)).unwrap();
2134        let footer = body_filter.finalize(&multi_range_info.boundary).unwrap();
2135
2136        // Convert to String so that we can inspect whole response
2137        let output_str = str::from_utf8(&output).unwrap();
2138        let final_boundary = str::from_utf8(&footer).unwrap();
2139        let boundary = &multi_range_info.boundary;
2140
2141        // Check part headers
2142        for (i, range) in ranges.iter().enumerate() {
2143            let header = &format!(
2144                "--{}\r\nContent-Range: bytes {}-{}/{}\r\n\r\n",
2145                boundary,
2146                range.start,
2147                range.end - 1,
2148                content_length
2149            );
2150            assert!(
2151                output_str.contains(header),
2152                "Missing part header {} in multipart body",
2153                i
2154            );
2155            // Check body matches
2156            let expected_body = &"0123456789"[range.clone()];
2157            assert!(
2158                output_str.contains(expected_body),
2159                "Missing body {} for range {:?}",
2160                expected_body,
2161                range
2162            )
2163        }
2164        // Check the final boundary footer
2165        assert_eq!(final_boundary, format!("\r\n--{}--\r\n", boundary));
2166
2167        // Test #2 - Test multipart ranges from multiple chunks
2168        let full_body = b"0123456789";
2169        let ranges = vec![0..2, 4..6, 8..9];
2170        let content_length = full_body.len();
2171        let content_type = "text/plain".to_string();
2172        let mut body_filter = RangeBodyFilter::new();
2173        body_filter.set(RangeType::new_multi(ranges.clone()));
2174
2175        body_filter
2176            .range
2177            .update_multirange_info(content_length, Some(content_type.clone()));
2178
2179        let multi_range_info = body_filter
2180            .range
2181            .get_multirange_info()
2182            .cloned()
2183            .expect("Multipart Ranges should have MultiPartInfo struct");
2184
2185        // Split the body into 4 chunks
2186        let chunk1 = Bytes::from_static(b"012");
2187        let chunk2 = Bytes::from_static(b"345");
2188        let chunk3 = Bytes::from_static(b"678");
2189        let chunk4 = Bytes::from_static(b"9");
2190
2191        let mut collected_bytes = BytesMut::new();
2192        for chunk in [chunk1, chunk2, chunk3, chunk4] {
2193            if let Some(filtered) = body_filter.filter_body(Some(chunk)) {
2194                collected_bytes.extend_from_slice(&filtered);
2195            }
2196        }
2197        if let Some(final_boundary) = body_filter.finalize(&multi_range_info.boundary) {
2198            collected_bytes.extend_from_slice(&final_boundary);
2199        }
2200
2201        let output_str = str::from_utf8(&collected_bytes).unwrap();
2202        let boundary = multi_range_info.boundary;
2203
2204        for (i, range) in ranges.iter().enumerate() {
2205            let header = &format!(
2206                "--{}\r\nContent-Type: {}\r\nContent-Range: bytes {}-{}/{}\r\n\r\n",
2207                boundary,
2208                content_type,
2209                range.start,
2210                range.end - 1,
2211                content_length
2212            );
2213            let expected_body = &full_body[range.clone()];
2214            let expected_output = format!("{}{}", header, str::from_utf8(expected_body).unwrap());
2215
2216            assert!(
2217                output_str.contains(&expected_output),
2218                "Missing or malformed part {} in multipart body. \n Expected: \n{}\n Got: \n{}",
2219                i,
2220                expected_output,
2221                output_str
2222            )
2223        }
2224
2225        assert!(
2226            output_str.ends_with(&format!("\r\n--{}--\r\n", boundary)),
2227            "Missing final boundary"
2228        );
2229
2230        // Test #3 - Test multipart ranges from multiple chunks, with ranges spanning chunks
2231        let full_body = b"abcdefghijkl";
2232        let ranges = vec![2..7, 9..11];
2233        let content_length = full_body.len();
2234        let content_type = "application/octet-stream".to_string();
2235        let mut body_filter = RangeBodyFilter::new();
2236        body_filter.set(RangeType::new_multi(ranges.clone()));
2237
2238        body_filter
2239            .range
2240            .update_multirange_info(content_length, Some(content_type.clone()));
2241
2242        let multi_range_info = body_filter
2243            .range
2244            .clone()
2245            .get_multirange_info()
2246            .cloned()
2247            .expect("Multipart Ranges should have MultiPartInfo struct");
2248
2249        // Split the body into 4 chunks
2250        let chunk1 = Bytes::from_static(b"abc");
2251        let chunk2 = Bytes::from_static(b"def");
2252        let chunk3 = Bytes::from_static(b"ghi");
2253        let chunk4 = Bytes::from_static(b"jkl");
2254
2255        let mut collected_bytes = BytesMut::new();
2256        for chunk in [chunk1, chunk2, chunk3, chunk4] {
2257            if let Some(filtered) = body_filter.filter_body(Some(chunk)) {
2258                collected_bytes.extend_from_slice(&filtered);
2259            }
2260        }
2261        if let Some(final_boundary) = body_filter.finalize(&multi_range_info.boundary) {
2262            collected_bytes.extend_from_slice(&final_boundary);
2263        }
2264
2265        let output_str = str::from_utf8(&collected_bytes).unwrap();
2266        let boundary = &multi_range_info.boundary;
2267
2268        let header1 = &format!(
2269            "--{}\r\nContent-Type: {}\r\nContent-Range: bytes {}-{}/{}\r\n\r\n",
2270            boundary,
2271            content_type,
2272            ranges[0].start,
2273            ranges[0].end - 1,
2274            content_length
2275        );
2276        let header2 = &format!(
2277            "--{}\r\nContent-Type: {}\r\nContent-Range: bytes {}-{}/{}\r\n\r\n",
2278            boundary,
2279            content_type,
2280            ranges[1].start,
2281            ranges[1].end - 1,
2282            content_length
2283        );
2284
2285        assert!(output_str.contains(header1));
2286        assert!(output_str.contains(header2));
2287
2288        let expected_body_slices = ["cdefg", "jk"];
2289
2290        assert!(
2291            output_str.contains(expected_body_slices[0]),
2292            "Missing expected sliced body {}",
2293            expected_body_slices[0]
2294        );
2295
2296        assert!(
2297            output_str.contains(expected_body_slices[1]),
2298            "Missing expected sliced body {}",
2299            expected_body_slices[1]
2300        );
2301
2302        assert!(
2303            output_str.ends_with(&format!("\r\n--{}--\r\n", boundary)),
2304            "Missing final boundary"
2305        );
2306    }
2307
2308    #[test]
2309    fn test_cache_multipart_advance_errors_when_reader_ends_part_early() {
2310        let ranges = vec![0..10, 20..30];
2311        let mut body_filter = RangeBodyFilter::new_range(RangeType::new_multi(ranges));
2312
2313        let first = body_filter.next_cache_multipart_range().unwrap();
2314        assert_eq!(first, 0..10);
2315        body_filter.set_current_cursor(first.start);
2316
2317        // The cache reader yielded only a prefix of the selected part before
2318        // reporting EOF. The filter therefore has not advanced past part 0.
2319        assert!(body_filter
2320            .filter_body(Some(Bytes::from_static(b"01234")))
2321            .is_some());
2322
2323        let err = body_filter.next_cache_multipart_range().unwrap_err();
2324        assert_eq!(err.etype(), &InternalError);
2325        assert!(err
2326            .to_string()
2327            .contains("cache multipart progress mismatch: body_filter_idx=0, cache_reader_idx=1"));
2328    }
2329
2330    #[test]
2331    fn test_cache_multipart_advance_errors_when_reader_overreads_part() {
2332        let ranges = vec![0..2, 4..6, 8..10];
2333        let mut body_filter = RangeBodyFilter::new_range(RangeType::new_multi(ranges));
2334
2335        let first = body_filter.next_cache_multipart_range().unwrap();
2336        assert_eq!(first, 0..2);
2337        body_filter.set_current_cursor(first.start);
2338
2339        // A seekable reader is expected to stop at the selected part's end.
2340        // This chunk spans all requested parts and advances the filter beyond
2341        // what the reader's seek state records.
2342        assert!(body_filter
2343            .filter_body(Some(Bytes::from_static(b"0123456789")))
2344            .is_some());
2345
2346        let err = body_filter.next_cache_multipart_range().unwrap_err();
2347        assert_eq!(err.etype(), &InternalError);
2348        assert!(err
2349            .to_string()
2350            .contains("cache multipart progress mismatch: body_filter_idx=3, cache_reader_idx=1"));
2351    }
2352}
2353
2354// a state machine for proxy logic to tell when to use cache in the case of
2355// miss/revalidation/error.
2356#[derive(Debug)]
2357pub(crate) enum ServeFromCache {
2358    // not using cache
2359    Off,
2360    // should serve cache header
2361    CacheHeader,
2362    // should serve cache header only
2363    CacheHeaderOnly,
2364    // should serve cache header only but upstream response should be admitted to cache
2365    CacheHeaderOnlyMiss,
2366    // should serve cache body with a bool to indicate if it has already called seek on the hit handler
2367    CacheBody(bool),
2368    // should serve cache header but upstream response should be admitted to cache
2369    // This is the starting state for misses, which go to CacheBodyMiss or
2370    // CacheHeaderOnlyMiss before ending at DoneMiss
2371    CacheHeaderMiss,
2372    // should serve cache body but upstream response should be admitted to cache, bool to indicate seek status
2373    CacheBodyMiss(bool),
2374    // done serving cache body
2375    Done,
2376    // done serving cache body, but upstream response should continue to be admitted to cache
2377    DoneMiss,
2378}
2379
2380impl ServeFromCache {
2381    pub fn new() -> Self {
2382        Self::Off
2383    }
2384
2385    pub fn is_on(&self) -> bool {
2386        !matches!(self, Self::Off)
2387    }
2388
2389    pub fn is_miss(&self) -> bool {
2390        matches!(
2391            self,
2392            Self::CacheHeaderMiss
2393                | Self::CacheHeaderOnlyMiss
2394                | Self::CacheBodyMiss(_)
2395                | Self::DoneMiss
2396        )
2397    }
2398
2399    pub fn is_miss_header(&self) -> bool {
2400        // NOTE: this check is for checking if miss was just enabled, so it is excluding
2401        // HeaderOnlyMiss
2402        matches!(self, Self::CacheHeaderMiss)
2403    }
2404
2405    pub fn is_miss_body(&self) -> bool {
2406        matches!(self, Self::CacheBodyMiss(_))
2407    }
2408
2409    pub fn should_discard_upstream(&self) -> bool {
2410        self.is_on() && !self.is_miss()
2411    }
2412
2413    pub fn should_send_to_downstream(&self) -> bool {
2414        !self.is_on()
2415    }
2416
2417    pub fn enable(&mut self) {
2418        *self = Self::CacheHeader;
2419    }
2420
2421    pub fn enable_miss(&mut self) {
2422        if !self.is_on() {
2423            *self = Self::CacheHeaderMiss;
2424        }
2425    }
2426
2427    pub fn enable_header_only(&mut self) {
2428        match self {
2429            Self::CacheBody(_) => *self = Self::Done, // TODO: make sure no body is read yet
2430            Self::CacheBodyMiss(_) => *self = Self::DoneMiss,
2431            _ => {
2432                if self.is_miss() {
2433                    *self = Self::CacheHeaderOnlyMiss;
2434                } else {
2435                    *self = Self::CacheHeaderOnly;
2436                }
2437            }
2438        }
2439    }
2440
2441    // This function is (best effort) cancel-safe to be used in select
2442    pub async fn next_http_task(
2443        &mut self,
2444        cache: &mut HttpCache,
2445        range: &mut RangeBodyFilter,
2446        upgraded: bool,
2447    ) -> Result<HttpTask> {
2448        fn body_task(data: Bytes, upgraded: bool) -> HttpTask {
2449            if upgraded {
2450                HttpTask::UpgradedBody(Some(data), false)
2451            } else {
2452                HttpTask::Body(Some(data), false)
2453            }
2454        }
2455
2456        if !cache.enabled() {
2457            // Cache is disabled due to internal error
2458            // TODO: if nothing is sent to eyeball yet, figure out a way to recovery by
2459            // fetching from upstream
2460            return Error::e_explain(InternalError, "Cache disabled");
2461        }
2462        match self {
2463            Self::Off => panic!("ProxyUseCache not enabled"),
2464            Self::CacheHeader => {
2465                *self = Self::CacheBody(true);
2466                Ok(HttpTask::Header(cache_hit_header(cache), false)) // false for now
2467            }
2468            Self::CacheHeaderMiss => {
2469                *self = Self::CacheBodyMiss(true);
2470                Ok(HttpTask::Header(cache_hit_header(cache), false)) // false for now
2471            }
2472            Self::CacheHeaderOnly => {
2473                *self = Self::Done;
2474                Ok(HttpTask::Header(cache_hit_header(cache), true))
2475            }
2476            Self::CacheHeaderOnlyMiss => {
2477                *self = Self::DoneMiss;
2478                Ok(HttpTask::Header(cache_hit_header(cache), true))
2479            }
2480            Self::CacheBody(should_seek) => {
2481                log::trace!("cache body should seek: {should_seek}");
2482                if *should_seek {
2483                    self.maybe_seek_hit_handler(cache, range)?;
2484                }
2485                loop {
2486                    if let Some(b) = cache.hit_handler().read_body().await? {
2487                        return Ok(body_task(b, upgraded));
2488                    }
2489                    // EOF from hit handler for body requested
2490                    // if multipart, then seek again
2491                    if range.should_cache_seek_again() {
2492                        self.maybe_seek_hit_handler(cache, range)?;
2493                    } else {
2494                        *self = Self::Done;
2495                        return Ok(HttpTask::Done);
2496                    }
2497                }
2498            }
2499            Self::CacheBodyMiss(should_seek) => {
2500                if *should_seek {
2501                    self.maybe_seek_miss_handler(cache, range)?;
2502                }
2503                // safety: caller of enable_miss() call it only if the async_body_reader exist
2504                loop {
2505                    if let Some(b) = cache.miss_body_reader().unwrap().read_body().await? {
2506                        return Ok(body_task(b, upgraded));
2507                    } else {
2508                        // EOF from hit handler for body requested
2509                        // if multipart, then seek again
2510                        if range.should_cache_seek_again() {
2511                            self.maybe_seek_miss_handler(cache, range)?;
2512                        } else {
2513                            *self = Self::DoneMiss;
2514                            return Ok(HttpTask::Done);
2515                        }
2516                    }
2517                }
2518            }
2519            Self::Done => Ok(HttpTask::Done),
2520            Self::DoneMiss => Ok(HttpTask::Done),
2521        }
2522    }
2523
2524    fn maybe_seek_miss_handler(
2525        &mut self,
2526        cache: &mut HttpCache,
2527        range_filter: &mut RangeBodyFilter,
2528    ) -> Result<()> {
2529        match &range_filter.range {
2530            RangeType::Single(range) => {
2531                // safety: called only if the async_body_reader exists
2532                if cache.miss_body_reader().unwrap().can_seek() {
2533                    cache
2534                        .miss_body_reader()
2535                        // safety: called only if the async_body_reader exists
2536                        .unwrap()
2537                        .seek(range.start, Some(range.end))
2538                        .or_err(InternalError, "cannot seek miss handler")?;
2539                    // Because the miss body reader is seeking, we no longer need the
2540                    // RangeBodyFilter's help to return the requested byte range.
2541                    range_filter.range = RangeType::None;
2542                }
2543            }
2544            // safety: called only if the async_body_reader exists
2545            RangeType::Multi(_info) if cache.miss_body_reader().unwrap().can_seek_multipart() => {
2546                let range = range_filter.next_cache_multipart_range()?;
2547                cache
2548                    .miss_body_reader()
2549                    .unwrap()
2550                    .seek_multipart(range.start, Some(range.end))
2551                    .or_err(InternalError, "cannot seek hit handler for multirange")?;
2552                // we still need RangeBodyFilter's help to transform the byte
2553                // range into a multipart response.
2554                range_filter.set_current_cursor(range.start);
2555            }
2556            _ => {}
2557        }
2558
2559        *self = Self::CacheBodyMiss(false);
2560        Ok(())
2561    }
2562
2563    fn maybe_seek_hit_handler(
2564        &mut self,
2565        cache: &mut HttpCache,
2566        range_filter: &mut RangeBodyFilter,
2567    ) -> Result<()> {
2568        match &range_filter.range {
2569            RangeType::Single(range) => {
2570                if cache.hit_handler().can_seek() {
2571                    cache
2572                        .hit_handler()
2573                        .seek(range.start, Some(range.end))
2574                        .or_err(InternalError, "cannot seek hit handler")?;
2575                    // Because the hit handler is seeking, we no longer need the
2576                    // RangeBodyFilter's help to return the requested byte range.
2577                    range_filter.range = RangeType::None;
2578                }
2579            }
2580            RangeType::Multi(_info) if cache.hit_handler().can_seek_multipart() => {
2581                let range = range_filter.next_cache_multipart_range()?;
2582                cache
2583                    .hit_handler()
2584                    .seek_multipart(range.start, Some(range.end))
2585                    .or_err(InternalError, "cannot seek hit handler for multirange")?;
2586                // we still need RangeBodyFilter's help to transform the byte
2587                // range into a multipart response.
2588                range_filter.set_current_cursor(range.start);
2589            }
2590            _ => {}
2591        }
2592        *self = Self::CacheBody(false);
2593        Ok(())
2594    }
2595}
2596
2597#[cfg(test)]
2598mod tests {
2599    use super::*;
2600    use pingora_cache::{
2601        predictor::{CacheablePredictor, Predictor},
2602        CacheKey, CacheMeta, CachePhase, MemCache, RespCacheable,
2603    };
2604    use pingora_http::ResponseHeader;
2605    use std::sync::{Arc, LazyLock};
2606    use std::time::Duration;
2607    use tokio::io::AsyncWriteExt;
2608
2609    /// Stands in for a caller-defined `NoCacheReason::Custom` that has nothing to do with size.
2610    const AUTHORIZATION_HEADER: &str = "AuthorizationHeader";
2611
2612    static CACHE_STORAGE: LazyLock<MemCache> = LazyLock::new(MemCache::new);
2613    static CACHE_PREDICTOR: LazyLock<Predictor<1>> = LazyLock::new(|| Predictor::new(10, None));
2614
2615    struct TestProxy;
2616
2617    #[async_trait]
2618    impl ProxyHttp for TestProxy {
2619        type CTX = ();
2620
2621        fn new_ctx(&self) -> Self::CTX {}
2622
2623        async fn upstream_peer(
2624            &self,
2625            _session: &mut Session,
2626            _ctx: &mut Self::CTX,
2627        ) -> Result<Box<HttpPeer>> {
2628            unreachable!("test drives cache_http_task directly")
2629        }
2630
2631        fn response_cache_filter(
2632            &self,
2633            _session: &Session,
2634            resp: &ResponseHeader,
2635            _ctx: &mut Self::CTX,
2636        ) -> Result<RespCacheable> {
2637            let now = SystemTime::now();
2638            Ok(RespCacheable::Cacheable(CacheMeta::new(
2639                now + Duration::from_secs(60),
2640                now,
2641                0,
2642                0,
2643                resp.clone(),
2644            )))
2645        }
2646    }
2647
2648    /// Build a session that the predictor has bypassed, exactly as `proxy_cache` would:
2649    /// the key is remembered as uncacheable for `reason`, so lookup is skipped.
2650    async fn bypassed_session(
2651        key: CacheKey,
2652        reason: NoCacheReason,
2653        max_file_size: usize,
2654    ) -> Session {
2655        let (mut client, server) = tokio::io::duplex(1024);
2656        client
2657            .write_all(b"GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")
2658            .await
2659            .expect("test request should be written");
2660
2661        let mut session = Session::new_h1(Box::new(server) as pingora_core::protocols::Stream);
2662        session
2663            .read_request()
2664            .await
2665            .expect("test request should parse");
2666        session
2667            .cache
2668            .enable(&*CACHE_STORAGE, None, Some(&*CACHE_PREDICTOR), None, None);
2669        session.cache.set_cache_key(key.clone());
2670        session.cache.set_max_file_size_bytes(max_file_size);
2671
2672        CACHE_PREDICTOR.mark_uncacheable(&key, reason);
2673        assert!(
2674            !session.cache.cacheable_prediction(),
2675            "predictor should bypass a key it just marked uncacheable"
2676        );
2677        session.cache.bypass();
2678        session
2679    }
2680
2681    /// A cacheable 200 with no `Content-Length`, i.e. an origin that chunks the body.
2682    fn chunked_cacheable_response() -> ResponseHeader {
2683        let mut resp = ResponseHeader::build(StatusCode::OK, None).unwrap();
2684        resp.insert_header(http::header::CACHE_CONTROL, "max-age=60")
2685            .unwrap();
2686        resp
2687    }
2688
2689    async fn run_header_task(session: &mut Session, resp: ResponseHeader) {
2690        let proxy = HttpProxy::new(TestProxy, Arc::new(ServerConf::default()));
2691        proxy
2692            .cache_http_task(
2693                session,
2694                &HttpTask::Header(Box::new(resp), true),
2695                &mut (),
2696                &mut ServeFromCache::new(),
2697            )
2698            .await
2699            .expect("cache_http_task should succeed");
2700    }
2701
2702    /// A predictor bypass that had nothing to do with size must not be reported as
2703    /// PredictedResponseTooLarge, and must not cost the request a wasted bypass.
2704    #[tokio::test]
2705    async fn non_size_bypass_admits_and_clears_predictor() {
2706        for reason in [
2707            NoCacheReason::Custom(AUTHORIZATION_HEADER),
2708            NoCacheReason::OriginNotCache,
2709        ] {
2710            let key = CacheKey::new(format!("/non-size-bypass/{}", reason.as_str()), "");
2711            let mut session = bypassed_session(key.clone(), reason, 1024).await;
2712
2713            run_header_task(&mut session, chunked_cacheable_response()).await;
2714
2715            assert_eq!(
2716                session.cache.phase(),
2717                CachePhase::Miss,
2718                "bypass remembered for {reason:?} should admit this response, got {:?}",
2719                session.cache.phase()
2720            );
2721            assert!(
2722                CACHE_PREDICTOR.cacheable_prediction(&key),
2723                "bypass remembered for {reason:?} should be cleared once the response came back cacheable"
2724            );
2725        }
2726    }
2727
2728    /// The deferral only exists to protect against a response that already blew the size
2729    /// limit once. That case still defers, and still reports the reason it acted on.
2730    #[tokio::test]
2731    async fn size_bypass_defers_when_length_is_unknown() {
2732        let key = CacheKey::new("/size-bypass-chunked", "");
2733        let mut session =
2734            bypassed_session(key.clone(), NoCacheReason::ResponseTooLarge, 1024).await;
2735
2736        run_header_task(&mut session, chunked_cacheable_response()).await;
2737
2738        assert_eq!(
2739            session.cache.phase(),
2740            CachePhase::Disabled(NoCacheReason::PredictedResponseTooLarge)
2741        );
2742        assert!(
2743            !CACHE_PREDICTOR.cacheable_prediction(&key),
2744            "the response body has not been measured yet, so the key stays marked"
2745        );
2746    }
2747
2748    /// With a Content-Length the size is known up front, so even a size bypass can admit.
2749    #[tokio::test]
2750    async fn size_bypass_admits_when_length_is_known() {
2751        let key = CacheKey::new("/size-bypass-with-length", "");
2752        let mut session =
2753            bypassed_session(key.clone(), NoCacheReason::ResponseTooLarge, 1024).await;
2754
2755        let mut resp = chunked_cacheable_response();
2756        resp.insert_header(CONTENT_LENGTH, "5").unwrap();
2757        run_header_task(&mut session, resp).await;
2758
2759        assert_eq!(session.cache.phase(), CachePhase::Miss);
2760    }
2761
2762    /// Without a predictor reason there is nothing to act on, so the conservative
2763    /// deferral stays in place.
2764    #[tokio::test]
2765    async fn unknown_bypass_reason_stays_conservative() {
2766        let key = CacheKey::new("/unknown-bypass-reason", "");
2767        // Bypass without the predictor remembering anything, e.g. a caller that bypassed
2768        // for its own reasons.
2769        let (mut client, server) = tokio::io::duplex(1024);
2770        client
2771            .write_all(b"GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")
2772            .await
2773            .unwrap();
2774        let mut session = Session::new_h1(Box::new(server) as pingora_core::protocols::Stream);
2775        session.read_request().await.unwrap();
2776        session
2777            .cache
2778            .enable(&*CACHE_STORAGE, None, Some(&*CACHE_PREDICTOR), None, None);
2779        session.cache.set_cache_key(key);
2780        session.cache.set_max_file_size_bytes(1024);
2781        session.cache.bypass();
2782        assert_eq!(session.cache.predicted_uncacheable_reason(), None);
2783
2784        run_header_task(&mut session, chunked_cacheable_response()).await;
2785
2786        assert_eq!(
2787            session.cache.phase(),
2788            CachePhase::Disabled(NoCacheReason::PredictedResponseTooLarge)
2789        );
2790    }
2791}