From 059a4fb1fb223deb983a6d313c9d287a722e4bf5 Mon Sep 17 00:00:00 2001 From: Miel Verkerken Date: Wed, 10 Jul 2024 11:15:15 +0200 Subject: [PATCH 1/6] Flow expiration based on packet time in pcap rather than server time. --- rustiflow/src/main.rs | 266 ++++++++++-------------------------------- 1 file changed, 61 insertions(+), 205 deletions(-) diff --git a/rustiflow/src/main.rs b/rustiflow/src/main.rs index 00012ea..70c8387 100644 --- a/rustiflow/src/main.rs +++ b/rustiflow/src/main.rs @@ -717,6 +717,7 @@ where export(&T::get_features()); } } + while let Ok(packet) = cap.next_packet() { let ts = packet.header.ts; @@ -743,6 +744,7 @@ where if amount_of_packets % 10_000 == 0 { info!("{} packets have been processed...", amount_of_packets); } + check_flows_to_expire(&flow_map_ipv4, lifespan, ts_datetime); redirect_packet_ipv4(&features_ipv4, &flow_map_ipv4, ts_instant, ts_datetime); } } @@ -754,6 +756,7 @@ where if amount_of_packets % 10_000 == 0 { info!("{} packets have been processed...", amount_of_packets); } + check_flows_to_expire(&flow_map_ipv6, lifespan, ts_datetime); redirect_packet_ipv6(&features_ipv6, &flow_map_ipv6, ts_instant, ts_datetime); } } @@ -771,12 +774,8 @@ where amount_of_packets ); } - redirect_packet_ipv4( - &features_ipv4, - &flow_map_ipv4, - ts_instant, - ts_datetime, - ); + check_flows_to_expire(&flow_map_ipv4, lifespan, ts_datetime); + redirect_packet_ipv4(&features_ipv4, &flow_map_ipv4, ts_instant, ts_datetime,); } } } @@ -790,12 +789,8 @@ where amount_of_packets ); } - redirect_packet_ipv6( - &features_ipv6, - &flow_map_ipv6, - ts_instant, - ts_datetime, - ); + check_flows_to_expire(&flow_map_ipv6, lifespan, ts_datetime); + redirect_packet_ipv6(&features_ipv6, &flow_map_ipv6, ts_instant, ts_datetime); } } } @@ -810,58 +805,6 @@ where } } - let flow_map_end_ipv4 = flow_map_ipv4.clone(); - let flow_map_end_ipv6 = flow_map_ipv6.clone(); - task::spawn(async move { - let mut interval = time::interval(Duration::from_secs(2)); - loop { - interval.tick().await; - let timestamp = Utc::now(); - - // Collect keys to remove - let mut keys_to_remove_ipv4 = Vec::new(); - for entry in flow_map_end_ipv4.iter() { - let flow = entry.value(); - let end = get_duration(flow.get_first_timestamp(), timestamp) / 1_000_000.0; - - if end >= lifespan as f64 { - if *NO_CONTAMINANT_FEATURES.lock().unwrap().deref() { - export(&flow.dump_without_contamination()); - } else { - export(&flow.dump()); - } - keys_to_remove_ipv4.push(entry.key().clone()); - } - } - - // Collect keys to remove - let mut keys_to_remove_ipv6 = Vec::new(); - for entry in flow_map_end_ipv6.iter() { - let flow = entry.value(); - let end = get_duration(flow.get_first_timestamp(), timestamp) / 1_000_000.0; - - if end >= lifespan as f64 { - if *NO_CONTAMINANT_FEATURES.lock().unwrap().deref() { - export(&flow.dump_without_contamination()); - } else { - export(&flow.dump()); - } - keys_to_remove_ipv6.push(entry.key().clone()); - } - } - - // Remove entries outside of the iteration - for key in keys_to_remove_ipv4 { - flow_map_end_ipv4.remove(&key); - } - - // Remove entries outside of the iteration - for key in keys_to_remove_ipv6 { - flow_map_end_ipv6.remove(&key); - } - } - }); - for entry in flow_map_ipv4.iter() { let flow = entry.value(); if *NO_CONTAMINANT_FEATURES.lock().unwrap().deref() { @@ -896,6 +839,33 @@ where ); } +fn check_flows_to_expire( + flow_map: &Arc>, + lifespan: u64, + ts_datetime: DateTime, +) where + T: Flow, +{ + let mut keys_to_remove = Vec::new(); + for entry in flow_map.iter() { + let flow = entry.value(); + let duration = get_duration(flow.get_first_timestamp(), ts_datetime) / 1_000_000.0; + + if duration >= lifespan as f64 { + if *NO_CONTAMINANT_FEATURES.lock().unwrap().deref() { + export(&flow.dump_without_contamination()); + } else { + export(&flow.dump()); + } + keys_to_remove.push(entry.key().clone()); + } + } + + for key in keys_to_remove { + flow_map.remove(&key); + } +} + /// Export the flow to the set export function. /// /// ### Arguments @@ -963,31 +933,11 @@ fn process_packet_ipv4( ) where T: Flow, { - let ts; - let ts_date; - if let Some(timestamp) = timestamp { - ts = timestamp; - } else { - ts = Instant::now(); - } - - if let Some(ts_datetime) = ts_datetime { - ts_date = ts_datetime; - } else { - ts_date = Utc::now(); - } - let destination = std::net::IpAddr::V4(Ipv4Addr::from(u32::from_le_bytes( - data.ipv4_source - .to_be_bytes() - .try_into() - .expect("Invalid IP length"), - ))); - let source = std::net::IpAddr::V4(Ipv4Addr::from(u32::from_le_bytes( - data.ipv4_destination - .to_be_bytes() - .try_into() - .expect("Invalid IP length"), - ))); + let ts = timestamp.unwrap_or_else(Instant::now); + let ts_date = ts_datetime.unwrap_or_else(Utc::now); + let destination = std::net::IpAddr::V4(Ipv4Addr::from(data.ipv4_destination)); + let source = std::net::IpAddr::V4(Ipv4Addr::from(data.ipv4_source)); + let combined_flags = data.combined_flags; let features = BasicFeatures { fin_flag: ((combined_flags & 0b00000001) != 0) as u8, @@ -1004,54 +954,24 @@ fn process_packet_ipv4( window_size: data.window_size, }; let flow_id = if fwd { - create_flow_id( - &source, - data.port_source, - &destination, - data.port_destination, - data.protocol, - ) + create_flow_id(&source, data.port_source, &destination, data.port_destination, data.protocol) } else { - create_flow_id( - &destination, - data.port_destination, - &source, - data.port_source, - data.protocol, - ) + create_flow_id(&destination, data.port_destination, &source, data.port_source, data.protocol) }; - let flow_id_clone = flow_id.clone(); - let flow_id_remove = flow_id.clone(); - let mut entry = flow_map.entry(flow_id).or_insert_with(|| { + let mut entry = flow_map.entry(flow_id.clone()).or_insert_with(|| { if fwd { - T::new( - flow_id_clone, - source, - data.port_source, - destination, - data.port_destination, - data.protocol, - ts_date, - ) + T::new(flow_id.clone(), source, data.port_source, destination, data.port_destination, data.protocol, ts_date) } else { - T::new( - flow_id_clone, - destination, - data.port_destination, - source, - data.port_source, - data.protocol, - ts_date, - ) + T::new(flow_id.clone(), destination, data.port_destination, source, data.port_source, data.protocol, ts_date) } }); let end = entry.update_flow(&features, &ts, ts_date, fwd); - if end.is_some() { - export(&end.unwrap()); + if let Some(flow) = end { + export(&flow); drop(entry); - flow_map.remove(&flow_id_remove); + flow_map.remove(&flow_id); } } @@ -1072,18 +992,8 @@ fn process_packet_ipv6( ) where T: Flow, { - let ts; - let ts_date; - if let Some(timestamp) = timestamp { - ts = timestamp; - } else { - ts = Instant::now(); - } - if let Some(ts_datetime) = ts_datetime { - ts_date = ts_datetime; - } else { - ts_date = Utc::now(); - } + let ts = timestamp.unwrap_or_else(Instant::now); + let ts_date = ts_datetime.unwrap_or_else(Utc::now); let destination = std::net::IpAddr::V6(Ipv6Addr::from(data.ipv6_destination)); let source = std::net::IpAddr::V6(Ipv6Addr::from(data.ipv6_source)); let combined_flags = data.combined_flags; @@ -1103,54 +1013,24 @@ fn process_packet_ipv6( }; let flow_id = if fwd { - create_flow_id( - &source, - data.port_source, - &destination, - data.port_destination, - data.protocol, - ) + create_flow_id(&source, data.port_source, &destination, data.port_destination, data.protocol) } else { - create_flow_id( - &destination, - data.port_destination, - &source, - data.port_source, - data.protocol, - ) + create_flow_id(&destination, data.port_destination, &source, data.port_source, data.protocol) }; - let flow_id_clone = flow_id.clone(); - let flow_id_remove = flow_id.clone(); - let mut entry = flow_map.entry(flow_id).or_insert_with(|| { + let mut entry = flow_map.entry(flow_id.clone()).or_insert_with(|| { if fwd { - T::new( - flow_id_clone, - source, - data.port_source, - destination, - data.port_destination, - data.protocol, - ts_date, - ) + T::new(flow_id.clone(), source, data.port_source, destination, data.port_destination, data.protocol, ts_date) } else { - T::new( - flow_id_clone, - destination, - data.port_destination, - source, - data.port_source, - data.protocol, - ts_date, - ) + T::new(flow_id.clone(), destination, data.port_destination, source, data.port_source, data.protocol, ts_date) } }); let end = entry.update_flow(&features, &ts, ts_date, fwd); - if end.is_some() { - export(&end.unwrap()); + if let Some(flow) = end { + export(&flow); drop(entry); - flow_map.remove(&flow_id_remove); + flow_map.remove(&flow_id); } } @@ -1169,13 +1049,6 @@ fn redirect_packet_ipv4( ) where T: Flow, { - let fwd_flow_id = create_flow_id( - &std::net::IpAddr::V4(Ipv4Addr::from(features_ipv4.ipv4_source)), - features_ipv4.port_source, - &std::net::IpAddr::V4(Ipv4Addr::from(features_ipv4.ipv4_destination)), - features_ipv4.port_destination, - features_ipv4.protocol, - ); let bwd_flow_id = create_flow_id( &std::net::IpAddr::V4(Ipv4Addr::from(features_ipv4.ipv4_destination)), features_ipv4.port_destination, @@ -1184,13 +1057,8 @@ fn redirect_packet_ipv4( features_ipv4.protocol, ); - if flow_map.contains_key(&fwd_flow_id) { - process_packet_ipv4(&features_ipv4, &flow_map, true, Some(timestamp), Some(ts_datetime)); - } else if flow_map.contains_key(&bwd_flow_id) { - process_packet_ipv4(&features_ipv4, &flow_map, false, Some(timestamp), Some(ts_datetime)); - } else { - process_packet_ipv4(&features_ipv4, &flow_map, true, Some(timestamp), Some(ts_datetime)); - } + let add_fwd_direction = !flow_map.contains_key(&bwd_flow_id); // If not exists, create a new flow in the forward direction or add to existing fwd flow + process_packet_ipv4(&features_ipv4, &flow_map, add_fwd_direction, Some(timestamp), Some(ts_datetime)); } /// Redirects an ipv6 packet to the correct flow. @@ -1208,13 +1076,6 @@ fn redirect_packet_ipv6( ) where T: Flow, { - let fwd_flow_id = create_flow_id( - &std::net::IpAddr::V6(Ipv6Addr::from(features_ipv6.ipv6_source)), - features_ipv6.port_source, - &std::net::IpAddr::V6(Ipv6Addr::from(features_ipv6.ipv6_destination)), - features_ipv6.port_destination, - features_ipv6.protocol, - ); let bwd_flow_id = create_flow_id( &std::net::IpAddr::V6(Ipv6Addr::from(features_ipv6.ipv6_destination)), features_ipv6.port_destination, @@ -1223,13 +1084,8 @@ fn redirect_packet_ipv6( features_ipv6.protocol, ); - if flow_map.contains_key(&fwd_flow_id) { - process_packet_ipv6(&features_ipv6, &flow_map, true, Some(timestamp), Some(ts_datetime)); - } else if flow_map.contains_key(&bwd_flow_id) { - process_packet_ipv6(&features_ipv6, &flow_map, false, Some(timestamp), Some(ts_datetime)); - } else { - process_packet_ipv6(&features_ipv6, &flow_map, true, Some(timestamp), Some(ts_datetime)); - } + let add_fwd_direction = !flow_map.contains_key(&bwd_flow_id); // If not exists, create a new flow in the forward direction or add to existing fwd flow + process_packet_ipv6(&features_ipv6, &flow_map, add_fwd_direction, Some(timestamp), Some(ts_datetime)); } /// Extracts the basic features of an ipv4 packet pnet struct. From 1b4d7f0fd56aa13ecb6490662feb837ca7dacc04 Mon Sep 17 00:00:00 2001 From: Miel Verkerken Date: Wed, 10 Jul 2024 13:24:11 +0200 Subject: [PATCH 2/6] Add support ICMPv4 and ICMPv6 --- rustiflow/src/main.rs | 65 +++++++++++++++++++++++++++++++++++++------ 1 file changed, 57 insertions(+), 8 deletions(-) diff --git a/rustiflow/src/main.rs b/rustiflow/src/main.rs index 70c8387..14cd1e2 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, IcmpTypes}, icmpv6::{Icmpv6Packet, Icmpv6Types}, ip::IpNextHeaderProtocols, ipv4::Ipv4Packet, ipv6::Ipv6Packet, tcp::TcpPacket, Packet }; use std::{ fs::{File, OpenOptions}, @@ -1107,8 +1102,8 @@ fn extract_ipv4_features(ipv4_packet: &Ipv4Packet) -> Option let mut combined_flags: u8 = 0; - let data_length: u16; - let header_length: u8; + let mut data_length: u16 = 0; + let mut header_length: u8 = 0; let length: u16; let mut window_size: u16 = 0; @@ -1139,6 +1134,31 @@ fn extract_ipv4_features(ipv4_packet: &Ipv4Packet) -> Option } else { return None; } + } else if protocol.0 == IpNextHeaderProtocols::Icmp.0 { + if let Some(icmp_packet) = IcmpPacket::new(ipv4_packet.payload()) { + destination_port = 0; // ICMP does not have ports + source_port = 0; + match icmp_packet.get_icmp_type() { + IcmpTypes::EchoRequest => { + if let Some(echo_request) = pnet::packet::icmp::echo_request::EchoRequestPacket::new(icmp_packet.packet()) { + data_length = echo_request.payload().len() as u16; + header_length = 8; + } + } + IcmpTypes::EchoReply => { + if let Some(echo_reply) = pnet::packet::icmp::echo_reply::EchoReplyPacket::new(icmp_packet.packet()) { + data_length = echo_reply.payload().len() as u16; + header_length = 8; + } + } + _ => { + return None; + } + } + length = ipv4_packet.get_total_length(); + } else { + return None; + } } else { return None; } @@ -1208,6 +1228,35 @@ 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()) { + source_port = 0; // ICMP does not have ports + destination_port = 0; + match icmpv6_packet.get_icmpv6_type() { + Icmpv6Types::EchoRequest => { + if let Some(echo_request) = pnet::packet::icmpv6::echo_request::EchoRequestPacket::new(icmpv6_packet.packet()) { + data_length = echo_request.payload().len() as u16; + header_length = 8; + } else { + return None; + } + } + Icmpv6Types::EchoReply => { + if let Some(echo_reply) = pnet::packet::icmpv6::echo_reply::EchoReplyPacket::new(icmpv6_packet.packet()) { + data_length = echo_reply.payload().len() as u16; + header_length = 8; + } else { + return None; + } + } + _ => { + return None; + } + } + length = ipv6_packet.packet().len() as u16; + } else { + return None; + } } else { return None; } From 725c196b1861270ed71517229d2a1bee3e15e2fa Mon Sep 17 00:00:00 2001 From: Miel Verkerken Date: Wed, 10 Jul 2024 14:24:06 +0200 Subject: [PATCH 3/6] support all icmp types --- rustiflow/src/main.rs | 66 +++++++++++++++---------------------------- 1 file changed, 22 insertions(+), 44 deletions(-) diff --git a/rustiflow/src/main.rs b/rustiflow/src/main.rs index 14cd1e2..2f0c243 100644 --- a/rustiflow/src/main.rs +++ b/rustiflow/src/main.rs @@ -1108,7 +1108,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(); @@ -1123,7 +1123,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(); @@ -1134,28 +1134,19 @@ fn extract_ipv4_features(ipv4_packet: &Ipv4Packet) -> Option } else { return None; } - } else if protocol.0 == IpNextHeaderProtocols::Icmp.0 { + } else if protocol == IpNextHeaderProtocols::Icmp { if let Some(icmp_packet) = IcmpPacket::new(ipv4_packet.payload()) { - destination_port = 0; // ICMP does not have ports - source_port = 0; - match icmp_packet.get_icmp_type() { - IcmpTypes::EchoRequest => { - if let Some(echo_request) = pnet::packet::icmp::echo_request::EchoRequestPacket::new(icmp_packet.packet()) { - data_length = echo_request.payload().len() as u16; - header_length = 8; - } - } - IcmpTypes::EchoReply => { - if let Some(echo_reply) = pnet::packet::icmp::echo_reply::EchoReplyPacket::new(icmp_packet.packet()) { - data_length = echo_reply.payload().len() as u16; - header_length = 8; - } - } - _ => { - return None; - } - } + // 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(); + + // Log the ICMP type and code, you can handle specific types/codes if needed + println!("ICMP Type: {:?}, Code: {:?}", icmp_type, icmp_code); } else { return None; } @@ -1230,30 +1221,17 @@ fn extract_ipv6_features(ipv6_packet: &Ipv6Packet) -> Option } } else if protocol == IpNextHeaderProtocols::Icmpv6 { if let Some(icmpv6_packet) = Icmpv6Packet::new(ipv6_packet.payload()) { - source_port = 0; // ICMP does not have ports + // 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; - match icmpv6_packet.get_icmpv6_type() { - Icmpv6Types::EchoRequest => { - if let Some(echo_request) = pnet::packet::icmpv6::echo_request::EchoRequestPacket::new(icmpv6_packet.packet()) { - data_length = echo_request.payload().len() as u16; - header_length = 8; - } else { - return None; - } - } - Icmpv6Types::EchoReply => { - if let Some(echo_reply) = pnet::packet::icmpv6::echo_reply::EchoReplyPacket::new(icmpv6_packet.packet()) { - data_length = echo_reply.payload().len() as u16; - header_length = 8; - } else { - return None; - } - } - _ => { - return None; - } - } + data_length = icmpv6_packet.payload().len() as u16; + header_length = 8; // ICMPv6 header length length = ipv6_packet.packet().len() as u16; + + // Log the ICMPv6 type and code, you can handle specific types/codes if needed + println!("ICMPv6 Type: {:?}, Code: {:?}", icmpv6_type, icmpv6_code); } else { return None; } From 6b17bb122cb416d32254b044490d6818e905ab42 Mon Sep 17 00:00:00 2001 From: Miel Verkerken Date: Wed, 10 Jul 2024 14:28:01 +0200 Subject: [PATCH 4/6] fix code scanning warnings --- rustiflow/src/main.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/rustiflow/src/main.rs b/rustiflow/src/main.rs index 2f0c243..90bc691 100644 --- a/rustiflow/src/main.rs +++ b/rustiflow/src/main.rs @@ -29,7 +29,7 @@ use flows::{ use lazy_static::lazy_static; use log::{debug, error, info}; use pnet::packet::{ - ethernet::{EtherTypes, EthernetPacket}, icmp::{IcmpPacket, IcmpTypes}, icmpv6::{Icmpv6Packet, Icmpv6Types}, 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}, @@ -1102,8 +1102,8 @@ fn extract_ipv4_features(ipv4_packet: &Ipv4Packet) -> Option let mut combined_flags: u8 = 0; - let mut data_length: u16 = 0; - let mut header_length: u8 = 0; + let data_length: u16; + let header_length: u8; let length: u16; let mut window_size: u16 = 0; From d63decbce9d4a4ac45d97a240aec8efaa19e993a Mon Sep 17 00:00:00 2001 From: Miel Verkerken Date: Wed, 10 Jul 2024 16:44:45 +0200 Subject: [PATCH 5/6] Remove debug logging --- rustiflow/src/main.rs | 6 ------ 1 file changed, 6 deletions(-) diff --git a/rustiflow/src/main.rs b/rustiflow/src/main.rs index 90bc691..da2ea61 100644 --- a/rustiflow/src/main.rs +++ b/rustiflow/src/main.rs @@ -1144,9 +1144,6 @@ fn extract_ipv4_features(ipv4_packet: &Ipv4Packet) -> Option data_length = icmp_packet.payload().len() as u16; header_length = 8; // ICMP header length length = ipv4_packet.get_total_length(); - - // Log the ICMP type and code, you can handle specific types/codes if needed - println!("ICMP Type: {:?}, Code: {:?}", icmp_type, icmp_code); } else { return None; } @@ -1229,9 +1226,6 @@ fn extract_ipv6_features(ipv6_packet: &Ipv6Packet) -> Option data_length = icmpv6_packet.payload().len() as u16; header_length = 8; // ICMPv6 header length length = ipv6_packet.packet().len() as u16; - - // Log the ICMPv6 type and code, you can handle specific types/codes if needed - println!("ICMPv6 Type: {:?}, Code: {:?}", icmpv6_type, icmpv6_code); } else { return None; } From d6b054f43dfaac2d92176fd18b7effe529b1f774 Mon Sep 17 00:00:00 2001 From: Matisse Callewaert Date: Thu, 11 Jul 2024 08:57:57 +0200 Subject: [PATCH 6/6] :sparkles: Add icmp support in ebpf programs --- ebpf-ipv4/src/main.rs | 29 +++++++++++++++++++++++++++++ ebpf-ipv6/src/main.rs | 29 +++++++++++++++++++++++++++++ rustiflow/src/main.rs | 8 ++++---- 3 files changed, 62 insertions(+), 4 deletions(-) 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 da2ea61..3b6459b 100644 --- a/rustiflow/src/main.rs +++ b/rustiflow/src/main.rs @@ -1137,8 +1137,8 @@ fn extract_ipv4_features(ipv4_packet: &Ipv4Packet) -> Option } 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(); + // 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; @@ -1219,8 +1219,8 @@ fn extract_ipv6_features(ipv6_packet: &Ipv6Packet) -> Option } 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(); + // 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;