1use 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 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 where
42 SV: ProxyHttp + Send + Sync + 'static,
43 SV::CTX: Send + Sync,
44 {
45 if let Err(e) = self.inner.request_cache_filter(session, ctx) {
47 warn!(
49 "Fail to request_cache_filter: {e}, {}",
50 self.inner.request_summary(session, ctx)
51 );
52 }
53
54 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 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 if self.inner.is_purge(session, ctx) {
73 return self.proxy_purge(session, ctx).await;
74 }
75
76 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 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 let cache_key = session.cache.cache_key();
97 if let Some(variance) = cache_key.variance_bin() {
98 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 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 if !session.cache.cache_vary_lookup(variance, &meta) {
112 continue;
115 }
116 } }
118
119 let now = SystemTime::now();
124 let is_fresh = meta.is_fresh(now);
125 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 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 meta.disable_serve_stale();
150 HitStatus::ForceExpired
151 }
152 Ok(Some(ForcedFreshness::ForceExpiredServeStale { expired_at })) => {
153 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 session.cache.cache_found(meta, handler, hit_status);
168 }
169
170 if hit_status_opt.is_none_or(HitStatus::is_treated_as_miss) {
171 if !session.cache.enabled() {
173 break None;
175 } else if session.cache.is_cache_locked() {
176 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 let hit_status = hit_status_opt.expect("None case handled as miss");
199
200 if !hit_status.is_fresh() {
201 if session.cache.is_cache_locked() {
203 if let Some(write_lock) = session
205 .subrequest_ctx
206 .as_mut()
207 .and_then(|ctx| ctx.take_write_lock())
208 {
209 session.cache.set_write_lock(write_lock);
211 session.cache.tag_as_subrequest();
212 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 session.cache.set_stale_updating();
234 } else if session.cache.is_cache_lock_writer() {
235 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 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 session.cache.set_stale_updating();
255 } else {
256 break None;
258 }
259 } else {
260 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 break Some((reuse, err));
274 }
275 Err(e) => {
276 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 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 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 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 let header_only = header_only || matches!(range_type, RangeType::Invalid);
334 debug!("header: {header:?}");
335
336 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 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 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 return (true, Some(e));
383 }
384 }
385 debug!("finished sending cached header to downstream");
386
387 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 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 continue;
443 }
444 Ok(false) => {
445 Error::explain(
452 InternalError,
453 "hit handler cannot seek for multipart again",
454 )
455 }
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) => { }
479 Err(e) => {
480 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 return (false, Some(e));
491 }
492
493 if !end && body.as_ref().is_none_or(|b| b.is_empty()) {
494 continue;
497 }
498
499 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 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 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 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 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 use_cache.enable_header_only();
567 }
568 }
569
570 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 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 let bypassed_over_size =
621 match session.cache.predicted_uncacheable_reason() {
622 Some(reason) => reason == NoCacheReason::ResponseTooLarge,
623 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 fill_cache = false;
652 session.cache.disable(NoCacheReason::Deferred);
653 }
654 }
655
656 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 session.ignore_downstream_range = true;
672 }
673 }
674 }
678 }
679 if fill_cache {
680 let req_header = session.req_header();
681 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 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() .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 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 match data {
718 Some(d) => {
719 if session.cache.enabled() {
720 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 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(_) => {} HttpTask::Done => {
759 if session.cache.enabled() {
760 self.finish_miss_handler_best_effort(session, ctx).await;
761 }
762 }
763 HttpTask::Failed(_) => {
764 }
766 }
767 Ok(())
768 }
769
770 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 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 return true;
805 }
806 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 let old_meta = session.cache.maybe_cache_meta().unwrap(); 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 warn!("revalidate_cache_meta failed {e:?}");
829 }
830 }
831 Ok(Uncacheable(reason)) => {
832 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 warn!("Error {e:?} response_cache_filter during revalidation");
853 session.cache.revalidate_uncacheable(
854 merged_header,
855 NoCacheReason::InternalError,
856 );
857 }
859 }
860 true
862 } else {
863 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 if !session.cache.can_serve_stale_error()
875 || session.response_written().is_some()
876 {
877 return false;
878 }
879
880 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 session
893 .cache
894 .release_write_lock(NoCacheReason::UpstreamError);
895 true
896 } else {
897 false
898 }
899 } else {
900 false }
902 }
903 _ => false, }
905 }
906
907 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 if !session.cache.can_serve_stale_error() {
921 return None;
922 }
923
924 if session.response_written().is_some() {
927 return None;
928 }
929
930 if !self.inner.should_serve_stale(session, ctx, Some(error)) {
932 return None;
933 }
934
935 warn!(
937 "Fail to proxy: {}, serving stale, {}",
938 error,
939 self.inner.request_summary(session, ctx)
940 );
941
942 session
944 .cache
945 .release_write_lock(NoCacheReason::UpstreamError);
946
947 Some(self.proxy_cache_hit(session, ctx).await)
948 }
949
950 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 LockWaitOutcome::Done => true,
964 LockWaitOutcome::TransientError => true,
966 LockWaitOutcome::GiveUp => {
968 session.cache.disable(NoCacheReason::CacheLockGiveUp);
969 false
970 }
971 LockWaitOutcome::Abandoned { reason, .. } => {
974 session.cache.disable(reason);
975 false
976 }
977 LockWaitOutcome::Dangling => {
979 warn!(
981 "Dangling cache lock, {}",
982 self.inner.request_summary(session, ctx)
983 );
984 true
985 }
986 LockWaitOutcome::WaitTimeout => {
989 warn!(
990 "Cache lock timeout, {}",
991 self.inner.request_summary(session, ctx)
992 );
993 session.cache.disable(NoCacheReason::CacheLockTimeout);
994 false
996 }
997 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 let no_body = matches!(header.status.as_u16(), 204 | 304);
1037
1038 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 header.set_version(Version::HTTP_11);
1051
1052 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
1065pub mod range_filter {
1067 use super::*;
1068 use bytes::BytesMut;
1069 use http::header::*;
1070 use std::ops::Range;
1071
1072 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 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 let range_str = match str::from_utf8(range) {
1090 Ok(s) => s,
1091 Err(_) => return RangeType::None,
1092 };
1093
1094 let mut parts = range_str.splitn(2, "=");
1096
1097 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 return RangeType::None;
1106 };
1107
1108 if ranges_str.trim().is_empty() {
1111 return RangeType::Invalid;
1112 }
1113
1114 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 return RangeType::None;
1122 }
1123 }
1124 }
1125 let mut ranges: Vec<Range<usize>> = Vec::with_capacity(range_count);
1126
1127 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 continue;
1148 }
1149 let end = std::cmp::min(end.unwrap_or(content_length - 1), content_length - 1) + 1;
1153 if end <= start {
1154 continue;
1156 }
1157 start..end
1158 } else {
1159 if let Some(end) = end {
1162 if content_length >= end {
1163 (content_length - end)..content_length
1164 } else {
1165 0..content_length
1167 }
1168 } else {
1169 continue;
1171 }
1172 };
1173 if range.start < last_range_end {
1176 return RangeType::None;
1177 }
1178 last_range_end = range.end;
1179 ranges.push(range);
1180 }
1181
1182 if ranges.is_empty() {
1192 RangeType::Invalid
1194 } else if ranges.len() == 1 {
1195 RangeType::Single(ranges[0].clone()) } 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 #[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 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 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 assert_eq!(
1286 parse_range_header(b"bytes=0-,-2", 10, None),
1287 RangeType::None,
1288 );
1289 assert!(parse_range_header(b"bytes=0-,-2", 10, None)
1291 .get_multirange_info()
1292 .is_none());
1293 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 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 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 assert_eq!(
1321 parse_range_header(b"bytes=5-0, 20-15, 30-25", 200, None),
1322 RangeType::Invalid
1323 );
1324
1325 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 let ranges = generate_range_header(201);
1343 assert_eq!(
1344 parse_range_header(&ranges, 1000, Some(200)),
1345 RangeType::None
1346 )
1347 }
1348
1349 #[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 pub fn new(ranges: Vec<Range<usize>>) -> Self {
1362 Self {
1363 ranges,
1364 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 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 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 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 #[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 pub fn range_header_filter(
1449 req: &RequestHeader,
1450 resp: &mut ResponseHeader,
1451 max_multipart_ranges: Option<usize>,
1452 ) -> RangeType {
1453 if resp.status != StatusCode::OK {
1457 return RangeType::None;
1458 }
1459
1460 let Some(content_length_bytes) = resp.headers.get(CONTENT_LENGTH) else {
1463 return RangeType::None;
1464 };
1465 let Some(content_length) = parse_number(content_length_bytes.as_bytes()) else {
1467 return RangeType::None;
1468 };
1469
1470 fn request_range_type(
1475 req: &RequestHeader,
1476 resp: &ResponseHeader,
1477 content_length: usize,
1478 max_multipart_ranges: Option<usize>,
1479 ) -> RangeType {
1480 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 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 resp.insert_header(&ACCEPT_RANGES, "bytes").unwrap();
1524 }
1525 RangeType::Single(r) => {
1526 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), )
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 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 ), )
1560 .unwrap();
1561 resp.remove_header(&CONTENT_RANGE);
1562 }
1563 RangeType::Invalid => {
1564 resp.set_status(StatusCode::RANGE_NOT_SATISFIABLE).unwrap();
1566 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 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 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 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 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 assert!(resp.headers.get("accept-ranges").is_none());
1646
1647 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 #[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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 None
1983 } else if current >= start && current + data.len() <= end {
1984 Some(data)
1986 } else {
1987 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 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 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 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 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 result.extend_from_slice(&sliced);
2065 if self.current_chunk_includes_range_end(range, current, data_len) {
2066 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 self.multipart_idx = Some(self.multipart_idx.expect("must be set") + 1);
2074 }
2075 } else {
2076 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 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 let output = body_filter.filter_body(Some(data)).unwrap();
2134 let footer = body_filter.finalize(&multi_range_info.boundary).unwrap();
2135
2136 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 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 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 assert_eq!(final_boundary, format!("\r\n--{}--\r\n", boundary));
2166
2167 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 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 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 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 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 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#[derive(Debug)]
2357pub(crate) enum ServeFromCache {
2358 Off,
2360 CacheHeader,
2362 CacheHeaderOnly,
2364 CacheHeaderOnlyMiss,
2366 CacheBody(bool),
2368 CacheHeaderMiss,
2372 CacheBodyMiss(bool),
2374 Done,
2376 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 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, 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 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 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)) }
2468 Self::CacheHeaderMiss => {
2469 *self = Self::CacheBodyMiss(true);
2470 Ok(HttpTask::Header(cache_hit_header(cache), false)) }
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 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 loop {
2505 if let Some(b) = cache.miss_body_reader().unwrap().read_body().await? {
2506 return Ok(body_task(b, upgraded));
2507 } else {
2508 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 if cache.miss_body_reader().unwrap().can_seek() {
2533 cache
2534 .miss_body_reader()
2535 .unwrap()
2537 .seek(range.start, Some(range.end))
2538 .or_err(InternalError, "cannot seek miss handler")?;
2539 range_filter.range = RangeType::None;
2542 }
2543 }
2544 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 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 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 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 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 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 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 #[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 #[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 #[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 #[tokio::test]
2765 async fn unknown_bypass_reason_stays_conservative() {
2766 let key = CacheKey::new("/unknown-bypass-reason", "");
2767 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}