[#93]: identity code refactoring. improved file organization

Author: LorenzoTettamanti

core/src/components/conntracker/src/main.rs

@@ -13,19 +13,33 @@
 // Imports
 #![no_std]
 #![no_main]
+#![allow(warnings)]
 
-use bytemuck::{Pod,Zeroable};
+//mod skbuff;
+
+use bytemuck::{ Pod, Zeroable };
 use aya_ebpf::{
-    bindings::{TC_ACT_OK,TC_ACT_SHOT},
-    macros::{ classifier, map },
+    bindings::{ TC_ACT_OK, TC_ACT_SHOT },
+    macros::{ classifier, map, kprobe, tracepoint },
     maps::PerfEventArray,
     maps::LruPerCpuHashMap,
-    programs::TcContext,
+    programs::{ TcContext, TracePointContext },
+    helpers::{ bpf_probe_read_kernel, bpf_ktime_get_ns },
 };
+use aya_ebpf::EbpfContext;
+//use crate::skbuff::{ sock, sock_common };
 use aya_log_ebpf::info;
-use core::mem;
-
-
+use core::{ mem, ptr };
+//use crate::skbuff::proto;
+//use crate::skbuff::{ iphdr };
+//use crate::skbuff::sk_buff;
+use network_types::{
+    eth::{ EthHdr, EtherType },
+    ip::{ IpProto, Ipv4Hdr },
+    tcp::TcpHdr,
+    udp::UdpHdr,
+};
+use core::ptr::addr_of;
 /* 
     * ETHERNET TYPE II FRAME:
     * Reference: https://it.wikipedia.org/wiki/Frame_Ethernet
@@ -43,18 +57,7 @@ use core::mem;
     * https://en.wikipedia.org/wiki/IPv4#Header
     *
     * Original reference:
-    * https://datatracker.ietf.org/doc/html/rfc791
-    *
-    *
-    * Ipv4 header datagram
 
-    0                   1                   2                   3                       TOT BYTES OFFSET (full length)     
-    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1      32 bit           
-   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
-   |Version|  IHL  |Type of Service|          Total Length         |     4 bytes            4             
-   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
-   |         Identification        |Flags|      Fragment Offset    |     4 bytes            8            
-   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+    
    |  Time to Live |    Protocol   |         Header Checksum       |     4 bytes            12             
    +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    |                       Source Address                          |     4 bytes            16             
@@ -95,16 +98,28 @@ pub struct PacketLog {
 }
 
 #[repr(C)]
-#[derive(Clone, Copy, Pod, Zeroable)]
-pub struct ConnArray{
-    pub event_id: u16,
-    pub connection_id: u16
+#[derive(Clone, Copy)]
+pub struct ConnArray {
+    pub src_ip: u32,
+    pub dst_ip: u32,
+    pub src_port: u16,
+    pub dst_port: u16,
+    pub proto: u8,
 }
 
-#[map(name="EventsMap")]
+#[map(name = "EventsMap")]
 static mut EVENTS: PerfEventArray<PacketLog> = PerfEventArray::new(0);
 #[map(name = "ConnectionMap")]
-pub static mut CONNMAP: LruPerCpuHashMap<u8,ConnArray> = LruPerCpuHashMap::with_max_entries(1024,0); //TODO: modify this to a LRU HASHMAP
+pub static mut ACTIVE_CONNECTIONS: LruPerCpuHashMap<
+    u16,
+    ConnArray
+> = LruPerCpuHashMap::with_max_entries(65536, 0);
+
+#[map(name = "ConnectionTrackerMap")]
+pub static mut CONNTRACKER: LruPerCpuHashMap<ConnArray, u8> = LruPerCpuHashMap::with_max_entries(
+    65536,
+    0
+);
 
 const IPV4_ETHERTYPE: u16 = 0x0800;
 
@@ -122,20 +137,35 @@ const ETHERTYPE_BYTES: usize = 2;
 const SRC_PORT_OFFSET_FROM_IP_HEADER: usize = 0;
 const DST_PORT_OFFSET_FROM_IP_HEADER: usize = 2;
 
+static ETH_STACK_BYTES: usize = SRC_MAC + DST_MAC + ETHERTYPE_BYTES;
+static DST_T0TAL_BYTES_OFFSET: usize = ETH_STACK_BYTES + DST_BYTE_OFFSET;
+static SRC_T0TAL_BYTES_OFFSET: usize = ETH_STACK_BYTES + SRC_BYTE_OFFSET;
+static PROTOCOL_T0TAL_BYTES_OFFSET: usize = ETH_STACK_BYTES + IPV4_PROTOCOL_OFFSET;
+
+const AF_INET: u16 = 2; //ipv4
+const AF_INET6: u16 = 10; //ipv6
 
-static ETH_STACK_BYTES :usize = SRC_MAC+DST_MAC+ETHERTYPE_BYTES;
-static DST_T0TAL_BYTES_OFFSET :usize = ETH_STACK_BYTES + DST_BYTE_OFFSET;
-static SRC_T0TAL_BYTES_OFFSET :usize = ETH_STACK_BYTES + SRC_BYTE_OFFSET;
-static PROTOCOL_T0TAL_BYTES_OFFSET :usize = ETH_STACK_BYTES + IPV4_PROTOCOL_OFFSET;
+const IPPROTO_UDP: u8 = 17;
+const IPPROTO_TCP: u8 = 6;
 
+//TODO: add kprobe tracing for process ID
+//kprobe docs: https://docs.kernel.org/trace/kprobes.html
 
+/* constants */
+const HOST_NETNS_INUM: u32 = 4026531993;
+const KUBE_POD_CIDR: u32 = 0x0af40000; // 10.244.0.0/16
 
+/* Helper Functions */
+#[inline]
+unsafe fn is_kube_internal(ip: u32) -> bool {
+    (ip & 0xffff0000) == KUBE_POD_CIDR
+}
 
 #[classifier]
 pub fn identity_classifier(ctx: TcContext) -> i32 {
     match try_identity_classifier(ctx) {
         Ok(_) => TC_ACT_OK,
-        Err(_) => TC_ACT_SHOT,//block packets that returns errors
+        Err(_) => TC_ACT_SHOT, //block packets that returns errors
     }
 }
 
@@ -145,59 +175,81 @@ fn try_identity_classifier(ctx: TcContext) -> Result<(), i64> {
     //only ipv4 protcol allowed
     if eth_proto != IPV4_ETHERTYPE {
         return Ok(());
-    } 
-    
+    }
+
     //read if the packets has Options
     let first_ipv4_byte = u8::from_be(ctx.load::<u8>(ETH_STACK_BYTES).map_err(|_| 1)?);
-    let ihl = (first_ipv4_byte & 0x0F) as usize; /* 0x0F=00001111 &=AND bit a bit operator to extract the last 4 bit*/
-    let ip_header_len= ihl*4; //returns the header lenght in bytes
+    let ihl = (first_ipv4_byte &
+        0x0f) as usize; /* 0x0F=00001111 &=AND bit a bit operator to extract the last 4 bit*/
+    let ip_header_len = ihl * 4; //returns the header lenght in bytes
 
     //get the source ip,destination ip and connection id
     let src_ip = u32::from_be(ctx.load::<u32>(SRC_T0TAL_BYTES_OFFSET).map_err(|_| 1)?); // ETH+SOURCE_ADDRESS
-    let src_port = u16::from_be(ctx.load::<u16>(ETH_STACK_BYTES+ip_header_len+SRC_PORT_OFFSET_FROM_IP_HEADER).map_err(|_| 1)?); //14+IHL-Lenght+0
+    let src_port = u16::from_be(
+        ctx
+            .load::<u16>(ETH_STACK_BYTES + ip_header_len + SRC_PORT_OFFSET_FROM_IP_HEADER)
+            .map_err(|_| 1)?
+    ); //14+IHL-Lenght+0
     let dst_ip = u32::from_be(ctx.load::<u32>(DST_T0TAL_BYTES_OFFSET).map_err(|_| 1)?); // ETH+ DESTINATION_ADDRESS
-    let dst_port = u16::from_be(ctx.load::<u16>(ETH_STACK_BYTES+ip_header_len+DST_PORT_OFFSET_FROM_IP_HEADER).map_err(|_| 1)?); //14+IHL-Lenght+0
+    let dst_port = u16::from_be(
+        ctx
+            .load::<u16>(ETH_STACK_BYTES + ip_header_len + DST_PORT_OFFSET_FROM_IP_HEADER)
+            .map_err(|_| 1)?
+    ); //14+IHL-Lenght+0
     let proto = u8::from_be(ctx.load::<u8>(PROTOCOL_T0TAL_BYTES_OFFSET).map_err(|_| 1)?);
-    
 
     //not logging internal communication packets
     //TODO: do not log internal communications such as minikube dashboard packets or kubectl api packets
-    let ip_to_block = u32::from_be_bytes([192,168,49,1]); //inverted requence
-    let dst_ip_to_block = u32::from_be_bytes([192,168,49,2]);    
-    
-    
-    // XOR to generate the hash id for the given connection
-    let event_id = (src_ip ^ dst_ip ^ (src_port as u32) ^ (dst_port as u32) ^ (proto as u32)) as u16; //generate one for every event using a 'byte XOR' operation 
+    let ip_to_block = u32::from_be_bytes([192, 168, 49, 1]); //inverted requence
+    let dst_ip_to_block = u32::from_be_bytes([192, 168, 49, 2]);
 
-    let connection_id = (src_ip ^ dst_ip ^(proto as u32)) as u16; //added host_id to track the host to count every all the different connections
+    let key = ConnArray {
+        src_ip,
+        dst_ip,
+        src_port,
+        dst_port,
+        proto,
+    };
 
-    if src_ip == ip_to_block && dst_ip == dst_ip_to_block {
-        return Ok(());
-    }
-    else{
-        //log all other packets
-        let log = PacketLog {
-            proto,
-            src_ip,
-            src_port,
-            dst_ip,
-            dst_port,
-            event_id,
-        };
-        let connections = ConnArray{
-            event_id,
-            connection_id
-        };
-        unsafe {
-            EVENTS.output(&ctx, &log, 0); //output to userspace
-            //TODO: add more parameters to better identify the active connection (maybe timestamp?)
-            CONNMAP.insert(&proto,&connections, 0) //save hash_id to kernel space lru per cpu hashmap
-        };
+    // XOR to generate the hash id for the given connection
+    let event_id = (src_ip ^
+        dst_ip ^
+        (src_port as u32) ^
+        (dst_port as u32) ^
+        (proto as u32)) as u16; //generate one for every event using a 'byte XOR' operation
+
+    //let connection_id = (src_ip ^ dst_ip ^(proto as u32)) as u16; //added host_id to track the host to count every all the different connections
+
+    //if
+    //  (unsafe { is_kube_internal(src_ip) }) ||
+    //(unsafe { is_kube_internal(dst_ip) }) ||
+    // src_ip == ip_to_block ||
+    // src_ip == dst_ip_to_block
+    //{
+    //   return Ok(());
+    //} else {
+    //log all other packets
+    let log = PacketLog {
+        proto,
+        src_ip,
+        src_port,
+        dst_ip,
+        dst_port,
+        event_id,
+    };
+    //let connections = ConnArray{
+    //  event_id,
+    //connection_id
+    //};
+    unsafe {
+        EVENTS.output(&ctx, &log, 0); //output to userspace
+        //TODO: add more parameters to better identify the active connection (maybe timestamp?)
+        ACTIVE_CONNECTIONS.insert(&event_id, &key, 0);
     }
-        
+    //}
+
     Ok(())
 }
-
 #[panic_handler]
 fn panic(_info: &core::panic::PanicInfo) -> ! {
     loop {

core/src/components/identity/src/enums.rs

@@ -0,0 +1,11 @@
+/*
+ * IpProtocols enum to reconstruct the packet protocol based on the
+ * IPV4 Header Protocol code
+ */
+#[derive(Debug)]
+#[repr(u8)]
+pub enum IpProtocols {
+    ICMP = 1,
+    TCP = 6,
+    UDP = 17,
+}

core/src/components/identity/src/helpers.rs

@@ -0,0 +1,56 @@
+use aya::{ maps::{ perf::{ PerfEventArrayBuffer }, MapData } };
+use crate::structs::PacketLog;
+use bytes::BytesMut;
+use std::{ borrow::BorrowMut, net::Ipv4Addr, sync::{ atomic::{ AtomicBool, Ordering }, Arc } };
+use crate::enums::IpProtocols;
+use tracing::{ info, error, warn };
+
+pub async fn display_events<T: BorrowMut<MapData>>(
+    mut perf_buffers: Vec<PerfEventArrayBuffer<T>>,
+    running: Arc<AtomicBool>,
+    mut buffers: Vec<BytesMut>
+) {
+    while running.load(Ordering::SeqCst) {
+        for buf in perf_buffers.iter_mut() {
+            match buf.read_events(&mut buffers) {
+                Ok(events) => {
+                    for i in 0..events.read {
+                        let data = &buffers[i];
+                        if data.len() >= std::mem::size_of::<PacketLog>() {
+                            let pl: PacketLog = unsafe {
+                                std::ptr::read(data.as_ptr() as *const _)
+                            };
+                            let src = Ipv4Addr::from(u32::from_be(pl.src_ip));
+                            let dst = Ipv4Addr::from(u32::from_be(pl.dst_ip));
+                            let src_port = u16::from_be(pl.src_port as u16);
+                            let dst_port = u16::from_be(pl.dst_port as u16);
+                            let event_id = pl.event_id;
+
+                            match IpProtocols::try_from(pl.proto) {
+                                Ok(proto) => {
+                                    info!(
+                                        "Hash: {} Protocol: {:?} SRC: {}:{} -> DST: {}:{}",
+                                        event_id,
+                                        proto,
+                                        src,
+                                        src_port,
+                                        dst,
+                                        dst_port
+                                    );
+                                }
+                                Err(_) =>
+                                    info!("Hash: {} Protocol: Unknown ({})", event_id, pl.proto),
+                            };
+                        } else {
+                            warn!("Received packet data too small: {} bytes", data.len());
+                        }
+                    }
+                }
+                Err(e) => {
+                    error!("Error reading events: {:?}", e);
+                }
+            }
+        }
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+    }
+}