use crate::storage::HandleMiss;
use crate::MissHandler;
use async_trait::async_trait;
use bytes::Bytes;
use pingora_error::{Error, ErrorType};
pub struct MaxFileSizeMissHandler {
inner: MissHandler,
max_file_size_bytes: usize,
bytes_written: usize,
}
impl MaxFileSizeMissHandler {
pub fn new(inner: MissHandler, max_file_size_bytes: usize) -> MaxFileSizeMissHandler {
MaxFileSizeMissHandler {
inner,
max_file_size_bytes,
bytes_written: 0,
}
}
}
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<()> {
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
}
}