Skip to content

Instantly share code, notes, and snippets.

@kpp kpp/test.rs
Created Sep 26, 2017

Embed
What would you like to do?
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() } ) );
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.