pingora_cache/
max_file_size.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
// Copyright 2024 Cloudflare, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Set limit on the largest size to cache

use crate::storage::HandleMiss;
use crate::MissHandler;
use async_trait::async_trait;
use bytes::Bytes;
use pingora_error::{Error, ErrorType};

/// [MaxFileSizeMissHandler] wraps a MissHandler to enforce a maximum asset size that should be
/// written to the MissHandler.
///
/// This is used to enforce a maximum cache size for a request when the
/// response size is not known ahead of time (no Content-Length header). When the response size _is_
/// known ahead of time, it should be checked up front (when calculating cacheability) for efficiency.
/// Note: for requests with partial read support (where downstream reads the response from cache as
/// it is filled), this will cause the request as a whole to fail. The response will be remembered
/// as uncacheable, though, so downstream will be able to retry the request, since the cache will be
/// disabled for the retried request.
pub struct MaxFileSizeMissHandler {
    inner: MissHandler,
    max_file_size_bytes: usize,
    bytes_written: usize,
}

impl MaxFileSizeMissHandler {
    /// Create a new [MaxFileSizeMissHandler] wrapping the given [MissHandler]
    pub fn new(inner: MissHandler, max_file_size_bytes: usize) -> MaxFileSizeMissHandler {
        MaxFileSizeMissHandler {
            inner,
            max_file_size_bytes,
            bytes_written: 0,
        }
    }
}

/// Error type returned when the limit is reached.
pub const ERR_RESPONSE_TOO_LARGE: ErrorType = ErrorType::Custom("response too large");

#[async_trait]
impl HandleMiss for MaxFileSizeMissHandler {
    async fn write_body(&mut self, data: Bytes, eof: bool) -> pingora_error::Result<()> {
        // fail if writing the body would exceed the max_file_size_bytes
        if self.bytes_written + data.len() > self.max_file_size_bytes {
            return Error::e_explain(
                ERR_RESPONSE_TOO_LARGE,
                format!(
                    "writing data of size {} bytes would exceed max file size of {} bytes",
                    data.len(),
                    self.max_file_size_bytes
                ),
            );
        }

        self.bytes_written += data.len();
        self.inner.write_body(data, eof).await
    }

    async fn finish(self: Box<Self>) -> pingora_error::Result<usize> {
        self.inner.finish().await
    }

    fn streaming_write_tag(&self) -> Option<&[u8]> {
        self.inner.streaming_write_tag()
    }
}