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

use bytes::{Bytes, Buf, BufMut};
use crate::errors::*;
use std::io::Cursor;

mod std_impls;
mod write;

pub use self::write::*;

pub use dnsbox_derive::DnsPacketData;

pub trait DnsPacketData: Sized {
	fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self>;
	fn serialize(&self, context: &mut DnsPacketWriteContext, packet: &mut Vec<u8>) -> Result<()>;
}

pub fn deserialize_with<F, O>(data: Bytes, parser: F) -> Result<O>
where
	F: FnOnce(&mut Cursor<Bytes>) -> Result<O>,
{
	let mut c = Cursor::new(data);
	let result = parser(&mut c)?;
	failure::ensure!(!c.has_remaining(), "data remaining: {} bytes", c.remaining());
	Ok(result)
}

pub fn remaining_bytes(data: &mut Cursor<Bytes>) -> Bytes {
	let pos = data.position() as usize;
	let len = data.remaining();
	let result = data.get_ref().slice(pos, pos + len);
	data.advance(len);
	result
}

pub fn get_blob(data: &mut Cursor<Bytes>, len: usize) -> Result<Bytes> {
	check_enough_data!(data, len, "blob content");
	let pos = data.position() as usize;
	let blob = data.get_ref().slice(pos, pos + len);
	data.advance(len);
	Ok(blob)
}

pub fn short_blob(data: &mut Cursor<Bytes>) -> Result<Bytes> {
	check_enough_data!(data, 1, "short blob length");
	let blob_len = data.get_u8() as usize;
	get_blob(data, blob_len)
}

pub fn write_short_blob(data: &[u8], packet: &mut Vec<u8>) -> Result<()> {
	failure::ensure!(data.len() < 256, "short blob must be at most 255 bytes long");
	packet.reserve(data.len() + 1);
	packet.put_u8(data.len() as u8);
	packet.put_slice(data);
	Ok(())
}