use std::io;
use bytes::BytesMut;
use byteorder::{LittleEndian, ByteOrder};
use tokio_io::codec::{Decoder, Encoder};
use messages::{HEADER_LENGTH, MessageBuffer, RawMessage};
use super::error::other_error;
#[derive(Debug)]
pub struct MessagesCodec {
max_message_len: u32,
}
impl MessagesCodec {
pub fn new(max_message_len: u32) -> MessagesCodec {
MessagesCodec { max_message_len }
}
}
impl Decoder for MessagesCodec {
type Item = RawMessage;
type Error = io::Error;
fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, io::Error> {
if buf.len() < HEADER_LENGTH {
return Ok(None);
}
let total_len = LittleEndian::read_u32(&buf[6..10]) as usize;
if total_len as u32 > self.max_message_len {
return Err(other_error(format!(
"Received message is too long: {}, maximum allowed length is {} bytes",
total_len,
self.max_message_len,
)));
}
if total_len < HEADER_LENGTH {
return Err(other_error(format!(
"Received malicious message with insufficient \
size in header: {}, expected header size {}",
total_len,
HEADER_LENGTH
)));
}
if buf.len() >= total_len {
let data = buf.split_to(total_len).to_vec();
let raw = RawMessage::new(MessageBuffer::from_vec(data));
return Ok(Some(raw));
}
Ok(None)
}
}
impl Encoder for MessagesCodec {
type Item = RawMessage;
type Error = io::Error;
fn encode(&mut self, msg: Self::Item, buf: &mut BytesMut) -> io::Result<()> {
buf.extend_from_slice(msg.as_ref());
Ok(())
}
}
#[cfg(test)]
mod test {
use super::MessagesCodec;
use messages::{MessageBuffer, RawMessage};
use bytes::BytesMut;
use tokio_io::codec::Decoder;
#[test]
fn decode_message_valid_header_size() {
let data = vec![0u8, 0, 0, 0, 0, 0, 10, 0, 0, 0];
let mut bytes: BytesMut = data.as_slice().into();
let mut codec = MessagesCodec { max_message_len: 10000 };
match codec.decode(&mut bytes) {
Ok(Some(ref r)) if r == &RawMessage::new(MessageBuffer::from_vec(data)) => {}
_ => panic!("Wrong input"),
};
}
#[test]
fn decode_message_small_size_in_header() {
let data = vec![0u8, 0, 0, 0, 0, 0, 0, 0, 0, 0];
let mut bytes: BytesMut = data.as_slice().into();
let mut codec = MessagesCodec { max_message_len: 10000 };
assert!(codec.decode(&mut bytes).is_err());
}
}