Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- 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() } ) );
- }
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement