ctf_pwn/io/cache/
read.rs

1use crate::io::cache::CacheReader;
2use crate::io::AsyncCacheRead;
3
4use std::borrow::BorrowMut;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7use tokio::io;
8use tokio::io::{AsyncRead, ReadBuf};
9
10impl<R: AsyncRead> AsyncRead for CacheReader<R> {
11    fn poll_read(
12        self: Pin<&mut Self>,
13        cx: &mut Context<'_>,
14        buf: &mut ReadBuf<'_>,
15    ) -> Poll<Result<(), io::Error>> {
16        let mut this = self.project();
17        if !this.cache.is_empty() {
18            let remaining = usize::min(buf.remaining(), this.cache.len());
19            buf.put_slice(&this.cache[..remaining]);
20            this.cache.drain(..remaining);
21            return Poll::Ready(Ok(()));
22        }
23        this.reader.poll_read(cx, buf)
24    }
25}
26
27impl<R: AsyncRead> AsyncCacheRead for CacheReader<R> {
28    fn poll_reader(
29        self: Pin<&mut Self>,
30        cx: &mut Context<'_>,
31        buf: &mut ReadBuf<'_>,
32    ) -> Poll<std::io::Result<()>> {
33        self.project().reader.poll_read(cx, buf)
34    }
35
36    fn consume(self: Pin<&mut Self>, amt: usize) {
37        self.project().borrow_mut().cache.drain(..amt);
38    }
39
40    fn restore(self: Pin<&mut Self>, data: &[u8]) {
41        self.project().borrow_mut().cache.extend_from_slice(data)
42    }
43}