diff --git a/ebpf-ipv4/src/main.rs b/ebpf-ipv4/src/main.rs index 9264ea7..b6f1bb7 100644 --- a/ebpf-ipv4/src/main.rs +++ b/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 { 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 Result { + let icmphdr = ctx + .load::(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 + } +} diff --git a/ebpf-ipv6/src/main.rs b/ebpf-ipv6/src/main.rs index b22a0ec..295602c 100644 --- a/ebpf-ipv6/src/main.rs +++ b/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 { 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 Result { + let icmphdr = ctx + .load::(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 + } +} diff --git a/rustiflow/src/main.rs b/rustiflow/src/main.rs index a3f9249..89f32fb 100644 --- a/rustiflow/src/main.rs +++ b/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 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 } 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 } 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 } 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; }