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
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
//! IoBuffer library
//!
//! Copyright 2020 Metaswitch Networks
//!
//! Memory-based buffer which implements Write and Read traits.

use std::cmp;
use std::default::Default;
use std::io;
use std::result::Result::Ok;
use std::sync::{Arc, Mutex};

/// Simple object which implements both `std::io::Write` and `std::io::Read`.
/// Uses an internal buffer. Thread-safe and cloneable.
#[derive(Clone)]
pub struct IoBuffer {
    /// The actual shared contents of the buffer.
    inner: Arc<Mutex<IoBufferInner>>,
}

/// A simple read/write buffer.
struct IoBufferInner {
    /// All content that has been written so far.
    buf: Vec<u8>,

    /// The current reading cursor.
    pos: usize,
}

impl IoBuffer {
    /// Create a new empty buffer.
    pub fn new() -> Self {
        IoBuffer {
            inner: Arc::new(Mutex::new(IoBufferInner {
                buf: vec![],
                pos: 0,
            })), // LCOV_EXCL_LINE kcov bug?
        }
    }
}

impl Default for IoBuffer {
    fn default() -> Self {
        Self::new()
    }
}

impl io::Write for IoBuffer {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        // Trivial implementation - add all the data onto the end.
        let mut lock = self.inner.lock().expect("lock poisoned");
        lock.buf.extend_from_slice(buf);
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        // Nothing to do.
        Ok(())
    }
}

impl io::Read for IoBuffer {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        // Trivial implementation - read all the available data that can fit, and
        // advance the cursor.
        let mut lock = self.inner.lock().expect("lock poisoned");
        let len = cmp::min(lock.buf.len() - lock.pos, buf.len());
        let pos = lock.pos;
        buf[0..len].copy_from_slice(&lock.buf[pos..pos + len]);
        lock.pos += len;
        Ok(len)
    }
}

impl IoBuffer {
    /// Read a full line (terminated by the indicated byte), if any.
    /// Any partial line is left unread. The terminator is not included
    /// in the result.
    pub fn read_full_line(&mut self, terminator: u8) -> Option<Vec<u8>> {
        let mut lock = self.inner.lock().expect("lock poisoned");
        let mut p = lock.buf[lock.pos..].split(|c| *c == terminator);
        match p.next() {
            Some(line) => {
                match p.next() {
                    Some(_rest) => {
                        let line = line.to_vec();
                        lock.pos += line.len() + 1;
                        Some(line)
                    }
                    None => None, // incomplete line
                }
            }
            None => None, // no data
        }
    }

    /// Iterator of full lines (as returned by `read_full_line`).
    pub fn lines(&mut self) -> impl Iterator<Item = Vec<u8>> {
        let mut buf = self.clone();
        std::iter::from_fn(move || buf.read_full_line(b'\n'))
    }
}

#[cfg(test)]
mod tests {
    #![allow(unused_imports)] // these are required, but there's a warning for some reason
    use super::IoBuffer;
    use std::io::{Read, Write};

    #[test]
    fn test_simple_buffer_usage() {
        let s1 = "This is some unexciting test data";
        let s2 = "This is some more unexciting test data";

        let mut buf = IoBuffer::new();

        let mut dest = Vec::new();
        let rc = buf.read_to_end(&mut dest).unwrap();
        assert!(rc == 0, "{}", rc);

        write!(buf, "{}", s1).unwrap();
        write!(buf, "{}", s2).unwrap();
        buf.flush().unwrap();

        let rc = buf.read_to_end(&mut dest).unwrap();
        assert!(
            rc == s1.len() + s2.len(),
            "{} != {}",
            rc,
            s1.len() + s2.len()
        );
        let s_out = String::from_utf8(dest).unwrap();
        assert!(s_out == (s1.to_string() + s2), "{}", s_out);
    }

    fn next_full_line(buf: &mut IoBuffer) -> Option<String> {
        buf.read_full_line(b'\n')
            .map(|x| String::from_utf8(x).unwrap())
    }

    #[test]
    fn test_partial_lines() {
        let mut buf = IoBuffer::new();

        assert_eq!(None, next_full_line(&mut buf));

        write!(buf, "{}", "abc").unwrap();
        assert_eq!(None, next_full_line(&mut buf));

        write!(buf, "{}", "d\n").unwrap();
        assert_eq!(Some("abcd".to_string()), next_full_line(&mut buf));
        assert_eq!(None, next_full_line(&mut buf));

        write!(buf, "{}", "e\n\nfghi\nj").unwrap();
        assert_eq!(Some("e".to_string()), next_full_line(&mut buf));
        assert_eq!(Some("".to_string()), next_full_line(&mut buf));
        assert_eq!(Some("fghi".to_string()), next_full_line(&mut buf));
        assert_eq!(None, next_full_line(&mut buf));

        write!(buf, "{}", "\n").unwrap();
        assert_eq!(Some("j".to_string()), next_full_line(&mut buf));
        assert_eq!(None, next_full_line(&mut buf));
        assert_eq!(None, next_full_line(&mut buf));
    }
}