ippusb_bridge/src/util.rs

// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::io::{self, BufRead};

/// Read from `reader` until `delimiter` is seen or EOF is reached.
/// Returns read data.
pub fn read_until_delimiter(reader: &mut dyn BufRead, delimiter: &[u8]) -> io::Result<Vec<u8>> {
    let mut result: Vec<u8> = Vec::new();
    loop {
        let buf = match reader.fill_buf() {
            Ok(buf) => buf,
            Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
            Err(e) => return Err(e),
        };

        if buf.is_empty() {
            return Ok(result);
        }

        // First check if our delimiter spans the old buffer and the new buffer.
        for split in 1..delimiter.len() {
            let (first_delimiter, second_delimiter) = delimiter.split_at(split);
            if first_delimiter.len() > result.len() || second_delimiter.len() > buf.len() {
                continue;
            }

            let first = result.get(result.len() - first_delimiter.len()..);
            let second = buf.get(..second_delimiter.len());
            if let (Some(first), Some(second)) = (first, second) {
                if first == first_delimiter && second == second_delimiter {
                    result.extend_from_slice(second);
                    reader.consume(second_delimiter.len());
                    return Ok(result);
                }
            }
        }

        // Then check if our delimiter occurs in the new buffer.
        if let Some(i) = buf
            .windows(delimiter.len())
            .position(|window| window == delimiter)
        {
            result.extend_from_slice(&buf[..i + delimiter.len()]);
            reader.consume(i + delimiter.len());
            return Ok(result);
        }

        // Otherwise just copy the entire buffer into result.
        let consumed = buf.len();
        result.extend_from_slice(&buf);
        reader.consume(consumed);
    }
}

#[cfg(test)]
mod tests {
    use crate::util::read_until_delimiter;
    use std::io::{BufReader, Cursor};

    #[test]
    fn test_read_until_delimiter() {
        let mut source = Cursor::new(&b"abdcdef"[..]);
        let v = read_until_delimiter(&mut source, b"20").unwrap();
        assert_eq!(v, b"abdcdef");

        let mut source = Cursor::new(&b"abdcdef"[..]);
        let v = read_until_delimiter(&mut source, b"de").unwrap();
        assert_eq!(v, b"abdcde");

        let mut source = Cursor::new(&b"abdcdef"[..]);
        let v = read_until_delimiter(&mut source, b"dc").unwrap();
        assert_eq!(v, b"abdc");

        let mut source = Cursor::new(&b"abdcdef"[..]);
        let v = read_until_delimiter(&mut source, b"abd").unwrap();
        assert_eq!(v, b"abd");

        let mut source = BufReader::with_capacity(2, Cursor::new(&b"abdcdeffegh"[..]));
        let v = read_until_delimiter(&mut source, b"bdc").unwrap();
        assert_eq!(v, b"abdc");

        let v = read_until_delimiter(&mut source, b"ef").unwrap();
        assert_eq!(v, b"def");

        let v = read_until_delimiter(&mut source, b"g").unwrap();
        assert_eq!(v, b"feg");
    }
}