read-to-timeout 0.3.1

/// 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 is implemented for all types that implements
/// [std::io::Read](std::io::Read), provides
/// [read_to_timeout](ReadToTimeout::read_to_timeout),
/// [read_to_timeout_or_bytes](ReadToTimeout::read_to_timeout_or_bytes), and
/// [read_to_timeout_or_pattern](ReadToTimeout::read_to_timeout_or_pattern)
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 the received byte count reaches `max_byte`, this function
    /// returns `Ok(bytes_read)` immediately
    fn read_to_timeout_or_bytes(
        &mut self,
        buf: &mut Vec<u8>,
        max_byte: usize,
    ) -> std::io::Result<usize>;

    /// Similar to [read_to_timeout](ReadToTimeout::read_to_timeout)
    ///
    /// But when a specified pattern is reached, this function 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_timeout_or_pattern(
        &mut self,
        buf: &mut Vec<u8>,
        pattern: &[u8],
    ) -> std::io::Result<usize>;

    /// **DEPRECATED**
    ///
    /// Use [read_to_timeout_or_pattern](ReadToTimeout::read_to_timeout_or_pattern)
    /// instead
    #[deprecated]
    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_timeout_or_bytes(
        &mut self,
        buf: &mut Vec<u8>,
        max_byte: usize,
    ) -> 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]);

                    if (buf.len() - old_len) == max_byte {
                        return Ok(buf.len() - old_len);
                    }
                }
                Err(e) => {
                    return match e.kind() {
                        std::io::ErrorKind::TimedOut => Ok(buf.len() - old_len),
                        _ => Err(e),
                    }
                }
            }
        }
    }

    fn read_to_timeout_or_pattern(
        &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);
                    }
                },
            }
        }
    }

    fn read_to_pattern_or_timeout(
        &mut self,
        buf: &mut Vec<u8>,
        pattern: &[u8],
    ) -> std::io::Result<usize> {
        self.read_to_timeout_or_pattern(buf, pattern)
    }
}

#[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_timeout_or_pattern(&mut buffer, &[7, 6]).unwrap(),
            8
        );
        assert_eq!(buffer, &[1, 2, 3, 4, 5, 6, 7, 6]);
        
        assert_eq!(bytes.read_to_timeout_or_bytes(&mut buffer, 4).unwrap(), 4);
        assert_eq!(buffer, &[1, 2, 3, 4, 5, 6, 7, 6, 5, 4, 3, 2]);
    }
}