use std::io::{Write, Read, Cursor};
use model::error::{RdpResult, RdpErrorKind, RdpError, Error};
use byteorder::{WriteBytesExt, ReadBytesExt, LittleEndian, BigEndian};
use indexmap::IndexMap;
use std::collections::{HashSet, HashMap};
pub enum DataType<'a> {
Component(&'a Component),
Trame(&'a Trame),
U32(u32),
U16(u16),
U8(u8),
Slice(&'a [u8]),
None
}
#[macro_export]
macro_rules! cast {
($ident:path, $expr:expr) => (match $expr.visit() {
$ident(e) => Ok(e),
_ => Err(Error::RdpError(RdpError::new(RdpErrorKind::InvalidCast, "Invalid Cast")))
})
}
pub enum MessageOption {
SkipField(String),
Size(String, usize),
None
}
pub trait Message : Send {
fn write(&self, writer: &mut dyn Write) -> RdpResult<()>;
fn read(&mut self, reader: &mut dyn Read) -> RdpResult<()>;
fn length(&self) -> u64;
fn visit(&self) -> DataType;
fn options(&self) -> MessageOption;
}
impl Message for u8 {
fn write(&self, writer: &mut dyn Write) -> RdpResult<()> {
Ok(writer.write_u8(*self)?)
}
fn read(&mut self, reader: &mut dyn Read) -> RdpResult<()> {
*self = reader.read_u8()?;
Ok(())
}
fn length(&self) -> u64 {
1
}
fn visit(&self) -> DataType {
DataType::U8(*self)
}
fn options(&self) -> MessageOption {
MessageOption::None
}
}
pub type Trame = Vec<Box<dyn Message>>;
#[macro_export]
macro_rules! trame {
() => { Trame::new() };
($( $val: expr ),*) => {{
let mut vec = Trame::new();
$( vec.push(Box::new($val)); )*
vec
}}
}
impl Message for Trame {
fn write(&self, writer: &mut dyn Write) -> RdpResult<()>{
for v in self {
v.write(writer)?;
}
Ok(())
}
fn read(&mut self, reader: &mut dyn Read) -> RdpResult<()>{
for v in self {
v.read(reader)?;
}
Ok(())
}
fn length(&self) -> u64 {
let mut sum : u64 = 0;
for v in self {
sum += v.length();
}
sum
}
fn visit(&self) -> DataType {
DataType::Trame(self)
}
fn options(&self) -> MessageOption {
MessageOption::None
}
}
pub type Component = IndexMap<String, Box<dyn Message>>;
#[macro_export]
macro_rules! component {
() => { Component::new() };
($( $key: expr => $val: expr ),*) => {{
let mut map = Component::new();
$( map.insert($key.to_string(), Box::new($val)) ; )*
map
}}
}
impl Message for Component {
fn write(&self, writer: &mut dyn Write) -> RdpResult<()>{
let mut filtering_key = HashSet::new();
for (name, value) in self.iter() {
if filtering_key.contains(name) {
continue;
}
value.write(writer)?;
if let MessageOption::SkipField(field) = value.options() {
filtering_key.insert(field);
}
}
Ok(())
}
fn read(&mut self, reader: &mut dyn Read) -> RdpResult<()>{
let mut filtering_key = HashSet::new();
let mut dynamic_size = HashMap::new();
for (name, value) in self.into_iter() {
if filtering_key.contains(name) {
continue;
}
if dynamic_size.contains_key(name) {
let mut local =vec![0; dynamic_size[name]];
reader.read_exact(&mut local)?;
value.read(&mut Cursor::new(local))?;
}
else {
value.read(reader)?;
}
match value.options() {
MessageOption::SkipField(field) => { filtering_key.insert(field); },
MessageOption::Size(field, size) => { dynamic_size.insert(field, size); },
MessageOption::None => ()
}
}
Ok(())
}
fn length(&self) -> u64 {
let mut sum : u64 = 0;
let mut filtering_key = HashSet::new();
for (name, value) in self.iter() {
if filtering_key.contains(name) {
continue;
}
if let MessageOption::SkipField(field) = value.options() {
filtering_key.insert(field);
}
sum += value.length();
}
sum
}
fn visit(&self) -> DataType {
DataType::Component(self)
}
fn options(&self) -> MessageOption {
MessageOption::None
}
}
#[derive(Copy, Clone)]
pub enum Value<Type> {
BE(Type),
LE(Type)
}
impl<Type: Copy + PartialEq> Value<Type> {
pub fn inner(&self) -> Type {
match self {
Value::<Type>::BE(e) | Value::<Type>::LE(e) => *e
}
}
}
impl<Type: Copy + PartialEq> PartialEq for Value<Type> {
fn eq(&self, other: &Self) -> bool {
return self.inner() == other.inner()
}
}
pub type U16 = Value<u16>;
impl Message for U16 {
fn write(&self, writer: &mut dyn Write) -> RdpResult<()>{
match self {
U16::BE(value) => Ok(writer.write_u16::<BigEndian>(*value)?),
U16::LE(value) => Ok(writer.write_u16::<LittleEndian>(*value)?)
}
}
fn read(&mut self, reader: &mut dyn Read) -> RdpResult<()>{
match self {
U16::BE(value) => *value = reader.read_u16::<BigEndian>()?,
U16::LE(value) => *value = reader.read_u16::<LittleEndian>()?
}
Ok(())
}
fn length(&self) -> u64 {
2
}
fn visit(&self) -> DataType {
DataType::U16(self.inner())
}
fn options(&self) -> MessageOption {
MessageOption::None
}
}
pub type U32 = Value<u32>;
impl Message for U32 {
fn write(&self, writer: &mut dyn Write) -> RdpResult<()> {
match self {
U32::BE(value) => Ok(writer.write_u32::<BigEndian>(*value)?),
U32::LE(value) => Ok(writer.write_u32::<LittleEndian>(*value)?)
}
}
fn read(&mut self, reader: &mut dyn Read) -> RdpResult<()> {
match self {
U32::BE(value) => *value = reader.read_u32::<BigEndian>()?,
U32::LE(value) => *value = reader.read_u32::<LittleEndian>()?
}
Ok(())
}
fn length(&self) -> u64 {
4
}
fn visit(&self) -> DataType {
DataType::U32(self.inner())
}
fn options(&self) -> MessageOption {
MessageOption::None
}
}
pub struct Check<T> {
value: T
}
impl<T> Check<T> {
pub fn new(value: T) -> Self{
Check {
value
}
}
}
impl<T: Message + Clone + PartialEq> Message for Check<T> {
fn write(&self, writer: &mut dyn Write) -> RdpResult<()> {
self.value.write(writer)
}
fn read(&mut self, reader: &mut dyn Read) -> RdpResult<()> {
let old = self.value.clone();
self.value.read(reader)?;
if old != self.value {
return Err(Error::RdpError(RdpError::new(RdpErrorKind::InvalidConst, "Invalid constness of data")))
}
Ok(())
}
fn length(&self) -> u64 {
self.value.length()
}
fn visit(&self) -> DataType {
self.value.visit()
}
fn options(&self) -> MessageOption {
MessageOption::None
}
}
impl Message for Vec<u8> {
fn write(&self, writer: &mut dyn Write) -> RdpResult<()> {
writer.write(self)?;
Ok(())
}
fn read(&mut self, reader: &mut dyn Read) -> RdpResult<()> {
if self.len() == 0 {
reader.read_to_end(self)?;
}
else {
reader.read_exact(self)?;
}
Ok(())
}
fn length(&self) -> u64 {
self.len() as u64
}
fn visit(&self) -> DataType {
DataType::Slice(self.as_slice())
}
fn options(&self) -> MessageOption {
MessageOption::None
}
}
pub type DynOptionFnSend<T> = dyn Fn(&T) -> MessageOption + Send;
pub struct DynOption<T> {
inner: T,
filter: Box<DynOptionFnSend<T>>
}
impl<T> DynOption<T> {
pub fn new<F: 'static>(current: T, filter: F) -> Self
where F: Fn(&T) -> MessageOption, F: Send {
DynOption {
inner: current,
filter : Box::new(filter)
}
}
}
impl<T: Message> Message for DynOption<T> {
fn write(&self, writer: &mut dyn Write) -> RdpResult<()> {
self.inner.write(writer)
}
fn read(&mut self, reader: &mut dyn Read) -> RdpResult<()> {
self.inner.read(reader)
}
fn length(&self) -> u64 {
self.inner.length()
}
fn visit(&self) -> DataType {
self.inner.visit()
}
fn options(&self) -> MessageOption {
(self.filter)(&self.inner)
}
}
pub fn to_vec(message: &dyn Message) -> Vec<u8> {
let mut stream = Cursor::new(Vec::new());
message.write(&mut stream).unwrap();
stream.into_inner()
}
#[macro_export]
macro_rules! is_none {
($expr:expr) => (match $expr.visit() {
DataType::None => true,
_ => false
})
}
impl<T: Message> Message for Option<T> {
fn write(&self, writer: &mut dyn Write) -> RdpResult<()> {
Ok(if let Some(value) = self {
value.write(writer)?
})
}
fn read(&mut self, reader: &mut dyn Read) -> RdpResult<()> {
if let Some(value) = self {
if value.read(reader).is_err() {
*self = None
}
}
Ok(())
}
fn length(&self) -> u64 {
if let Some(value) = self {
value.length()
}
else {
0
}
}
fn visit(&self) -> DataType {
if let Some(value) = self {
value.visit()
}
else {
DataType::None
}
}
fn options(&self) -> MessageOption {
MessageOption::None
}
}
pub type ArrayFnSend<T> = dyn Fn() -> T + Send;
pub struct Array<T> {
inner: Trame,
factory: Box<ArrayFnSend<T>>
}
impl<T: Message> Array<T> {
pub fn new<F: 'static>(factory: F) -> Self
where F: Fn() -> T, F: Send {
Array {
inner: trame![],
factory: Box::new(factory)
}
}
pub fn from_trame(inner: Trame) -> Self {
Array {
inner,
factory: Box::new(|| panic!("Try reading a non empty array"))
}
}
pub fn inner(&self) -> &Trame {
&self.inner
}
}
impl<T: 'static + Message> Message for Array<T> {
fn write(&self, writer: &mut dyn Write) -> RdpResult<()> {
self.inner.write(writer)
}
fn read(&mut self, reader: &mut dyn Read) -> RdpResult<()> {
loop {
let mut element = Some((self.factory)());
element.read(reader)?;
if let Some(e) = element {
self.inner.push(Box::new(e))
} else {
break;
}
}
Ok(())
}
fn length(&self) -> u64 {
self.inner.length()
}
fn visit(&self) -> DataType {
self.inner.visit()
}
fn options(&self) -> MessageOption {
MessageOption::None
}
}
impl<T> AsRef<Trame> for Array<T> {
fn as_ref(&self) -> &Trame {
&self.inner
}
}
#[cfg(test)]
mod test {
use super::*;
use std::io::Cursor;
#[test]
fn test_data_u8_write() {
let mut stream = Cursor::new(Vec::<u8>::new());
let x = 1 as u8;
x.write(&mut stream).unwrap();
assert_eq!(stream.get_ref().as_slice(), [1])
}
}