Support icmp traffic (#54)

* Flow expiration based on packet time in pcap rather than server time.

* Add support ICMPv4 and ICMPv6

* support all icmp types

* fix code scanning warnings

* Remove debug logging

* ✨ Add icmp support in ebpf programs

---------

Co-authored-by: Matisse Callewaert <matisse.callewaert@gmail.com>

ebpf-ipv4/src/main.rs

@@ -15,6 +15,7 @@ use network_types::{
     ip::{IpProto, Ipv4Hdr},
     tcp::TcpHdr,
     udp::UdpHdr,
+    icmp::IcmpHdr,
 };
 
 #[panic_handler]
@@ -49,6 +50,7 @@ fn process_ipv4_packet(ctx: &TcContext) -> Result<i32, ()> {
     match ipv4hdr.proto {
         IpProto::Tcp => process_tcp_packet(ctx, packet_info),
         IpProto::Udp => process_udp_packet(ctx, packet_info),
+        IpProto::Icmp => process_icmp_packet(ctx, packet_info),
         _ => Ok(TC_ACT_PIPE),
     }
 }
@@ -71,6 +73,15 @@ fn process_udp_packet(ctx: &TcContext, packet_info: PacketInfo) -> Result<i32, (
     Ok(TC_ACT_PIPE)
 }
 
+fn process_icmp_packet(ctx: &TcContext, packet_info: PacketInfo) -> Result<i32, ()> {
+    let icmphdr = ctx
+        .load::<IcmpHdr>(EthHdr::LEN + Ipv4Hdr::LEN)
+        .map_err(|_| ())?;
+    let packet_log = packet_info.to_packet_log(&icmphdr);
+    EVENTS_IPV4.output(ctx, &packet_log, 0);
+    Ok(TC_ACT_PIPE)
+}
+
 struct PacketInfo {
     ipv4_source: u32,
     ipv4_destination: u32,
@@ -154,3 +165,21 @@ impl NetworkHeader for UdpHdr {
         UdpHdr::LEN as u8
     }
 }
+
+impl NetworkHeader for IcmpHdr {
+    fn source_port(&self) -> u16 {
+        0
+    }
+    fn destination_port(&self) -> u16 {
+        0
+    }
+    fn window_size(&self) -> u16 {
+        0
+    }
+    fn combined_flags(&self) -> u8 {
+        0
+    }
+    fn header_length(&self) -> u8 {
+        IcmpHdr::LEN as u8
+    }
+}

ebpf-ipv6/src/main.rs

@@ -15,6 +15,7 @@ use network_types::{
     ip::{IpProto, Ipv6Hdr},
     tcp::TcpHdr,
     udp::UdpHdr,
+    icmp::IcmpHdr,
 };
 
 #[panic_handler]
@@ -49,6 +50,7 @@ fn process_ipv6_packet(ctx: &TcContext) -> Result<i32, ()> {
     match ipv6hdr.next_hdr {
         IpProto::Tcp => process_tcp_packet(ctx, packet_info),
         IpProto::Udp => process_udp_packet(ctx, packet_info),
+        IpProto::Icmp => process_icmp_packet(ctx, packet_info),
         _ => Ok(TC_ACT_PIPE),
     }
 }
@@ -71,6 +73,15 @@ fn process_udp_packet(ctx: &TcContext, packet_info: PacketInfo) -> Result<i32, (
     Ok(TC_ACT_PIPE)
 }
 
+fn process_icmp_packet(ctx: &TcContext, packet_info: PacketInfo) -> Result<i32, ()> {
+    let icmphdr = ctx
+        .load::<IcmpHdr>(EthHdr::LEN + Ipv6Hdr::LEN)
+        .map_err(|_| ())?;
+    let packet_log = packet_info.to_packet_log(&icmphdr);
+    EVENTS_IPV6.output(ctx, &packet_log, 0);
+    Ok(TC_ACT_PIPE)
+}
+
 struct PacketInfo {
     ipv6_source: u128,
     ipv6_destination: u128,
@@ -154,3 +165,21 @@ impl NetworkHeader for UdpHdr {
         UdpHdr::LEN as u8
     }
 }
+
+impl NetworkHeader for IcmpHdr {
+    fn source_port(&self) -> u16 {
+        0
+    }
+    fn destination_port(&self) -> u16 {
+        0
+    }
+    fn window_size(&self) -> u16 {
+        0
+    }
+    fn combined_flags(&self) -> u8 {
+        0
+    }
+    fn header_length(&self) -> u8 {
+        IcmpHdr::LEN as u8
+    }
+}

rustiflow/src/main.rs

@@ -29,12 +29,7 @@ use flows::{
 use lazy_static::lazy_static;
 use log::{debug, error, info};
 use pnet::packet::{
-    ethernet::{EtherTypes, EthernetPacket},
-    ip::IpNextHeaderProtocols,
-    ipv4::Ipv4Packet,
-    ipv6::Ipv6Packet,
-    tcp::TcpPacket,
-    Packet,
+    ethernet::{EtherTypes, EthernetPacket}, icmp::IcmpPacket, icmpv6::Icmpv6Packet, ip::IpNextHeaderProtocols, ipv4::Ipv4Packet, ipv6::Ipv6Packet, tcp::TcpPacket, Packet
 };
 use std::{
     fs::{File, OpenOptions},
@@ -1123,7 +1118,7 @@ fn extract_ipv4_features(ipv4_packet: &Ipv4Packet) -> Option<BasicFeaturesIpv4>
 
     let mut window_size: u16 = 0;
 
-    if protocol.0 == IpNextHeaderProtocols::Tcp.0 {
+    if protocol == IpNextHeaderProtocols::Tcp {
         if let Some(tcp_packet) = TcpPacket::new(ipv4_packet.payload()) {
             source_port = tcp_packet.get_source();
             destination_port = tcp_packet.get_destination();
@@ -1138,7 +1133,7 @@ fn extract_ipv4_features(ipv4_packet: &Ipv4Packet) -> Option<BasicFeaturesIpv4>
         } else {
             return None;
         }
-    } else if protocol.0 == IpNextHeaderProtocols::Udp.0 {
+    } else if protocol == IpNextHeaderProtocols::Udp {
         if let Some(udp_packet) = pnet::packet::udp::UdpPacket::new(ipv4_packet.payload()) {
             source_port = udp_packet.get_source();
             destination_port = udp_packet.get_destination();
@@ -1149,6 +1144,19 @@ fn extract_ipv4_features(ipv4_packet: &Ipv4Packet) -> Option<BasicFeaturesIpv4>
         } else {
             return None;
         }
+    } else if protocol == IpNextHeaderProtocols::Icmp {
+        if let Some(icmp_packet) = IcmpPacket::new(ipv4_packet.payload()) {
+            // For ICMP, we will extract the type and code, along with data length
+            // let icmp_type = icmp_packet.get_icmp_type();
+            // let icmp_code = icmp_packet.get_icmp_code();
+            source_port = 0; // ICMPv6 does not have ports
+            destination_port = 0;
+            data_length = icmp_packet.payload().len() as u16;
+            header_length = 8; // ICMP header length
+            length = ipv4_packet.get_total_length();
+        } else {
+            return None;
+        }
     } else {
         return None;
     }
@@ -1218,6 +1226,19 @@ fn extract_ipv6_features(ipv6_packet: &Ipv6Packet) -> Option<BasicFeaturesIpv6>
         } else {
             return None;
         }
+    } else if protocol == IpNextHeaderProtocols::Icmpv6 {
+        if let Some(icmpv6_packet) = Icmpv6Packet::new(ipv6_packet.payload()) {
+            // For ICMPv6, we will just extract the type and code for now, along with data length
+            // let icmpv6_type = icmpv6_packet.get_icmpv6_type();
+            // let icmpv6_code = icmpv6_packet.get_icmpv6_code();
+            source_port = 0; // ICMPv6 does not have ports
+            destination_port = 0;
+            data_length = icmpv6_packet.payload().len() as u16;
+            header_length = 8; // ICMPv6 header length
+            length = ipv6_packet.packet().len() as u16;
+        } else {
+            return None;
+        }
     } else {
         return None;
     }