try fix ports

Author: adgaultier

Justfile

@@ -16,7 +16,7 @@ show interface:
 
 # Run oryx
 run:
-    cargo xtask run
+    cargo xtask run --release
 
 # Build oryx
 build:

oryx-common/src/ip.rs

@@ -1,5 +1,7 @@
 use core::{fmt::Display, net::IpAddr};
 
+use network_types::{icmp::IcmpHdr, tcp::TcpHdr, udp::UdpHdr};
+
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
 pub struct TcpPacket {
@@ -58,6 +60,14 @@ pub enum IpPacket {
     Icmp(IcmpPacket),
 }
 
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub enum ProtoHdr {
+    Tcp(TcpHdr),
+    Udp(UdpHdr),
+    Icmp(IcmpHdr),
+}
+
 impl Display for IpPacket {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         match self {

oryx-common/src/lib.rs

@@ -21,7 +21,7 @@ pub mod ip;
 
 #[repr(C)]
 pub enum RawPacket {
-    Ip(IpHdr),
+    Ip(IpHdr, ip::ProtoHdr),
     Arp(ArpHdr),
 }
 
@@ -53,40 +53,40 @@ impl From<[u8; RawPacket::LEN]> for AppPacket {
     fn from(value: [u8; RawPacket::LEN]) -> Self {
         let raw_packet = value.as_ptr() as *const RawPacket;
         match unsafe { &*raw_packet } {
-            RawPacket::Ip(packet) => match packet {
-                IpHdr::V4(ipv4_packet) => match ipv4_packet.proto {
-                    IpProto::Tcp => {
+            RawPacket::Ip(packet, proto) => match packet {
+                IpHdr::V4(ipv4_packet) => match proto {
+                    ip::ProtoHdr::Tcp(header) => {
                         //FIX: This does not work
-                        let tcphdr =
-                            unsafe { raw_packet.add(mem::size_of::<Ipv4Hdr>()) } as *const TcpHdr;
+                        // let tcphdr =
+                        //     unsafe { raw_packet.add(mem::size_of::<Ipv4Hdr>()) } as *const TcpHdr;
 
                         let tcp_packet = TcpPacket {
                             src_ip: IpAddr::V4(Ipv4Addr::from(u32::from_be(ipv4_packet.src_addr))),
-                            src_port: u16::from_be(unsafe { (*tcphdr).source }),
+                            src_port: u16::from_be(header.source),
                             dst_ip: IpAddr::V4(Ipv4Addr::from(u32::from_be(ipv4_packet.dst_addr))),
-                            dst_port: u16::from_be(unsafe { (*tcphdr).dest }),
+                            dst_port: u16::from_be(header.dest),
                         };
                         AppPacket::Ip(IpPacket::Tcp(tcp_packet))
                     }
-                    IpProto::Udp => {
-                        let udphdr = unsafe {
-                            raw_packet.offset(Ipv4Hdr::LEN.try_into().unwrap()) as *const UdpHdr
-                        };
+                    ip::ProtoHdr::Udp(header) => {
+                        // let udphdr = unsafe {
+                        //     raw_packet.offset(Ipv4Hdr::LEN.try_into().unwrap()) as *const UdpHdr
+                        // };
 
                         let udp_packet = UdpPacket {
                             src_ip: IpAddr::V4(Ipv4Addr::from(u32::from_be(ipv4_packet.src_addr))),
-                            src_port: u16::from_be(unsafe { (*udphdr).source }),
+                            src_port: u16::from_be(header.source),
                             dst_ip: IpAddr::V4(Ipv4Addr::from(u32::from_be(ipv4_packet.dst_addr))),
-                            dst_port: u16::from_be(unsafe { (*udphdr).dest }),
+                            dst_port: u16::from_be(header.dest),
                         };
                         Self::Ip(IpPacket::Udp(udp_packet))
                     }
-                    IpProto::Icmp => {
-                        let icmphdr = unsafe {
-                            raw_packet.offset(Ipv4Hdr::LEN.try_into().unwrap()) as *const IcmpHdr
-                        };
+                    ip::ProtoHdr::Icmp(header) => {
+                        // let icmphdr = unsafe {
+                        //     raw_packet.offset(Ipv4Hdr::LEN.try_into().unwrap()) as *const IcmpHdr
+                        // };
 
-                        let icmp_type = match unsafe { (*icmphdr).type_ } {
+                        let icmp_type = match header.type_ {
                             0 => IcmpType::EchoRequest,
                             1 => IcmpType::EchoReply,
                             _ => IcmpType::DestinationUnreachable,

oryx-ebpf/src/main.rs

@@ -7,13 +7,16 @@ use aya_ebpf::{
     maps::RingBuf,
     programs::TcContext,
 };
-
+use core::mem;
 use network_types::{
     arp::ArpHdr,
     eth::{EthHdr, EtherType},
+    icmp::IcmpHdr,
     ip::{IpHdr, IpProto, Ipv4Hdr, Ipv6Hdr},
+    tcp::TcpHdr,
+    udp::UdpHdr,
 };
-use oryx_common::RawPacket;
+use oryx_common::{ip::ProtoHdr, RawPacket};
 
 #[map]
 static DATA: RingBuf = RingBuf::with_byte_size(4096 * 40, 0);
@@ -33,7 +36,18 @@ fn submit(packet: RawPacket) {
         buf.submit(0);
     }
 }
+#[inline]
+fn ptr_at<T>(ctx: &TcContext, offset: usize) -> Result<*const T, ()> {
+    let start = ctx.data();
+    let end = ctx.data_end();
+    let len = mem::size_of::<T>();
 
+    if start + offset + len > end {
+        return Err(());
+    }
+
+    Ok((start + offset) as *const T)
+}
 #[inline]
 fn process(ctx: TcContext) -> Result<i32, ()> {
     let ethhdr: EthHdr = ctx.load(0).map_err(|_| ())?;
@@ -42,17 +56,53 @@ fn process(ctx: TcContext) -> Result<i32, ()> {
         EtherType::Ipv4 => {
             let header: Ipv4Hdr = ctx.load(EthHdr::LEN).map_err(|_| ())?;
             match header.proto {
-                IpProto::Tcp | IpProto::Udp | IpProto::Icmp => {
-                    submit(RawPacket::Ip(IpHdr::V4(header)));
+                IpProto::Tcp => {
+                    let tcphdr: *const TcpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
+                    submit(RawPacket::Ip(
+                        IpHdr::V4(header),
+                        ProtoHdr::Tcp(unsafe { *tcphdr }),
+                    ));
+                }
+                IpProto::Udp => {
+                    let udphdr: *const UdpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
+                    submit(RawPacket::Ip(
+                        IpHdr::V4(header),
+                        ProtoHdr::Udp(unsafe { *udphdr }),
+                    ));
+                }
+                IpProto::Icmp => {
+                    let icmphdr: *const IcmpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
+                    submit(RawPacket::Ip(
+                        IpHdr::V4(header),
+                        ProtoHdr::Icmp(unsafe { *icmphdr }),
+                    ));
                 }
                 _ => {}
             }
         }
         EtherType::Ipv6 => {
             let header: Ipv6Hdr = ctx.load(EthHdr::LEN).map_err(|_| ())?;
             match header.next_hdr {
-                IpProto::Tcp | IpProto::Udp | IpProto::Icmp => {
-                    submit(RawPacket::Ip(IpHdr::V6(header)));
+                IpProto::Tcp => {
+                    let tcphdr: *const TcpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
+                    submit(RawPacket::Ip(
+                        IpHdr::V6(header),
+                        ProtoHdr::Tcp(unsafe { *tcphdr }),
+                    ));
+                }
+                IpProto::Udp => {
+                    let udphdr: *const UdpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
+                    submit(RawPacket::Ip(
+                        IpHdr::V6(header),
+                        ProtoHdr::Udp(unsafe { *udphdr }),
+                    ));
+                }
+                IpProto::Icmp => {
+                    let icmphdr: *const IcmpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
+                    submit(RawPacket::Ip(
+                        IpHdr::V6(header),
+                        ProtoHdr::Icmp(unsafe { *icmphdr }),
+                    ));
                 }
                 _ => {}
             }

oryx-tui/src/event.rs

@@ -12,7 +12,7 @@ pub enum Event {
     Key(KeyEvent),
     Mouse(MouseEvent),
     Resize(u16, u16),
-    Packet([u8; 48]),
+    Packet([u8; 72]),
     Notification(Notification),
     Reset,
 }