Jean-Gabriel Gill-Couture
16aac0e0df
Some checks failed
Run Check Script / check (pull_request) Failing after -42s
41 lines
1001 B
Rust
41 lines
1001 B
Rust
use std::pin::Pin;
|
|
use std::task::{Context, Poll};
|
|
|
|
use tokio::io::AsyncWrite;
|
|
|
|
pub struct BoundedWriter {
|
|
pub bytes: Vec<u8>,
|
|
limit: u64,
|
|
}
|
|
|
|
impl BoundedWriter {
|
|
pub fn new(limit: u64) -> Self {
|
|
Self {
|
|
bytes: Vec::new(),
|
|
limit,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl AsyncWrite for BoundedWriter {
|
|
fn poll_write(
|
|
mut self: Pin<&mut Self>,
|
|
_cx: &mut Context<'_>,
|
|
buffer: &[u8],
|
|
) -> Poll<std::io::Result<usize>> {
|
|
if self.bytes.len() as u64 + buffer.len() as u64 > self.limit {
|
|
return Poll::Ready(Err(std::io::Error::other("artifact exceeds byte limit")));
|
|
}
|
|
self.bytes.extend_from_slice(buffer);
|
|
Poll::Ready(Ok(buffer.len()))
|
|
}
|
|
|
|
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
|
|
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|