ctf_pwn/io/util/timeout/
read_until_timeout.rs

1use crate::io::ReadUntil;
2use crate::io::{get_deadline, read_until, AsyncCacheRead};
3use pin_project_lite::pin_project;
4use std::future::Future;
5use std::io;
6use std::io::ErrorKind;
7use std::marker::PhantomPinned;
8use std::pin::Pin;
9use std::task::{Context, Poll};
10use std::time::Duration;
11
12use crate::io::util::timeout::timeout;
13use tokio::time::Instant;
14
15pin_project! {
16    /// The delimiter is included in the resulting vector.
17    #[derive(Debug)]
18    #[must_use = "futures do nothing unless you `.await` or poll them"]
19    pub struct ReadUntilTimeout<'a, R: ?Sized, D:AsRef<[u8]>> {
20        #[pin]
21        read_until: ReadUntil<'a,R,D>,
22        deadline: Instant,
23        #[pin]
24        _pin: PhantomPinned,
25    }
26}
27
28pub(crate) fn read_until_timeout<'a, R, D: AsRef<[u8]>>(
29    reader: &'a mut R,
30    delimiter: D,
31    buf: &'a mut Vec<u8>,
32    timeout: Duration,
33) -> ReadUntilTimeout<'a, R, D>
34where
35    R: AsyncCacheRead + ?Sized + Unpin,
36{
37    let read_until = read_until(reader, delimiter, buf);
38    let deadline = get_deadline(timeout);
39    ReadUntilTimeout {
40        read_until,
41        deadline,
42        _pin: PhantomPinned,
43    }
44}
45
46impl<R: AsyncCacheRead + ?Sized + Unpin, D: AsRef<[u8]>> Future for ReadUntilTimeout<'_, R, D> {
47    type Output = io::Result<usize>;
48
49    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
50        let me = self.project();
51        if *me.deadline < Instant::now() {
52            return Err(timeout()).into();
53        }
54
55        me.read_until.poll(cx)
56    }
57}
ctf_pwn/io/util/timeout/read_until_timeout.rs

ctf_pwn/io/util/timeout/
read_until_timeout.rs
