Created
September 29, 2017 15:09
-
-
Save kpp/b4818e4007916e833c792f6be7bb104f to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() } ) ); | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment