ctf_pwn/io/util/timeout/
read_until_regex_timeout.rs

1use crate::io::{get_deadline, AsyncCacheRead};
2use crate::io::{read_until_regex, ReadUntilRegex};
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 ReadUntilRegexTimeout<'a, R: ?Sized> {
20        #[pin]
21        read_until_regex: ReadUntilRegex<'a,R>,
22        deadline: Instant,
23        #[pin]
24        _pin: PhantomPinned,
25    }
26}
27
28pub(crate) fn read_until_regex_timeout<'a, R>(
29    reader: &'a mut R,
30    pattern: &str,
31    buf: &'a mut Vec<u8>,
32    timeout: Duration,
33) -> Result<ReadUntilRegexTimeout<'a, R>, regex::Error>
34where
35    R: AsyncCacheRead + ?Sized + Unpin,
36{
37    let read_until_regex = read_until_regex(reader, pattern, buf)?;
38    let deadline = get_deadline(timeout);
39    Ok(ReadUntilRegexTimeout {
40        read_until_regex,
41        deadline,
42        _pin: PhantomPinned,
43    })
44}
45
46impl<R: AsyncCacheRead + ?Sized + Unpin> Future for ReadUntilRegexTimeout<'_, R> {
47    type Output = io::Result<(usize, 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_regex.poll(cx)
56    }
57}
ctf_pwn/io/util/timeout/read_until_regex_timeout.rs

ctf_pwn/io/util/timeout/
read_until_regex_timeout.rs
