read_to_timeout/
lib.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
/// The [std::io::Read](std::io::Read) trait implements many read operations,
/// but it doesn't contain a simple read method where timeout is the expected
/// behaviour
///
/// This trait provides [read_to_timeout](ReadToTimeout::read_to_timeout) and
/// [read_to_pattern_or_timeout](ReadToTimeout::read_to_pattern_or_timeout)
/// that are implemented for all types that implements
/// [std::io::Read](std::io::Read)
pub trait ReadToTimeout {
    /// Similar to [`read_to_end`](std::io::Read::read_to_end)
    ///
    /// But when timeout, instead of returning error, this function returns Ok(bytes_read)
    fn read_to_timeout(&mut self, buf: &mut Vec<u8>) -> std::io::Result<usize>;

    /// Similar to [`read_to_timeout`](ReadToTimeout::read_to_timeout)
    ///
    /// But when a specified pattern is reached, return Ok(bytes_read) immediately
    ///
    /// # Note
    /// If the provided buffer is non-empty, while **at least one byte** must be
    /// read before any pattern match, it is possible for pattern to match on
    /// old bytes.
    fn read_to_pattern_or_timeout(
        &mut self,
        buf: &mut Vec<u8>,
        pattern: &[u8],
    ) -> std::io::Result<usize>;
}

// impl ReadTimeout for all T that impl Read
impl<T: std::io::Read + ?Sized> ReadToTimeout for T {
    fn read_to_timeout(&mut self, buf: &mut Vec<u8>) -> std::io::Result<usize> {
        let old_len = buf.len();

        loop {
            let mut read_buf = [0];
            match self.read_exact(&mut read_buf) {
                Ok(_) => {
                    buf.push(read_buf[0]);
                }
                Err(e) => {
                    return match e.kind() {
                        std::io::ErrorKind::TimedOut => Ok(buf.len() - old_len),
                        _ => Err(e),
                    }
                }
            }
        }
    }

    fn read_to_pattern_or_timeout(
        &mut self,
        buf: &mut Vec<u8>,
        pattern: &[u8],
    ) -> std::io::Result<usize> {
        let old_len = buf.len();

        loop {
            let mut byte = [0];
            match self.read_exact(&mut byte) {
                Ok(_) => {
                    buf.push(byte[0]);
                    if buf.len() >= pattern.len()
                        && &buf[(buf.len() - pattern.len())..] == pattern
                    {
                        break Ok(buf.len() - old_len);
                    }
                }
                Err(err) => match err.kind() {
                    std::io::ErrorKind::TimedOut => {
                        break Ok(buf.len() - old_len);
                    }
                    _ => {
                        break Err(err);
                    }
                },
            }
        }
    }
}

#[cfg(test)]
mod test {
    use crate::ReadToTimeout;

    #[test]
    fn read_pattern() {
        let mut bytes: &[u8] = &[1, 2, 3, 4, 5, 6, 7, 6, 5, 4, 3, 2, 1];

        let mut buffer = Vec::new();
        assert_eq!(
            bytes.read_to_pattern_or_timeout(&mut buffer, &[7, 6]).unwrap(),
            8
        );
        assert_eq!(buffer, &[1, 2, 3, 4, 5, 6, 7, 6]);
    }
}