move filters to ebpf

pythops · Sep 17, 2024 · 19eb710 · 19eb710
1 parent 82abd75
commit 19eb710
Show file tree

Hide file tree

Showing 13 changed files with 351 additions and 203 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/oryx-common/src/lib.rs b/oryx-common/src/lib.rs
@@ -12,6 +12,7 @@ use network_types::{arp::ArpHdr, ip::IpHdr};
 
 pub mod arp;
 pub mod ip;
+pub mod protocols;
 
 #[repr(C)]
 pub enum RawPacket {

diff --git a/oryx-common/src/protocols.rs b/oryx-common/src/protocols.rs
@@ -0,0 +1,89 @@
+use core::fmt::Display;
+
+#[derive(Debug, Copy, Clone)]
+#[repr(C, u8)]
+pub enum Protocol {
+    Transport(TransportProtocol) = 0,
+    Network(NetworkProtocol) = 1,
+    Link(LinkProtocol) = 2,
+}
+
+// Transport Protocols
+
+pub const NB_TRANSPORT_PROTOCOL: u16 = 2;
+
+#[derive(Debug, Copy, Clone, PartialEq)]
+#[repr(C)]
+pub enum TransportProtocol {
+    TCP = 0,
+    UDP = 1,
+}
+
+impl TransportProtocol {
+    pub fn all() -> [TransportProtocol; 2] {
+        [TransportProtocol::TCP, TransportProtocol::UDP]
+    }
+}
+
+impl Display for TransportProtocol {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self {
+            TransportProtocol::TCP => write!(f, "Tcp"),
+            TransportProtocol::UDP => write!(f, "Udp"),
+        }
+    }
+}
+
+// Network Protocols
+
+pub const NB_NETWORK_PROTOCOL: u16 = 3;
+
+#[derive(Debug, Copy, Clone, PartialEq)]
+#[repr(C)]
+pub enum NetworkProtocol {
+    Ipv4 = 0,
+    Ipv6 = 1,
+    Icmp = 2,
+}
+
+impl NetworkProtocol {
+    pub fn all() -> [NetworkProtocol; 3] {
+        [
+            NetworkProtocol::Ipv4,
+            NetworkProtocol::Ipv6,
+            NetworkProtocol::Icmp,
+        ]
+    }
+}
+
+impl Display for NetworkProtocol {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self {
+            NetworkProtocol::Ipv4 => write!(f, "Ipv4"),
+            NetworkProtocol::Ipv6 => write!(f, "Ipv6"),
+            NetworkProtocol::Icmp => write!(f, "Icmp"),
+        }
+    }
+}
+
+// Link Protocols
+
+pub const NB_LINK_PROTOCOL: u16 = 1;
+
+#[derive(Debug, Copy, Clone, PartialEq)]
+#[repr(C)]
+pub enum LinkProtocol {
+    Arp = 0,
+}
+
+impl LinkProtocol {
+    pub fn all() -> [LinkProtocol; 1] {
+        [LinkProtocol::Arp]
+    }
+}
+
+impl Display for LinkProtocol {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        write!(f, "Arp")
+    }
+}
diff --git a/oryx-ebpf/src/main.rs b/oryx-ebpf/src/main.rs
@@ -4,7 +4,7 @@
 use aya_ebpf::{
     bindings::TC_ACT_PIPE,
     macros::{classifier, map},
-    maps::RingBuf,
+    maps::{Array, RingBuf},
     programs::TcContext,
 };
 use core::mem;
@@ -16,11 +16,24 @@ use network_types::{
     tcp::TcpHdr,
     udp::UdpHdr,
 };
-use oryx_common::{ip::ProtoHdr, RawPacket};
+use oryx_common::{
+    ip::ProtoHdr,
+    protocols::{LinkProtocol, NetworkProtocol, Protocol, TransportProtocol},
+    RawPacket,
+};
 
 #[map]
 static DATA: RingBuf = RingBuf::with_byte_size(4096 * RawPacket::LEN as u32, 0);
 
+#[map]
+static NETWORK_FILTERS: Array<u32> = Array::with_max_entries(8, 0);
+
+#[map]
+static TRANSPORT_FILTERS: Array<u32> = Array::with_max_entries(8, 0);
+
+#[map]
+static LINK_FILTERS: Array<u32> = Array::with_max_entries(8, 0);
+
 #[classifier]
 pub fn oryx(ctx: TcContext) -> i32 {
     match process(ctx) {
@@ -48,29 +61,64 @@ fn ptr_at<T>(ctx: &TcContext, offset: usize) -> Result<*const T, ()> {
 
     Ok((start + offset) as *const T)
 }
+
+#[inline]
+fn filter_packet(protocol: Protocol) -> bool {
+    match protocol {
+        Protocol::Network(p) => {
+            if let Some(v) = NETWORK_FILTERS.get(p as u32) {
+                return *v == 1;
+            }
+        }
+        Protocol::Transport(p) => {
+            if let Some(v) = TRANSPORT_FILTERS.get(p as u32) {
+                return *v == 1;
+            }
+        }
+        Protocol::Link(p) => {
+            if let Some(v) = LINK_FILTERS.get(p as u32) {
+                return *v == 1;
+            }
+        }
+    }
+    false
+}
+
 #[inline]
 fn process(ctx: TcContext) -> Result<i32, ()> {
     let ethhdr: EthHdr = ctx.load(0).map_err(|_| ())?;
 
     match ethhdr.ether_type {
         EtherType::Ipv4 => {
+            if filter_packet(Protocol::Network(NetworkProtocol::Ipv4)) {
+                return Ok(TC_ACT_PIPE);
+            }
             let header: Ipv4Hdr = ctx.load(EthHdr::LEN).map_err(|_| ())?;
             match header.proto {
                 IpProto::Tcp => {
+                    if filter_packet(Protocol::Transport(TransportProtocol::TCP)) {
+                        return Ok(TC_ACT_PIPE);
+                    }
                     let tcphdr: *const TcpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
                     submit(RawPacket::Ip(
                         IpHdr::V4(header),
                         ProtoHdr::Tcp(unsafe { *tcphdr }),
                     ));
                 }
                 IpProto::Udp => {
+                    if filter_packet(Protocol::Transport(TransportProtocol::UDP)) {
+                        return Ok(TC_ACT_PIPE);
+                    }
                     let udphdr: *const UdpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
                     submit(RawPacket::Ip(
                         IpHdr::V4(header),
                         ProtoHdr::Udp(unsafe { *udphdr }),
                     ));
                 }
                 IpProto::Icmp => {
+                    if filter_packet(Protocol::Network(NetworkProtocol::Icmp)) {
+                        return Ok(TC_ACT_PIPE);
+                    }
                     let icmphdr: *const IcmpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
                     submit(RawPacket::Ip(
                         IpHdr::V4(header),
@@ -81,23 +129,35 @@ fn process(ctx: TcContext) -> Result<i32, ()> {
             }
         }
         EtherType::Ipv6 => {
+            if filter_packet(Protocol::Network(NetworkProtocol::Ipv6)) {
+                return Ok(TC_ACT_PIPE);
+            }
             let header: Ipv6Hdr = ctx.load(EthHdr::LEN).map_err(|_| ())?;
             match header.next_hdr {
                 IpProto::Tcp => {
+                    if filter_packet(Protocol::Transport(TransportProtocol::TCP)) {
+                        return Ok(TC_ACT_PIPE);
+                    }
                     let tcphdr: *const TcpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv6Hdr::LEN)?;
                     submit(RawPacket::Ip(
                         IpHdr::V6(header),
                         ProtoHdr::Tcp(unsafe { *tcphdr }),
                     ));
                 }
                 IpProto::Udp => {
+                    if filter_packet(Protocol::Transport(TransportProtocol::UDP)) {
+                        return Ok(TC_ACT_PIPE);
+                    }
                     let udphdr: *const UdpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv6Hdr::LEN)?;
                     submit(RawPacket::Ip(
                         IpHdr::V6(header),
                         ProtoHdr::Udp(unsafe { *udphdr }),
                     ));
                 }
                 IpProto::Icmp => {
+                    if filter_packet(Protocol::Network(NetworkProtocol::Icmp)) {
+                        return Ok(TC_ACT_PIPE);
+                    }
                     let icmphdr: *const IcmpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv6Hdr::LEN)?;
                     submit(RawPacket::Ip(
                         IpHdr::V6(header),
@@ -108,6 +168,9 @@ fn process(ctx: TcContext) -> Result<i32, ()> {
             }
         }
         EtherType::Arp => {
+            if filter_packet(Protocol::Link(LinkProtocol::Arp)) {
+                return Ok(TC_ACT_PIPE);
+            }
             let header: ArpHdr = ctx.load(EthHdr::LEN).map_err(|_| ())?;
             submit(RawPacket::Arp(header));
         }

diff --git a/oryx-tui/Cargo.toml b/oryx-tui/Cargo.toml
@@ -9,6 +9,7 @@ homepage.workspace = true
 edition.workspace = true
 
 [dependencies]
+crossterm = { version = "0.28", default-features = false }
 ratatui = "0.28"
 tui-big-text = "0.6"
 tui-input = "0.10"