pingora_cache/max_file_size.rs
1// Copyright 2025 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
15//! Set limit on the largest size to cache
16
17use crate::storage::{HandleMiss, MissFinishType};
18use crate::MissHandler;
19use async_trait::async_trait;
20use bytes::Bytes;
21use pingora_error::{Error, ErrorType};
22
23/// [MaxFileSizeMissHandler] wraps a MissHandler to enforce a maximum asset size that should be
24/// written to the MissHandler.
25///
26/// This is used to enforce a maximum cache size for a request when the
27/// response size is not known ahead of time (no Content-Length header). When the response size _is_
28/// known ahead of time, it should be checked up front (when calculating cacheability) for efficiency.
29/// Note: for requests with partial read support (where downstream reads the response from cache as
30/// it is filled), this will cause the request as a whole to fail. The response will be remembered
31/// as uncacheable, though, so downstream will be able to retry the request, since the cache will be
32/// disabled for the retried request.
33pub struct MaxFileSizeMissHandler {
34 inner: MissHandler,
35 max_file_size_bytes: usize,
36 bytes_written: usize,
37}
38
39impl MaxFileSizeMissHandler {
40 /// Create a new [MaxFileSizeMissHandler] wrapping the given [MissHandler]
41 pub fn new(inner: MissHandler, max_file_size_bytes: usize) -> MaxFileSizeMissHandler {
42 MaxFileSizeMissHandler {
43 inner,
44 max_file_size_bytes,
45 bytes_written: 0,
46 }
47 }
48}
49
50/// Error type returned when the limit is reached.
51pub const ERR_RESPONSE_TOO_LARGE: ErrorType = ErrorType::Custom("response too large");
52
53#[async_trait]
54impl HandleMiss for MaxFileSizeMissHandler {
55 async fn write_body(&mut self, data: Bytes, eof: bool) -> pingora_error::Result<()> {
56 // fail if writing the body would exceed the max_file_size_bytes
57 if self.bytes_written + data.len() > self.max_file_size_bytes {
58 return Error::e_explain(
59 ERR_RESPONSE_TOO_LARGE,
60 format!(
61 "writing data of size {} bytes would exceed max file size of {} bytes",
62 data.len(),
63 self.max_file_size_bytes
64 ),
65 );
66 }
67
68 self.bytes_written += data.len();
69 self.inner.write_body(data, eof).await
70 }
71
72 async fn finish(self: Box<Self>) -> pingora_error::Result<MissFinishType> {
73 self.inner.finish().await
74 }
75
76 fn streaming_write_tag(&self) -> Option<&[u8]> {
77 self.inner.streaming_write_tag()
78 }
79}