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>(&self, buf: (&'a mut [u8], usize)) -> 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>(&self, buf: (&'a mut [u8], usize)) -> Result<(&'a mut [u8], usize), GenError> {
        do_gen!(buf,
            gen_be_u8!(0x00) >>
            gen_be_u8!(self.key.len() as u8) >>
            gen_slice!(&self.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>(&self, buf: (&'a mut [u8], usize)) -> Result<(&'a mut [u8], usize), GenError> {
        do_gen!(buf,
            gen_be_u8!(0x01) >>
            gen_be_u8!(self.key.len() as u8) >>
            gen_slice!(&self.key) >>
            gen_be_u8!(self.value.len() as u8) >>
            gen_slice!(&self.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>(&self, buf: (&'a mut [u8], usize)) -> Result<(&'a mut [u8], usize), GenError> {
        do_gen!(buf,
            gen_be_u8!(self.id) >>
            gen_be_u8!(self.payload.len() as u8) >>
            gen_slice!(&self.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>(&self, buf: (&'a mut [u8], usize)) -> Result<(&'a mut [u8], usize), GenError> {
        match self {
            &Packet::Get(ref inner) => inner.to_bytes(buf),
            &Packet::Set(ref inner) => inner.to_bytes(buf),
            &Packet::Data(ref inner) => inner.to_bytes(buf),
        }
    }
}

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

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

impl ToBytes for EncryptedPacket {
    fn to_bytes<'a>(&self, buf: (&'a mut [u8], usize)) -> Result<(&'a mut [u8], usize), GenError> {
        do_gen!(buf,
            gen_be_u16!(self.payload.len()) >>
            gen_slice!(self.payload.as_slice())
        )
    }
}


/// implements tokio-io's Decoder and Encoder to deal with Packet
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)).map(|tup| tup.1) {
            Ok(produced) => {
                buf.extend_from_slice(&stack_buf[..produced]);
                // TODO memzero stack_buf ?
                trace!("serialized packet: {} bytes", produced);
                Ok(())
            },
            Err(e) => {
                Err(io::Error::new(io::ErrorKind::Other, format!("encode error: {:?}", e)))
            }
        }
    }
}

/// implements tokio-io's Decoder and Encoder to deal with EncryptedPacket
pub struct EncryptedCodec;

impl Decoder for EncryptedCodec {
    type Item = EncryptedPacket;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        let (consumed, packet) = match EncryptedPacket::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)
            }
        };

        buf.split_to(consumed);

        Ok(Some(packet))
    }
}

impl Encoder for EncryptedCodec {
    type Item = EncryptedPacket;
    type Error = io::Error;

    fn encode(&mut self, packet: Self::Item, buf: &mut BytesMut) -> Result<(), Self::Error> {
        let mut stack_buf = [0; 2050];
        match packet.to_bytes((&mut stack_buf, 0)).map(|tup| tup.1) {
            Ok(produced) => {
                buf.extend_from_slice(&stack_buf[..produced]);
                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::*;
    fn check_packet(packet: Packet) {
        let mut buf = BytesMut::new();
        let mut codec = Codec {};
        codec.encode(packet.clone() , &mut buf ).expect("Should encode");
        // zip buf
        // create EncryptedPacket with payload = zipped buf
        // encrypted_codec.encode(encrypted_packet, buf)
        // -- send via network
        // encrypted_packet = encrypted_codec.decode(buf)
        // unzip buf
        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() } ) );
    }
}