diff --git a/Cargo.lock b/Cargo.lock index 6169ff4..8f3a08d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1546,6 +1546,7 @@ dependencies = [ "url", "uuid", "xsk-rs", + "zerocopy 0.8.33", ] [[package]] @@ -2693,9 +2694,9 @@ dependencies = [ [[package]] name = "sysinfo" -version = "0.39.3" +version = "0.39.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "21d0d938c10fcda3e897e28aaddf4ab462375d411f4378cd63b1c945f69aba96" +checksum = "14311e7e9a03114cd4b65eedd54e8fed2945e17f08586ae97ef53bc0669f9581" dependencies = [ "libc", "memchr", diff --git a/mantis/Cargo.toml b/mantis/Cargo.toml index a1cdae1..939d20d 100644 --- a/mantis/Cargo.toml +++ b/mantis/Cargo.toml @@ -21,6 +21,7 @@ parking_lot = "0.12.5" rust-embed = "8.7.2" serde = { workspace = true } serde_json = "1.0.143" +zerocopy = { workspace = true } sysinfo = "0.39.2" thiserror = "2.0.3" tokio = { version = "1.40.0", features = ["full", "macros"] } diff --git a/mantis/src/core/ebpf/xsk_manager.rs b/mantis/src/core/ebpf/xsk_manager.rs index 187aed0..7bab57e 100644 --- a/mantis/src/core/ebpf/xsk_manager.rs +++ b/mantis/src/core/ebpf/xsk_manager.rs @@ -61,8 +61,8 @@ impl XskManager { let combined_queue_count = config.combined_queue_count; for queue_id in 0..combined_queue_count { - let (ingress_to_egress_tx, ingress_to_egress_rx) = bounded(config.channel_size); - let (egress_to_ingress_tx, egress_to_ingress_rx) = bounded(config.channel_size); + let (ingress_to_egress_tx, ingress_to_egress_rx) = bounded::>>(config.channel_size); + let (egress_to_ingress_tx, egress_to_ingress_rx) = bounded::>>(config.channel_size); let ingress_xsk = XskPair::new( config.clone(), @@ -258,8 +258,8 @@ impl XskPair { pub fn run( mut self, - forward_tx: Sender>, - forward_rx: Receiver>, + forward_tx: Sender>>, + forward_rx: Receiver>>, ) -> Result, EbpfError> { let (shutdown_tx, shutdown_rx) = oneshot::channel(); @@ -363,26 +363,26 @@ impl XskPair { Ok(nb_completed) } - fn process_rx_queue(&mut self, forward_tx: &Sender>) -> Result { + fn process_rx_queue(&mut self, forward_tx: &Sender>>) -> Result { let mut rx_descs = vec![FrameDesc::default(); 256]; let rx_count = unsafe { self.rx.consume(&mut rx_descs) }; if rx_count > 0 { for rx_desc in rx_descs.iter().take(rx_count) { - let lengths = rx_desc.lengths(); - let packet_len = lengths.data() as usize; - + let packet_len = rx_desc.lengths().data() as usize; let data = unsafe { self.umem.data(rx_desc) }; - let packet_data = data.contents()[..packet_len].to_vec(); - - if let Some(ref se) = self.suricata_engine { - se.inject(packet_data.clone()); - } + let packet_slice = &data.contents()[..packet_len]; + // ML engine reads directly from UMEM — no copy for this consumer if let Some(ref engine) = self.engine { - engine.process_packet(&packet_data, self.direction == Direction::Ingress); + engine.process_packet(packet_slice, self.direction == Direction::Ingress); } + // Single copy shared between Suricata and forward_tx via Arc + let packet_data = Arc::new(packet_slice.to_vec()); + if let Some(ref se) = self.suricata_engine { + se.inject(Arc::clone(&packet_data)); + } if let Err(e) = forward_tx.try_send(packet_data) { match e { crossbeam::channel::TrySendError::Full(_) => { @@ -455,7 +455,7 @@ impl XskPair { } } - fn process_tx_queue(&mut self, forward_rx: &Receiver>) -> Result { + fn process_tx_queue(&mut self, forward_rx: &Receiver>>) -> Result { // Drain completed TX frames first to maximise pool availability. let _ = self.process_comp_queue(); @@ -474,7 +474,7 @@ impl XskPair { // Consume at most min(pool_size, 64) packets so we never over-commit. let max_to_send = pool_size.min(64); - let mut packets_to_send = Vec::with_capacity(max_to_send); + let mut packets_to_send: Vec>> = Vec::with_capacity(max_to_send); while let Ok(packet) = forward_rx.try_recv() { packets_to_send.push(packet); if packets_to_send.len() >= max_to_send { diff --git a/mantis/src/detection/suricata/engine.rs b/mantis/src/detection/suricata/engine.rs index 2490fb8..1331095 100644 --- a/mantis/src/detection/suricata/engine.rs +++ b/mantis/src/detection/suricata/engine.rs @@ -22,7 +22,7 @@ const SURICATA_LOG: &str = "/tmp/suricata.log"; const CHANNEL_CAP: usize = 4096; pub struct SuricataEngine { - tx: Sender>, + tx: Sender>>, child: std::sync::Mutex, } @@ -133,7 +133,7 @@ impl SuricataEngine { }) .map_err(|e| SuricataError::ProcessSpawnFailed { reason: e.to_string() })?; - let (tx, rx) = bounded::>(CHANNEL_CAP); + let (tx, rx) = bounded::>>(CHANNEL_CAP); thread::Builder::new() .name("suricata-mirror".into()) @@ -182,7 +182,7 @@ impl SuricataEngine { } /* Non-blocking: drops silently when the channel is full under load. */ - pub fn inject(&self, data: Vec) { + pub fn inject(&self, data: Arc>) { match self.tx.try_send(data) { Ok(()) => {} Err(crossbeam::channel::TrySendError::Full(_)) => { diff --git a/mantis/src/utils/packet_parser.rs b/mantis/src/utils/packet_parser.rs index be4f2a9..1f2faec 100644 --- a/mantis/src/utils/packet_parser.rs +++ b/mantis/src/utils/packet_parser.rs @@ -1,166 +1,190 @@ -use std::mem; use std::time; use common::model::event::{Event, IPv4Event, IPv6Event, TcpFlags}; use network_types::ip::IpProto; +use zerocopy::byteorder::{BigEndian, U16, U32}; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Ref}; + +#[repr(C, packed)] +#[derive(FromBytes, IntoBytes, KnownLayout, Immutable, Clone, Copy)] +struct EthHdr { + _dst_mac: [u8; 6], + _src_mac: [u8; 6], + ether_type: U16, +} + +#[repr(C, packed)] +#[derive(FromBytes, IntoBytes, KnownLayout, Immutable, Clone, Copy)] +struct Ipv4Hdr { + version_ihl: u8, + _dscp_ecn: u8, + total_len: U16, + _ident: U16, + _flags_frag: U16, + _ttl: u8, + protocol: u8, + _checksum: U16, + src_addr: U32, + dst_addr: U32, +} + +#[repr(C, packed)] +#[derive(FromBytes, IntoBytes, KnownLayout, Immutable, Clone, Copy)] +struct Ipv6Hdr { + _version_tc_fl: U32, + payload_len: U16, + next_hdr: u8, + _hop_limit: u8, + src_addr: [u8; 16], + dst_addr: [u8; 16], +} + +#[repr(C, packed)] +#[derive(FromBytes, IntoBytes, KnownLayout, Immutable, Clone, Copy)] +struct TcpHdr { + src_port: U16, + dst_port: U16, + _seq_num: U32, + _ack_num: U32, + data_off_flags: U16, + window: U16, + _checksum: U16, + _urgent_ptr: U16, +} + +#[repr(C, packed)] +#[derive(FromBytes, IntoBytes, KnownLayout, Immutable, Clone, Copy)] +struct UdpHdr { + src_port: U16, + dst_port: U16, + _length: U16, + _checksum: U16, +} + +const ETH_HDR_LEN: usize = 14; +const IPV6_HDR_LEN: usize = 40; pub fn parse_packet(packet_data: &[u8]) -> Option<(Event, usize)> { - if packet_data.len() < 14 { - return None; - } - - let eth_type = u16::from_be_bytes([packet_data[12], packet_data[13]]); + let (eth_hdr, ip_rest) = Ref::<&[u8], EthHdr>::from_prefix(packet_data).ok()?; let timestamp_us = time::SystemTime::now() .duration_since(time::UNIX_EPOCH) .ok()? .as_micros() as u64; - match eth_type { - 0x0800 => parse_ipv4(packet_data, timestamp_us), - 0x86DD => parse_ipv6(packet_data, timestamp_us), + match eth_hdr.ether_type.get() { + 0x0800 => parse_ipv4(ip_rest, timestamp_us), + 0x86DD => parse_ipv6(ip_rest, timestamp_us), _ => None, } } -fn parse_ipv4(packet_data: &[u8], timestamp_us: u64) -> Option<(Event, usize)> { - if packet_data.len() < 34 { - return None; - } - - let ip_header = &packet_data[14..]; - - let protocol_byte = ip_header[9]; - let protocol = unsafe { mem::transmute::(protocol_byte) }; - - let src_ip = u32::from_be_bytes([ip_header[12], ip_header[13], ip_header[14], ip_header[15]]); - let dst_ip = u32::from_be_bytes([ip_header[16], ip_header[17], ip_header[18], ip_header[19]]); +fn parse_ipv4(ip_bytes: &[u8], timestamp_us: u64) -> Option<(Event, usize)> { + let (ipv4_hdr, _) = Ref::<&[u8], Ipv4Hdr>::from_prefix(ip_bytes).ok()?; + let protocol_byte = ipv4_hdr.protocol; if protocol_byte != 6 && protocol_byte != 17 { return None; } + let protocol = unsafe { std::mem::transmute::(protocol_byte) }; - let ihl = (ip_header[0] & 0x0F) as usize * 4; - let total_len = u16::from_be_bytes([ip_header[2], ip_header[3]]) as u32; + let src_ip = ipv4_hdr.src_addr.get(); + let dst_ip = ipv4_hdr.dst_addr.get(); + let total_len = ipv4_hdr.total_len.get() as u32; + let ihl = (ipv4_hdr.version_ihl & 0x0F) as usize * 4; - if packet_data.len() < 14 + ihl + 4 { + if ip_bytes.len() < ihl + 4 { return None; } + let transport_bytes = &ip_bytes[ihl..]; - let transport_header = &ip_header[ihl..]; - let src_port = u16::from_be_bytes([transport_header[0], transport_header[1]]); - let dst_port = u16::from_be_bytes([transport_header[2], transport_header[3]]); - - let (tcp_flags, tcp_window_size, header_length) = if protocol_byte == 6 { - if packet_data.len() < 14 + ihl + 20 { - return None; - } - - let data_offset = (transport_header[12] >> 4) as u16 * 4; - let flags = TcpFlags::from_byte(transport_header[13]); - let window = u16::from_be_bytes([transport_header[14], transport_header[15]]); - - (flags, window, data_offset) - } else if protocol_byte == 17 { - (TcpFlags::default(), 0, 8) + let (src_port, dst_port, tcp_flags, tcp_window_size, header_length) = if protocol_byte == 6 { + let (tcp_hdr, _) = Ref::<&[u8], TcpHdr>::from_prefix(transport_bytes).ok()?; + let dof = tcp_hdr.data_off_flags.get(); + let data_offset = (dof >> 12) as u16 * 4; + let flags = TcpFlags::from_byte((dof & 0xFF) as u8); + (tcp_hdr.src_port.get(), tcp_hdr.dst_port.get(), flags, tcp_hdr.window.get(), data_offset) } else { - (TcpFlags::default(), 0, 0) + let (udp_hdr, _) = Ref::<&[u8], UdpHdr>::from_prefix(transport_bytes).ok()?; + (udp_hdr.src_port.get(), udp_hdr.dst_port.get(), TcpFlags::default(), 0u16, 8u16) }; let payload_length = total_len.saturating_sub(ihl as u32 + header_length as u32); - let payload_start = 14 + ihl + header_length as usize; + let payload_start = ETH_HDR_LEN + ihl + header_length as usize; - let event = IPv4Event { - protocol, - src_ip, - dst_ip, - src_port, - dst_port, - packet_length: total_len, - payload_length, - header_length, - timestamp_us, - tcp_flags, - tcp_window_size, - is_forward: false, - }; - - Some((Event::IPv4(event), payload_start)) + Some(( + Event::IPv4(IPv4Event { + protocol, + src_ip, + dst_ip, + src_port, + dst_port, + packet_length: total_len, + payload_length, + header_length, + timestamp_us, + tcp_flags, + tcp_window_size, + is_forward: false, + }), + payload_start, + )) } -fn parse_ipv6(packet_data: &[u8], timestamp_us: u64) -> Option<(Event, usize)> { - if packet_data.len() < 54 { - return None; - } - - let ip_header = &packet_data[14..]; - - let protocol_byte = ip_header[6]; - let protocol = unsafe { mem::transmute::(protocol_byte) }; - - let mut source_ip_bytes = [0u8; 16]; - source_ip_bytes.copy_from_slice(&ip_header[8..24]); - let src_ip = u128::from_be_bytes(source_ip_bytes); - - let mut dest_ip_bytes = [0u8; 16]; - dest_ip_bytes.copy_from_slice(&ip_header[24..40]); - let dst_ip = u128::from_be_bytes(dest_ip_bytes); +fn parse_ipv6(ip_bytes: &[u8], timestamp_us: u64) -> Option<(Event, usize)> { + let (ipv6_hdr, transport_bytes) = Ref::<&[u8], Ipv6Hdr>::from_prefix(ip_bytes).ok()?; + let protocol_byte = ipv6_hdr.next_hdr; if protocol_byte != 6 && protocol_byte != 17 { return None; } + let protocol = unsafe { std::mem::transmute::(protocol_byte) }; - let payload_len = u16::from_be_bytes([ip_header[4], ip_header[5]]) as u32; + let src_ip = u128::from_be_bytes(ipv6_hdr.src_addr); + let dst_ip = u128::from_be_bytes(ipv6_hdr.dst_addr); + let payload_len = ipv6_hdr.payload_len.get() as u32; let total_len = payload_len + 40; - if packet_data.len() < 54 + 4 { + if transport_bytes.len() < 4 { return None; } - let transport_header = &ip_header[40..]; - let src_port = u16::from_be_bytes([transport_header[0], transport_header[1]]); - let dst_port = u16::from_be_bytes([transport_header[2], transport_header[3]]); - - let (tcp_flags, tcp_window_size, header_length) = if protocol_byte == 6 { - if packet_data.len() < 54 + 20 { - return None; - } - - let data_offset = (transport_header[12] >> 4) as u16 * 4; - let flags = TcpFlags::from_byte(transport_header[13]); - let window = u16::from_be_bytes([transport_header[14], transport_header[15]]); - - (flags, window, data_offset) - } else if protocol_byte == 17 { - (TcpFlags::default(), 0, 8) + let (src_port, dst_port, tcp_flags, tcp_window_size, header_length) = if protocol_byte == 6 { + let (tcp_hdr, _) = Ref::<&[u8], TcpHdr>::from_prefix(transport_bytes).ok()?; + let dof = tcp_hdr.data_off_flags.get(); + let data_offset = (dof >> 12) as u16 * 4; + let flags = TcpFlags::from_byte((dof & 0xFF) as u8); + (tcp_hdr.src_port.get(), tcp_hdr.dst_port.get(), flags, tcp_hdr.window.get(), data_offset) } else { - (TcpFlags::default(), 0, 0) + let (udp_hdr, _) = Ref::<&[u8], UdpHdr>::from_prefix(transport_bytes).ok()?; + (udp_hdr.src_port.get(), udp_hdr.dst_port.get(), TcpFlags::default(), 0u16, 8u16) }; let payload_length = total_len.saturating_sub(40 + header_length as u32); - let payload_start = 14 + 40 + header_length as usize; + let payload_start = ETH_HDR_LEN + IPV6_HDR_LEN + header_length as usize; - let event = IPv6Event { - protocol, - src_ip, - dst_ip, - src_port, - dst_port, - packet_length: total_len, - payload_length, - header_length, - timestamp_us, - tcp_flags, - tcp_window_size, - is_forward: false, - }; - - Some((Event::IPv6(event), payload_start)) + Some(( + Event::IPv6(IPv6Event { + protocol, + src_ip, + dst_ip, + src_port, + dst_port, + packet_length: total_len, + payload_length, + header_length, + timestamp_us, + tcp_flags, + tcp_window_size, + is_forward: false, + }), + payload_start, + )) } pub fn format_ipv4(addr: u32) -> String { let bytes = addr.to_be_bytes(); - format!("{}.{}.{}.{}", bytes[0], bytes[1], bytes[2], bytes[3],) + format!("{}.{}.{}.{}", bytes[0], bytes[1], bytes[2], bytes[3]) } pub fn format_ipv6(addr: u128) -> String { @@ -182,6 +206,6 @@ pub fn format_ipv6(addr: u128) -> String { bytes[12], bytes[13], bytes[14], - bytes[15] + bytes[15], ) }