use crate::BufResult;
use crate::tcp::Port;
use crate::{Buf, BufMut};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct HeaderMapping<T: Buf> {
buffer: T,
}
impl<T: Buf> HeaderMapping<T> {
pub fn new(buffer: T) -> Self {
Self { buffer }
}
pub fn into_inner(self) -> T {
self.buffer
}
}
impl<T: Buf> HeaderMapping<T> {
pub fn read_source_port(&self) -> BufResult<Port> {
Ok((self.buffer.get_u16_be(0)?).into())
}
pub fn read_destination_port(&self) -> BufResult<Port> {
Ok((self.buffer.get_u16_be(2)?).into())
}
pub fn read_sequence_number(&self) -> BufResult<u32> {
self.buffer.get_u32_be(4)
}
pub fn read_acknowledgment_number(&self) -> BufResult<u32> {
self.buffer.get_u32_be(8)
}
fn read_offset_reserved_flags(&self) -> BufResult<u16> {
self.buffer.get_u16_be(12)
}
pub fn read_data_offset(&self) -> BufResult<u8> {
Ok((self.read_offset_reserved_flags()? >> 12) as u8)
}
pub fn read_rsrvd(&self) -> BufResult<u8> {
Ok(((self.read_offset_reserved_flags()? >> 9) & 0x07) as u8)
}
fn read_flags(&self) -> BufResult<u16> {
Ok(self.read_offset_reserved_flags()? & 0x01FF)
}
pub fn ns(&self) -> BufResult<bool> {
Ok((self.read_flags()? & 0x0100) != 0)
}
pub fn cwr(&self) -> BufResult<bool> {
Ok((self.read_flags()? & 0x0080) != 0)
}
pub fn ece(&self) -> BufResult<bool> {
Ok((self.read_flags()? & 0x0040) != 0)
}
pub fn urg(&self) -> BufResult<bool> {
Ok((self.read_flags()? & 0x0020) != 0)
}
pub fn ack(&self) -> BufResult<bool> {
Ok((self.read_flags()? & 0x0010) != 0)
}
pub fn psh(&self) -> BufResult<bool> {
Ok((self.read_flags()? & 0x0008) != 0)
}
pub fn rst(&self) -> BufResult<bool> {
Ok((self.read_flags()? & 0x0004) != 0)
}
pub fn syn(&self) -> BufResult<bool> {
Ok((self.read_flags()? & 0x0002) != 0)
}
pub fn fin(&self) -> BufResult<bool> {
Ok((self.read_flags()? & 0x0001) != 0)
}
pub fn read_window(&self) -> BufResult<u16> {
self.buffer.get_u16_be(14)
}
pub fn read_checksum(&self) -> BufResult<u16> {
self.buffer.get_u16_be(16)
}
pub fn read_urgent_pointer(&self) -> BufResult<u16> {
self.buffer.get_u16_be(18)
}
}
impl<T: BufMut> HeaderMapping<T> {
pub fn write_source_port(&mut self, port: u16) -> BufResult<()> {
self.buffer.set_u16_be(0, port)
}
pub fn write_destination_port(&mut self, port: u16) -> BufResult<()> {
self.buffer.set_u16_be(2, port)
}
pub fn write_sequence_number(&mut self, seq: u32) -> BufResult<()> {
self.buffer.set_u32_be(4, seq)
}
pub fn write_acknowledgment_number(&mut self, ack: u32) -> BufResult<()> {
self.buffer.set_u32_be(8, ack)
}
fn write_offset_reserved_flags(&mut self, val: u16) -> BufResult<()> {
self.buffer.set_u16_be(12, val)
}
pub fn write_data_offset(&mut self, offset: u8) -> BufResult<()> {
let current = self.read_offset_reserved_flags()?;
let new_val = (current & 0x0FFF) | ((offset as u16 & 0x0F) << 12);
self.write_offset_reserved_flags(new_val)
}
pub fn write_flags(&mut self, flags: u16) -> BufResult<()> {
let current = self.read_offset_reserved_flags()?;
let new_val = (current & 0xFE00) | (flags & 0x01FF);
self.write_offset_reserved_flags(new_val)
}
fn set_flag(&mut self, mask: u16, value: bool) -> BufResult<()> {
let mut flags = self.read_flags()?;
if value {
flags |= mask;
} else {
flags &= !mask;
}
self.write_flags(flags)
}
pub fn write_ns(&mut self, value: bool) -> BufResult<()> {
self.set_flag(0x0100, value)
}
pub fn write_cwr(&mut self, value: bool) -> BufResult<()> {
self.set_flag(0x0080, value)
}
pub fn write_ece(&mut self, value: bool) -> BufResult<()> {
self.set_flag(0x0040, value)
}
pub fn write_urg(&mut self, value: bool) -> BufResult<()> {
self.set_flag(0x0020, value)
}
pub fn write_ack(&mut self, value: bool) -> BufResult<()> {
self.set_flag(0x0010, value)
}
pub fn write_psh(&mut self, value: bool) -> BufResult<()> {
self.set_flag(0x0008, value)
}
pub fn write_rst(&mut self, value: bool) -> BufResult<()> {
self.set_flag(0x0004, value)
}
pub fn write_syn(&mut self, value: bool) -> BufResult<()> {
self.set_flag(0x0002, value)
}
pub fn write_fin(&mut self, value: bool) -> BufResult<()> {
self.set_flag(0x0001, value)
}
pub fn write_window(&mut self, window: u16) -> BufResult<()> {
self.buffer.set_u16_be(14, window)
}
pub fn write_checksum(&mut self, csum: u16) -> BufResult<()> {
self.buffer.set_u16_be(16, csum)
}
pub fn write_urgent_pointer(&mut self, ptr: u16) -> BufResult<()> {
self.buffer.set_u16_be(18, ptr)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tcp_mapping_ports_and_seq() {
let mut buffer = [0u8; 20];
let mut mapping = HeaderMapping::new(&mut buffer[..]);
mapping.write_source_port(49152).unwrap();
mapping.write_destination_port(443).unwrap();
mapping.write_sequence_number(0xDEADBEEF).unwrap();
mapping.write_acknowledgment_number(0xCAFEBABE).unwrap();
let src: u16 = mapping.read_source_port().unwrap().into();
let dst: u16 = mapping.read_destination_port().unwrap().into();
assert_eq!(src, 49152);
assert_eq!(dst, 443);
assert_eq!(mapping.read_sequence_number().unwrap(), 0xDEADBEEF);
assert_eq!(mapping.read_acknowledgment_number().unwrap(), 0xCAFEBABE);
}
#[test]
fn test_tcp_mapping_flags_and_offsets() {
let mut buffer = [0u8; 20];
let mut mapping = HeaderMapping::new(&mut buffer[..]);
mapping.write_data_offset(5).unwrap();
assert_eq!(mapping.read_data_offset().unwrap(), 5);
mapping.write_syn(true).unwrap();
mapping.write_ack(true).unwrap();
mapping.write_fin(false).unwrap();
assert!(mapping.syn().unwrap());
assert!(mapping.ack().unwrap());
assert!(!mapping.fin().unwrap());
assert!(!mapping.rst().unwrap());
mapping.write_data_offset(10).unwrap();
assert_eq!(mapping.read_data_offset().unwrap(), 10);
assert!(mapping.syn().unwrap());
assert!(mapping.ack().unwrap());
mapping.write_syn(false).unwrap();
assert!(!mapping.syn().unwrap());
}
#[test]
fn test_tcp_mapping_other_fields() {
let mut buffer = [0u8; 20];
let mut mapping = HeaderMapping::new(&mut buffer[..]);
mapping.write_window(65535).unwrap();
mapping.write_checksum(0x1234).unwrap();
mapping.write_urgent_pointer(0x9999).unwrap();
assert_eq!(mapping.read_window().unwrap(), 65535);
assert_eq!(mapping.read_checksum().unwrap(), 0x1234);
assert_eq!(mapping.read_urgent_pointer().unwrap(), 0x9999);
}
}