test.rs

use nom::*;
use cookie_factory::*;

use bytes::BytesMut;
use std::io;
use tokio_io::codec::{Decoder, Encoder};

trait FromBytes : Sized {
    fn from_bytes(i: &[u8]) -> IResult<&[u8], Self>;
}

trait ToBytes : Sized {
    fn to_bytes<'a>(buf: (&'a mut [u8], usize), packet: Self) -> Result<(&'a mut [u8], usize), GenError>;
}

#[derive(Debug, PartialEq, Clone)]
pub struct Get { key: Vec<u8> }

#[derive(Debug, PartialEq, Clone)]
pub struct Set { key: Vec<u8>, value: Vec<u8> }

#[derive(Debug, PartialEq, Clone)]
pub struct Data { id: u8, payload: Vec<u8> }

#[derive(Debug, PartialEq, Clone)]
pub enum Packet {
    Get(Get),
    Set(Set),
    Data(Data)
}

impl FromBytes for Get {
    named!(from_bytes<Get>, do_parse!(
        tag!("\x00") >>
        key: length_data!(be_u8) >>
        ( Get { key: key.to_owned() } )
    ));
}

impl ToBytes for Get {
    fn to_bytes<'a>(buf: (&'a mut [u8], usize), packet: Get) -> Result<(&'a mut [u8], usize), GenError> {
        do_gen!(buf,
            gen_be_u8!(0x00) >>
            gen_be_u8!(packet.key.len() as u8) >>
            gen_slice!(&packet.key)
        )
    }
}

impl FromBytes for Set {
    named!(from_bytes<Set>, do_parse!(
        tag!("\x01") >>
        key: length_data!(be_u8) >>
        value: length_data!(be_u8) >>
        ( Set { key: key.to_owned(), value: value.to_owned() } )
    ));
}

impl ToBytes for Set {
    fn to_bytes<'a>(buf: (&'a mut [u8], usize), packet: Set) -> Result<(&'a mut [u8], usize), GenError> {
        do_gen!(buf,
            gen_be_u8!(0x01) >>
            gen_be_u8!(packet.key.len() as u8) >>
            gen_slice!(&packet.key) >>
            gen_be_u8!(packet.value.len() as u8) >>
            gen_slice!(&packet.value)
        )
    }
}

impl FromBytes for Data {
    named!(from_bytes<Data>, do_parse!(
        id: be_u8 >>
        payload: length_data!(be_u8) >>
        ( Data { id: id, payload: payload.to_owned() } )
    ));
}

impl ToBytes for Data {
    fn to_bytes<'a>(buf: (&'a mut [u8], usize), packet: Data) -> Result<(&'a mut [u8], usize), GenError> {
        do_gen!(buf,
            gen_be_u8!(packet.id) >>
            gen_be_u8!(packet.payload.len() as u8) >>
            gen_slice!(&packet.payload)
        )
    }
}

impl FromBytes for Packet {
    named!(from_bytes<Packet>, alt!(
        map!(Get::from_bytes, Packet::Get) |
        map!(Set::from_bytes, Packet::Set) |
        map!(Data::from_bytes, Packet::Data)
    ));
}

impl ToBytes for Packet {
    fn to_bytes<'a>(buf: (&'a mut [u8], usize), packet: Packet) -> Result<(&'a mut [u8], usize), GenError> {
        match packet {
            Packet::Get(inner) => ToBytes::to_bytes(buf, inner),
            Packet::Set(inner) => ToBytes::to_bytes(buf, inner),
            Packet::Data(inner) => ToBytes::to_bytes(buf, inner),
        }
    }
}


/// implements tokio-io's Decoder and Encoder
pub struct Codec;

impl Decoder for Codec {
    type Item = Packet;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Packet>, io::Error> {
        let (consumed, packet) = match Packet::from_bytes(buf) {
            IResult::Incomplete(_) => {
                return Ok(None)
            },
            IResult::Error(e) => {
                return Err(io::Error::new(io::ErrorKind::Other, format!("decode error: {:?}", e)))
            },
            IResult::Done(i, packet) => {
                (buf.offset(i), packet)
            }
        };

        // TODO memzero buf[..consumed] ?
        buf.split_to(consumed);

        Ok(Some(packet))
    }
}

impl Encoder for Codec {
    type Item = Packet;
    type Error = io::Error;

    fn encode(&mut self, packet: Packet, buf: &mut BytesMut) -> Result<(), Self::Error> {
        let mut stack_buf = [0; 2050];
        match Packet::to_bytes((&mut stack_buf, 0), packet).map(|tup| tup.1) {
            Ok(produced) => {
                buf.extend_from_slice(&stack_buf);
                // TODO memzero stack_buf ?
                trace!("serialized packet: {} bytes", produced);
                Ok(())
            },
            Err(e) => {
                Err(io::Error::new(io::ErrorKind::Other, format!("encode error: {:?}", e)))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use ::toxcore::tcp::parse_test::*;
    #[test]
    fn get_from_bytes() {
      assert_eq!(Get::from_bytes(b"\x00"), IResult::Incomplete(Needed::Size(1)));
      assert_eq!(Get::from_bytes(b"\x00\x03abcd"), IResult::Done(&b"d"[..], Get{ key: b"abc".to_vec() }));
    }
    fn check_packet(packet: Packet) {
        let mut buf = BytesMut::new();
        let mut codec = Codec {};
        codec.encode(packet.clone() , &mut buf ).expect("Should encode");
        let res = codec.decode(&mut buf).unwrap().expect("Should decode");
        assert_eq!(packet, res);
    }
    #[test]
    fn encoder_decoder() {
        check_packet( Packet::Get( Get { key: b"abc".to_vec() } ) );
        check_packet( Packet::Set( Set { key: b"abc".to_vec(), value: b"qwe".to_vec() } ) );
    }
}