diff options
Diffstat (limited to 'src')
-rw-r--r-- | src/count.rs | 67 | ||||
-rw-r--r-- | src/ethertype.rs | 20 | ||||
-rw-r--r-- | src/ippacket.rs | 99 | ||||
-rw-r--r-- | src/layer3.rs | 65 | ||||
-rw-r--r-- | src/main.rs | 172 |
5 files changed, 0 insertions, 423 deletions
diff --git a/src/count.rs b/src/count.rs deleted file mode 100644 index c3d92bf..0000000 --- a/src/count.rs +++ /dev/null @@ -1,67 +0,0 @@ -const U16_MAX: usize = u16::MAX as usize; - -pub struct Count { - tcp_tx: Box<[usize]>, - tcp_rx: Box<[usize]>, - udp_tx: Box<[usize]>, - udp_rx: Box<[usize]>, -} - -impl Count { - pub fn new() -> Self { - Self { - tcp_tx: Box::new([0; U16_MAX]), - tcp_rx: Box::new([0; U16_MAX]), - udp_tx: Box::new([0; U16_MAX]), - udp_rx: Box::new([0; U16_MAX]), - } - } - - pub fn push_tcp(&mut self, tx_flag: bool, src: u16, dst: u16, len: usize) { - // The unsafe code here is bad and not good. As far as I am aware, this - // is safe except for allowing race conditions. - if tx_flag { - let bad_mut = unsafe { &mut *(self.tcp_tx.as_ptr() as *mut [usize; U16_MAX]) }; - bad_mut[src as usize] += len; - } else { - let bad_mut = unsafe { &mut *(self.tcp_rx.as_ptr() as *mut [usize; U16_MAX]) }; - bad_mut[dst as usize] += len; - } - } - - pub fn push_udp(&mut self, tx_flag: bool, src: u16, dst: u16, len: usize) { - // The unsafe code here is bad and not good. As far as I am aware, this - // is safe except for allowing race conditions. - if tx_flag { - let bad_mut = unsafe { &mut *(self.udp_tx.as_ptr() as *mut [usize; U16_MAX]) }; - bad_mut[src as usize] += len; - } else { - let bad_mut = unsafe { &mut *(self.udp_rx.as_ptr() as *mut [usize; U16_MAX]) }; - bad_mut[dst as usize] += len; - } - } - - pub fn tcp_tx(&self, port: u16) -> usize { - self.tcp_tx[port as usize] - } - - pub fn tcp_rx(&self, port: u16) -> usize { - self.tcp_rx[port as usize] - } - - pub fn udp_tx(&self, port: u16) -> usize { - self.udp_tx[port as usize] - } - - pub fn udp_rx(&self, port: u16) -> usize { - self.udp_rx[port as usize] - } - - pub fn many_tcp_tx(&self, ports: &[u16]) -> usize { - ports.iter().fold(0, |acc, port| acc + self.tcp_tx[*port as usize]) - } - - pub fn many_tcp_rx(&self, ports: &[u16]) -> usize { - ports.iter().fold(0, |acc, port| acc + self.tcp_rx[*port as usize]) - } -} diff --git a/src/ethertype.rs b/src/ethertype.rs deleted file mode 100644 index 0d728f3..0000000 --- a/src/ethertype.rs +++ /dev/null @@ -1,20 +0,0 @@ -#[derive(Copy, Clone, Debug, PartialEq)] -pub enum EtherType { - /// The frame type is IEEE 802.3 and this is it's length - Length(u16), - IPv4, - IPv6, - Unknown(u16), -} - -impl EtherType { - //TODO: check ethertype is correct - pub fn new(n: u16) -> Self { - match n { - n if n <= 1500 => Self::Length(n), - 0x0800 => Self::IPv4, - 0x86DD => Self::IPv6, - n => Self::Unknown(n), - } - } -} diff --git a/src/ippacket.rs b/src/ippacket.rs deleted file mode 100644 index 60858d2..0000000 --- a/src/ippacket.rs +++ /dev/null @@ -1,99 +0,0 @@ -use std::net::Ipv4Addr; - -pub struct Ipv4Packet<'p> { - /// The computed size of the IP packet - header_length: usize, - /// The total length of the IP packet when all fragments are combined - total_length: u16, - more_fragments: bool, - fragment_offset: usize, - next_header: IpNextHeader, - pub src: Ipv4Addr, - pub dst: Ipv4Addr, - payload: &'p [u8], -} - -impl<'p> Ipv4Packet<'p> { - pub fn new(buffer: &'p [u8]) -> Self { - let ihl = buffer[0].to_be() & 0b0000_1111; - let header_length = ihl as usize * 4; - - let total_length = u16::from_be_bytes([buffer[2], buffer[3]]); - - // Fragmentation - let more_fragments = (buffer[6] & 0b0010_0000) > 0; - let fragment_offset = u16::from_be_bytes([buffer[6] & 0b0001_1111, buffer[7]]); - // Fragments are in units of 8 bytes - let true_frag_offset = fragment_offset as usize * 8; - - let next_header = IpNextHeader::new(buffer[9]); - - let src = Ipv4Addr::new(buffer[12], buffer[13], buffer[14], buffer[15]); - let dst = Ipv4Addr::new(buffer[16], buffer[17], buffer[18], buffer[19]); - - Self { - header_length, - total_length, - more_fragments, - fragment_offset: true_frag_offset, - next_header, - src, - dst, - payload: &buffer[header_length..], - } - } - - pub fn get_source(&self) -> Ipv4Addr { - self.src - } - - pub fn get_destination(&self) -> Ipv4Addr { - self.dst - } - - pub fn has_more_fragments(&self) -> bool { - self.more_fragments - } - - pub fn get_fragment_offset(&self) -> usize { - self.fragment_offset - } - - pub fn get_next_header(&self) -> IpNextHeader { - self.next_header - } - - pub fn get_payload(&self) -> &[u8] { - &self.payload - } - - pub fn get_packet_len(&self) -> usize { - self.total_length as usize - } - - pub fn get_payload_len(&self) -> usize { - // An IPv4 header is always 20 bytes long - self.total_length as usize - 20 - } -} - -#[derive(Copy, Clone, Debug, PartialEq)] -pub enum IpNextHeader { - Icmp, - Igmp, - Tcp, - Udp, - Unknown(u8), -} - -impl IpNextHeader { - pub fn new(n: u8) -> Self { - match n { - 1 => Self::Icmp, - 2 => Self::Igmp, - 6 => Self::Tcp, - 17 => Self::Udp, - n => Self::Unknown(n), - } - } -} diff --git a/src/layer3.rs b/src/layer3.rs deleted file mode 100644 index 7521c0d..0000000 --- a/src/layer3.rs +++ /dev/null @@ -1,65 +0,0 @@ -pub struct Tcp<'p> { - src: u16, - dst: u16, - payload: &'p [u8], -} - -impl<'p> Tcp<'p> { - pub fn new(buffer: &'p [u8]) -> Self { - let src = u16::from_be_bytes([buffer[0], buffer[1]]); - let dst = u16::from_be_bytes([buffer[2], buffer[3]]); - let data_offset = (buffer[12].to_be() & 0b1111_0000) >> 4; - // Offset is number of 32-bit words - let true_offset = data_offset as usize * 4; - - Self { - src, - dst, - payload: &buffer[true_offset..], - } - } - - pub fn source_port(&self) -> u16 { - self.src - } - - pub fn destination_port(&self) -> u16 { - self.dst - } - - pub fn payload(&self) -> &[u8] { - self.payload - } -} - -pub struct Udp<'p> { - src: u16, - dst: u16, - payload: &'p [u8], -} - -impl<'p> Udp<'p> { - pub fn new(buffer: &'p [u8]) -> Self { - let src = u16::from_be_bytes([buffer[0], buffer[1]]); - let dst = u16::from_be_bytes([buffer[2], buffer[3]]); - let _len = u16::from_be_bytes([buffer[4], buffer[5]]); - - Self { - src, - dst, - payload: &buffer[8..], - } - } - - pub fn source_port(&self) -> u16 { - self.src - } - - pub fn destination_port(&self) -> u16 { - self.dst - } - - pub fn payload(&self) -> &[u8] { - self.payload - } -} diff --git a/src/main.rs b/src/main.rs deleted file mode 100644 index 3b0f63c..0000000 --- a/src/main.rs +++ /dev/null @@ -1,172 +0,0 @@ -use std::{cell::UnsafeCell, error::Error, net::Ipv4Addr, rc::Rc, thread::JoinHandle}; - -use count::Count; -use ethertype::EtherType; -use ippacket::{IpNextHeader, Ipv4Packet}; -use layer3::{Tcp, Udp}; -use pnet_datalink::{Channel, Config}; -use scurvy::{Argument, Scurvy}; -use smol::{ - LocalExecutor, - channel::{self, Receiver, Sender}, - io::AsyncWriteExt, - net::TcpListener, -}; - -mod count; -mod ethertype; -mod ippacket; -mod layer3; - -fn main() -> Result<(), Box<dyn Error>> { - let args = vec![ - Argument::new(&["interface", "iface"]).arg("dev"), - Argument::new("ip").arg("addr"), - ]; - let scurvy = Scurvy::make(args); - - let interface_name = scurvy.get("interface").unwrap(); - let interface = match pnet_datalink::interfaces().iter().find(|i| i.name == interface_name) { - None => { - eprintln!("No interface found named '{interface_name}'"); - return Ok(()); - } - Some(i) => i.clone(), - }; - - let ip_want: Ipv4Addr = scurvy.get("ip").unwrap().parse()?; - - let mut channel = match pnet_datalink::channel(&interface, Config::default())? { - Channel::Ethernet(_tx, rx) => rx, - _ => unimplemented!(), - }; - - let stat = stat_thread(); - - loop { - let pkt = channel.next()?; - - let ethertype = EtherType::new(u16::from_be_bytes([pkt[12], pkt[13]])); - let eth_payload = &pkt[14..]; - - match ethertype { - EtherType::IPv4 => { - let ip = Ipv4Packet::new(eth_payload); - let ip_payload = ip.get_payload(); - - let ip_tx = if ip.src == ip_want { - true - } else if ip.dst == ip_want { - false - } else { - continue; - }; - - // 6 byte per MAC (x2), 2 byte ethertype, 2 byte crc - let total_l2_len = ip.get_packet_len() + 18; - - match ip.get_next_header() { - IpNextHeader::Tcp => { - let tcp = Tcp::new(ip_payload); - - stat.tx.send_blocking(Meow::Tcp { - tx: ip_tx, - src: tcp.source_port(), - dst: tcp.destination_port(), - len: total_l2_len, - })?; - } - IpNextHeader::Udp => { - let udp = Udp::new(ip_payload); - - stat.tx.send_blocking(Meow::Udp { - tx: ip_tx, - src: udp.source_port(), - dst: udp.destination_port(), - len: total_l2_len, - })?; - } - _ => (), - } - } - _ => (), - } - } -} - -#[allow(dead_code)] -struct StatHandle { - hwnd: JoinHandle<()>, - tx: Sender<Meow>, -} - -enum Meow { - Tcp { - tx: bool, - src: u16, - dst: u16, - len: usize, - }, - Udp { - tx: bool, - src: u16, - dst: u16, - len: usize, - }, -} - -fn stat_thread() -> StatHandle { - let (tx, rx) = channel::unbounded(); - - let hwnd = std::thread::spawn(|| stat(rx)); - - StatHandle { hwnd, tx } -} - -fn stat(rx: Receiver<Meow>) { - let cat = Rc::new(UnsafeCell::new(Count::new())); - - let ex = Rc::new(LocalExecutor::new()); - - ex.spawn(async { - loop { - match rx.recv().await { - Err(_e) => { - eprintln!("error receiving! breaking from loop"); - break; - } - Ok(Meow::Tcp { tx, src, dst, len }) => { - unsafe { &mut *cat.get() }.push_tcp(tx, src, dst, len); - } - Ok(Meow::Udp { tx, src, dst, len }) => { - unsafe { &mut *cat.get() }.push_udp(tx, src, dst, len); - } - } - } - }) - .detach(); - - let run_cat = Rc::clone(&cat); - let run_ex = Rc::clone(&ex); - let run = ex.run(async move { - let socket = TcpListener::bind("127.0.0.1:6666").await.unwrap(); - - loop { - let (mut stream, _caddr) = socket.accept().await.unwrap(); - let per_spawn_cat = Rc::clone(&run_cat); - - run_ex - .spawn(async move { - let cat = unsafe { &mut *per_spawn_cat.get() }; - let http_tx = cat.many_tcp_tx(&[80, 443]); - let http_rx = cat.many_tcp_rx(&[80, 443]); - - let str = format!("{http_tx}\n{http_rx}"); - stream.write_all(str.as_bytes()).await.unwrap(); - }) - .detach(); - } - }); - - smol::block_on(run); -} |