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}