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

use ::std::io::{Read, Write, Result};
use ::std::io::{Error, ErrorKind};

// TODO: get rid of this
const DIGIT_LIMIT: usize = 64;

pub trait ReadNetstring {
    fn read_netstring(&mut self) -> Result<String>;
}

pub trait WriteNetstring {
    fn write_netstring<S: AsRef<str>>(&mut self, value: S) -> Result<()>;
}

impl<R: Read> ReadNetstring for R {
    fn read_netstring(&mut self) -> Result<String> {
        let ln = try!(read_length(self));
        let mut data = vec![0u8;ln];

        let mut offset = 0usize;
        let mut done = false;

        while !done {
            let r = try!(self.read(data[offset..].as_mut()));
            offset = offset + r;
            if r == 0 || offset == ln {
                done = true;
            }
        }

        // TODO: there has to be a a cleaner way to do this...
        // read delimiter ","
        let mut t = vec![0u8].into_boxed_slice();
        try!(self.read(t[..].as_mut()));

        // Verify delimiter
        if t[0] != b","[0] {
            return Err(Error::new(ErrorKind::InvalidData, "Expected `,` delimiter."));
        }

        // return utf8 string
        match String::from_utf8(data) {
            Ok(s) => Ok(s),
            Err(err) => Err(Error::new(ErrorKind::InvalidData, err)),
        }
    }
}

impl<W: Write> WriteNetstring for W {
    fn write_netstring<S: AsRef<str>>(&mut self, value: S) -> Result<()> {
        let value = value.as_ref();
        let s = format!("{}:{},", value.len(), value);
        try!(self.write(s.as_bytes()));
        Ok(())
    }
}

fn read_length<R: Read>(r: &mut R) -> Result<usize> {
    let mut t = [0u8; DIGIT_LIMIT];
    let mut current = 0usize;
    let mut done = false;
    while !done {
        let r = try!(r.read(t[current..current+1].as_mut()));
        if r == 0 {
            return Err(Error::new(ErrorKind::ConnectionAborted, "Connection closed by target."));
        }
        if t[current] == 0x3A {
            done = true;
        } else {
            current += 1;
        }
    }

    let s = match String::from_utf8(t[..current].to_vec()) {
                Ok(s)  => s,
                Err(err) => return Err(Error::new(ErrorKind::InvalidData, err)),
            };

    let ln = match s.parse::<u64>() {
                Ok(x)  => x,
                Err(err) => return Err(Error::new(ErrorKind::InvalidData, err)),
                };


    Ok(ln as usize)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn basic_read() {
        let mut raw = "5:hello,".as_bytes();
        let x = raw.read_netstring().unwrap();
        assert_eq!("hello", x);
    }

    #[test]
    fn basic_write() {
        let mut raw = vec![];
        let _ = raw.write_netstring("hello").unwrap();
        assert_eq!(raw, b"5:hello,");
    }

    #[test]
    #[should_panic(expected="Expected `,` delimiter.")]
    fn invalid_delimiter() {
        let mut raw = "5:hello?".as_bytes();
        raw.read_netstring().unwrap();
    }

    #[test]
    #[should_panic(expected="Expected `,` delimiter.")]
    fn longer() {
        let mut raw = "10:hello,".as_bytes();
        raw.read_netstring().unwrap();
    }

    #[test]
    #[should_panic(expected="Expected `,` delimiter.")]
    fn shorter() {
        let mut raw = "2:hello,".as_bytes();
        raw.read_netstring().unwrap();
    }
}