From 80c93e65ea585d95d2e931e9dcdfcf41198785fc Mon Sep 17 00:00:00 2001 From: Anhad Singh Date: Sun, 24 Sep 2023 18:18:21 +1000 Subject: [PATCH] feat(tcp): add TCP options builder Signed-off-by: Anhad Singh --- src/lib.rs | 27 +++++++- src/transport/tcp.rs | 151 ++++++++++++++++++++++++++++++++++++------- 2 files changed, 153 insertions(+), 25 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index d9c43a8b..810bcd23 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,4 @@ -#![no_std] +// #![no_std] #![feature(allocator_api, trivial_bounds, new_uninit, associated_type_defaults)] extern crate alloc; @@ -287,7 +287,7 @@ impl<'a> PacketParser<'a> { pub mod tests { use alloc::vec::Vec; - use crate::transport::{TcpOption, TcpOptions}; + use crate::transport::{TcpOption, TcpOptions, TcpOptionsBuilder}; use super::data_link::{Eth, EthType, MacAddr}; use super::network::{Ipv4, Ipv4Addr, Ipv4Type}; @@ -402,4 +402,27 @@ pub mod tests { ] ); } + + #[test] + fn tcp_options_build() { + const RAW_PACKET: &[u8] = &[ + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 8, 0, 69, 0, 0, 64, 0, 0, 0, 0, 64, 6, 122, 185, + 255, 255, 255, 255, 255, 255, 255, 255, 31, 144, 0, 80, 0, 0, 0, 0, 0, 0, 0, 0, 160, 0, + 0, 0, 202, 14, 0, 0, 2, 4, 5, 180, 4, 2, 8, 10, 151, 174, 53, 222, 0, 0, 0, 0, 3, 3, 7, + 0, 69, 69, 69, 69, + ]; + + let eth = Eth::new(MacAddr::NULL, MacAddr::NULL, EthType::Ip); + let ip = Ipv4::new(Ipv4Addr::BROADCAST, Ipv4Addr::BROADCAST, Ipv4Type::Tcp); + let tcp = Tcp::new(8080, 80); + let mut options = TcpOptionsBuilder::new() + .with(TcpOption::MaxSegmentSize(1460)) + .with(TcpOption::SackPermitted) + .with(TcpOption::TimeStamp(2544776670, 0)) + .with(TcpOption::WindowScale(7)); + + let options = options.build(); + let packet = (eth / ip / tcp / options / [69u8; 4]).into_boxed_bytes(); + assert_eq!(&*packet, RAW_PACKET) + } } diff --git a/src/transport/tcp.rs b/src/transport/tcp.rs index b58a1ffa..d7c0cf02 100644 --- a/src/transport/tcp.rs +++ b/src/transport/tcp.rs @@ -35,19 +35,19 @@ const_assert_eq!(core::mem::size_of::(), 20); impl Tcp { pub fn new(src_port: u16, dest_port: u16) -> Self { - let mut flags = 0; - flags.set_bits(12..=15, core::mem::size_of::() as u16 / 4); - - Self { + let mut tcp = Self { src_port: src_port.into(), dest_port: dest_port.into(), seq_nr: 0.into(), ack_nr: 0.into(), - flags: flags.into(), + flags: 0.into(), window: 0.into(), checksum: 0.into(), urgent_ptr: 0.into(), - } + }; + + tcp.set_header_size(core::mem::size_of::() as u8); + tcp } #[inline] @@ -116,6 +116,19 @@ impl Tcp { header_size as u8 * core::mem::size_of::() as u8 } + /// Sets the header size to `size`. + /// + /// # Panics + /// This function panics if `size` is less than the size of the TCP header. + #[inline] + fn set_header_size(&mut self, size: u8) { + assert!(size >= core::mem::size_of::() as u8); + + let mut flags = self.flags.to_native(); + flags.set_bits(12..=15, size as u16 / 4); + self.flags = flags.into(); + } + #[inline] pub fn options_size(&self) -> u8 { self.header_size() - core::mem::size_of::() as u8 @@ -148,6 +161,24 @@ crate::impl_stack!(@make Tcp { } }); +const TCP_OPTEOL: u8 = 0; + +const TCP_OPTNOP: u8 = 1; +const TCP_LENNOP: usize = 1; + +const TCP_OPTMSS: u8 = 2; +const TCP_LENMSS: usize = 4; + +const TCP_OPTWINDOW_SCALE: u8 = 3; +const TCP_LENWINDOW_SCALE: usize = 3; + +const TCP_OPTSACKPERM: u8 = 4; +const TCP_LENSACKPERM: usize = 2; + +// const TCP_OPTSACK: u8 = 5; +const TCP_OPTTIMESTAMPS: u8 = 8; +const TCP_LENTIMESTAMPS: usize = 10; + #[derive(Debug)] pub enum TcpOptionErr { /// The TCP option had an invalid size. @@ -176,14 +207,6 @@ impl<'a> Iterator for TcpOptionsIter<'a> { type Item = Result; fn next(&mut self) -> Option { - const TCP_OPTEOL: u8 = 0; - const TCP_OPTNOP: u8 = 1; - const TCP_OPTMSS: u8 = 2; - const TCP_OPTWINDOW_SCALE: u8 = 3; - const TCP_OPTSACKPERM: u8 = 4; - // const TCP_OPTSACK: u8 = 5; - const TCP_OPTTIMESTAMPS: u8 = 8; - // FIXME(andypython): In this case, it is an unexpected end of options list (i.e, the // EOL option was not found). Should we return an error instead? let kind = *self.options.get(self.cursor)?; @@ -193,7 +216,7 @@ impl<'a> Iterator for TcpOptionsIter<'a> { return None; } else if kind == TCP_OPTNOP { // NOP (used for padding) - self.cursor += 1; + self.cursor += TCP_LENNOP; return self.next(); } @@ -220,7 +243,7 @@ impl<'a> Iterator for TcpOptionsIter<'a> { }; match kind { - TCP_OPTMSS => match ensure_size(4) { + TCP_OPTMSS => match ensure_size(TCP_LENMSS) { Ok(data) => { let mss = u16::from_be_bytes([data[0], data[1]]); Some(Ok(TcpOption::MaxSegmentSize(mss))) @@ -229,12 +252,17 @@ impl<'a> Iterator for TcpOptionsIter<'a> { Err(err) => Some(Err(err)), }, - TCP_OPTWINDOW_SCALE => match ensure_size(3) { + TCP_OPTWINDOW_SCALE => match ensure_size(TCP_LENWINDOW_SCALE) { Ok(data) => Some(Ok(TcpOption::WindowScale(data[0]))), Err(err) => Some(Err(err)), }, - TCP_OPTTIMESTAMPS => match ensure_size(10) { + TCP_OPTSACKPERM => match ensure_size(TCP_LENSACKPERM) { + Ok(_) => Some(Ok(TcpOption::SackPermitted)), + Err(err) => Some(Err(err)), + }, + + TCP_OPTTIMESTAMPS => match ensure_size(TCP_LENTIMESTAMPS) { Ok(data) => { let init_ts = u32::from_be_bytes([data[0], data[1], data[2], data[3]]); let echo_ts = u32::from_be_bytes([data[4], data[5], data[6], data[7]]); @@ -244,11 +272,6 @@ impl<'a> Iterator for TcpOptionsIter<'a> { Err(err) => Some(Err(err)), }, - TCP_OPTSACKPERM => match ensure_size(2) { - Ok(_) => Some(Ok(TcpOption::SackPermitted)), - Err(err) => Some(Err(err)), - }, - _ => { if option_size == 0 { // Invalid option length. @@ -267,6 +290,71 @@ impl<'a> Iterator for TcpOptionsIter<'a> { } } +pub struct TcpOptionsBuilder { + // XXX: The maximum size of the TCP options is 40 bytes. + options: [u8; 40], + cursor: usize, +} + +impl TcpOptionsBuilder { + pub fn new() -> Self { + Self { + options: [0; 40], + cursor: 0, + } + } + + pub fn with(mut self, option: TcpOption) -> Self { + match option { + TcpOption::MaxSegmentSize(mss) => { + self.options[self.cursor] = TCP_OPTMSS; + self.options[self.cursor + 1] = TCP_LENMSS as u8; + self.options[self.cursor + 2..self.cursor + TCP_LENMSS] + .copy_from_slice(&mss.to_be_bytes()); + + self.cursor += TCP_LENMSS; + } + + TcpOption::WindowScale(scale) => { + self.options[self.cursor] = TCP_OPTWINDOW_SCALE; + self.options[self.cursor + 1] = TCP_LENWINDOW_SCALE as u8; + self.options[self.cursor + 2] = scale; + + self.cursor += TCP_LENWINDOW_SCALE; + } + + TcpOption::SackPermitted => { + self.options[self.cursor] = TCP_OPTSACKPERM; + self.options[self.cursor + 1] = TCP_LENSACKPERM as u8; + + self.cursor += TCP_LENSACKPERM; + } + + TcpOption::TimeStamp(init_ts, echo_ts) => { + let init_ts = init_ts.to_be_bytes(); + let echo_ts = echo_ts.to_be_bytes(); + + self.options[self.cursor] = TCP_OPTTIMESTAMPS; + self.options[self.cursor + 1] = TCP_LENTIMESTAMPS as u8; + + self.options[self.cursor + 2..self.cursor + 6].copy_from_slice(&init_ts); + self.options[self.cursor + 6..self.cursor + 10].copy_from_slice(&echo_ts); + + self.cursor += TCP_LENTIMESTAMPS; + } + } + + self + } + + pub fn build<'a>(&'a mut self) -> TcpOptions<'a> { + self.options[self.cursor] = TCP_OPTEOL; + self.cursor += 1; + + TcpOptions(&self.options[..self.cursor]) + } +} + #[derive(Debug)] pub struct TcpOptions<'a>(&'a [u8]); @@ -277,6 +365,7 @@ impl<'a> TcpOptions<'a> { self.0 } + #[inline] pub fn iter(&self) -> TcpOptionsIter<'a> { TcpOptionsIter { options: self.0, @@ -285,6 +374,13 @@ impl<'a> TcpOptions<'a> { } } +impl<'a> From<&'a [u8]> for TcpOptions<'a> { + #[inline] + fn from(value: &'a [u8]) -> Self { + Self(value) + } +} + unsafe impl<'a> Parsable<'a> for TcpOptions<'a> { type Output = TcpOptions<'a>; @@ -348,6 +444,15 @@ unsafe impl<'a> Protocol for TcpOptions<'a> { self.0.len(), ); } + + let header_size = core::mem::size_of::(); + + // SAFETY: We only implement Stack for [`TcpOptions`] and we do not + // implement [`IsSafeToWrite`], so we know that the the parent layer is TCP. In addition, + // the caller must guarantee that we have exclusive access to the whole packet thus, the + // pointer dereference is valid. + let tcp = unsafe { mem.sub(header_size).cast::().as_mut() }; + tcp.set_header_size(self.0.len() as u8 + header_size as u8); } #[inline]