rust-dnsbox/lib/dnsbox-base/src/ser/packet/std_impls.rs

99 lines
2.5 KiB
Rust

use bytes::{Bytes, Buf, BufMut, BigEndian};
use errors::*;
use ser::packet::{DnsPacketData, DnsPacketWriteContext};
use std::io::{Cursor, Read};
use std::mem::size_of;
use std::net::{Ipv4Addr, Ipv6Addr};
impl DnsPacketData for u8 {
fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
check_enough_data!(data, size_of::<Self>(), "u8");
Ok(data.get_u8())
}
fn serialize(&self, _context: &mut DnsPacketWriteContext, packet: &mut Vec<u8>) -> Result<()> {
packet.reserve(size_of::<Self>());
packet.put_u8(*self);
Ok(())
}
}
impl DnsPacketData for u16 {
fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
check_enough_data!(data, size_of::<Self>(), "u16");
Ok(data.get_u16::<BigEndian>())
}
fn serialize(&self, _context: &mut DnsPacketWriteContext, packet: &mut Vec<u8>) -> Result<()> {
packet.reserve(size_of::<Self>());
packet.put_u16::<BigEndian>(*self);
Ok(())
}
}
impl DnsPacketData for u32 {
fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
check_enough_data!(data, size_of::<Self>(), "u32");
Ok(data.get_u32::<BigEndian>())
}
fn serialize(&self, _context: &mut DnsPacketWriteContext, packet: &mut Vec<u8>) -> Result<()> {
packet.reserve(size_of::<Self>());
packet.put_u32::<BigEndian>(*self);
Ok(())
}
}
impl DnsPacketData for Ipv4Addr {
fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
check_enough_data!(data, size_of::<u32>(), "Ipv4Addr");
Ok(data.get_u32::<BigEndian>().into())
}
fn serialize(&self, _context: &mut DnsPacketWriteContext, packet: &mut Vec<u8>) -> Result<()> {
let data = self.octets();
debug_assert!(data.len() == 4);
packet.reserve(data.len());
packet.put_slice(&data);
Ok(())
}
}
impl DnsPacketData for Ipv6Addr {
fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
check_enough_data!(data, 16, "Ipv6Addr");
let mut buf = [0u8; 16];
data.read_exact(&mut buf)?;
Ok(buf.into())
}
fn serialize(&self, _context: &mut DnsPacketWriteContext, packet: &mut Vec<u8>) -> Result<()> {
let data = self.octets();
debug_assert!(data.len() == 16);
packet.reserve(data.len());
packet.put_slice(&data);
Ok(())
}
}
#[cfg(test)]
mod tests {
use errors::*;
fn deserialize<T: super::DnsPacketData>(data: &'static [u8]) -> Result<T> {
use bytes::{Bytes,Buf};
use std::io::Cursor;
let mut c = Cursor::new(Bytes::from_static(data));
let result = T::deserialize(&mut c)?;
if c.remaining() != 0 {
bail!("data remaining: {}", c.remaining())
}
Ok(result)
}
#[test]
fn test_u16() {
assert!(deserialize::<u16>(b"\x80\x08").unwrap() == 0x8008);
}
}