From 20f7090c796a6e7a6e9802ced454114fc005a91e Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Sat, 16 Nov 2024 19:09:57 +0900 Subject: [PATCH] refactor: refactor the parsing logic --- src/net/parser/fluereflows.rs | 177 +++++++++++++++++++++------------- src/net/parser/keys.rs | 71 ++++++++------ src/net/parser/ports.rs | 35 +++++-- 3 files changed, 180 insertions(+), 103 deletions(-) diff --git a/src/net/parser/fluereflows.rs b/src/net/parser/fluereflows.rs index ab15a67..42598fc 100644 --- a/src/net/parser/fluereflows.rs +++ b/src/net/parser/fluereflows.rs @@ -7,6 +7,7 @@ use log::trace; use pnet::packet::{ arp::ArpPacket, ethernet::{EtherTypes, EthernetPacket}, + ip::IpNextHeaderProtocols, ipv4::Ipv4Packet, ipv6::Ipv6Packet, udp::UdpPacket, @@ -38,67 +39,74 @@ pub fn parse_fluereflow(packet: pcap::Packet) -> Result<(usize, [u8; 9], FluereR let is_udp: bool = match ethernet_packet_unpack.get_ethertype() { EtherTypes::Ipv6 => { - let i = Ipv6Packet::new(ethernet_packet_unpack.payload()).unwrap(); - if i.payload().is_empty() { - return Err(NetError::EmptyPacket); - } - let is_udp = UdpPacket::new(i.payload()).is_some(); - - is_udp + let i = match Ipv6Packet::new(ethernet_packet_unpack.payload()) { + Some(packet) => packet, + None => return Err(NetError::InvalidPacket), + }; + UdpPacket::new(i.payload()).is_some() } EtherTypes::Ipv4 => { - let i = Ipv4Packet::new(ethernet_packet_unpack.payload()).unwrap(); - if i.payload().is_empty() { - return Err(NetError::EmptyPacket); - } - - let is_udp = UdpPacket::new(i.payload()).is_some(); - - is_udp + let i = match Ipv4Packet::new(ethernet_packet_unpack.payload()) { + Some(packet) => packet, + None => return Err(NetError::InvalidPacket), + }; + UdpPacket::new(i.payload()).is_some() } _ => false, }; + let mut decapsulated_data: Option> = None; if is_udp { let udp_payload = match ethernet_packet_unpack.get_ethertype() { EtherTypes::Ipv6 => { - let i = Ipv6Packet::new(ethernet_packet_unpack.payload()).unwrap(); - if i.payload().is_empty() { - return Err(NetError::EmptyPacket); + let i = match Ipv6Packet::new(ethernet_packet_unpack.payload()) { + Some(packet) => packet, + None => return Err(NetError::InvalidPacket), + }; + + match UdpPacket::new(i.payload()) { + Some(udp) => { + trace!("UDP payload length: {}", udp.payload().len()); + udp.payload().to_vec() + } + None => return Err(NetError::InvalidPacket), } - - UdpPacket::new(i.payload()).unwrap().payload().to_vec() } EtherTypes::Ipv4 => { - let i = Ipv4Packet::new(ethernet_packet_unpack.payload()).unwrap(); - if i.payload().is_empty() { - return Err(NetError::EmptyPacket); + let i = match Ipv4Packet::new(ethernet_packet_unpack.payload()) { + Some(packet) => packet, + None => return Err(NetError::InvalidPacket), + }; + + match UdpPacket::new(i.payload()) { + Some(udp) => { + trace!("UDP payload length: {}", udp.payload().len()); + udp.payload().to_vec() + } + None => return Err(NetError::InvalidPacket), } - - UdpPacket::new(i.payload()).unwrap().payload().to_vec() } _ => Vec::new(), }; - if udp_payload.is_empty() { - return Err(NetError::EmptyPacket); + + // Only check if UDP payload is empty when we expect it to have content + if udp_payload.is_empty() && is_udp { + trace!("Empty UDP payload detected"); } - //UdpPacket::new(ethernet_packet_unpack.payload()).unwrap().payload().to_vec(); - //println!("UDP payload: {:?}", udp_payload); + decapsulated_data = decapsulate_vxlan(&udp_payload); } - let ethernet_packet_decapsulated = if let Some(data) = &decapsulated_data { + let ethernet_packet = if let Some(data) = &decapsulated_data { match EthernetPacket::new(data) { - None => return Err(NetError::EmptyPacket), + None => ethernet_packet_unpack, // Fall back to original packet if decapsulation fails Some(e) => e, } } else { ethernet_packet_unpack }; - let ethernet_packet = ethernet_packet_decapsulated; - let time = parse_microseconds( packet.header.ts.tv_sec as u64, packet.header.ts.tv_usec as u64, @@ -106,45 +114,46 @@ pub fn parse_fluereflow(packet: pcap::Packet) -> Result<(usize, [u8; 9], FluereR let record_result = match ethernet_packet.get_ethertype() { EtherTypes::Ipv4 => { - let i = Ipv4Packet::new(ethernet_packet.payload()).unwrap(); - - if i.payload().is_empty() { - return Err(NetError::EmptyPacket); - } - + let i = match Ipv4Packet::new(ethernet_packet.payload()) { + Some(packet) => packet, + None => { + trace!("Failed to parse IPv4 packet"); + return Err(NetError::InvalidPacket); + } + }; ipv4_packet(time, i) } EtherTypes::Ipv6 => { - let i = Ipv6Packet::new(ethernet_packet.payload()).unwrap(); - if i.payload().is_empty() { - return Err(NetError::EmptyPacket); - } - + let i = match Ipv6Packet::new(ethernet_packet.payload()) { + Some(packet) => packet, + None => { + trace!("Failed to parse IPv6 packet"); + return Err(NetError::InvalidPacket); + } + }; ipv6_packet(time, i) } EtherTypes::Arp => { - let i = ArpPacket::new(ethernet_packet.payload()).unwrap(); - //if i.payload().is_empty() { - // return Err(NetError::EmptyPacket); - //} - + let i = match ArpPacket::new(ethernet_packet.payload()) { + Some(packet) => packet, + None => { + trace!("Failed to parse ARP packet"); + return Err(NetError::InvalidPacket); + } + }; arp_packet(time, i) } - _ => { + trace!("Unknown EtherType: {}", ethernet_packet.get_ethertype()); return Err(NetError::UnknownEtherType( ethernet_packet.get_ethertype().to_string(), - )) + )); } - }; + }?; - match record_result { - Ok(_) => {} - Err(e) => return Err(e), - } - let (doctets, raw_flags, record) = record_result.unwrap(); - Ok((doctets, raw_flags, record)) + Ok(record_result) } + fn arp_packet(time: u64, packet: ArpPacket) -> Result<(usize, [u8; 9], FluereRecord), NetError> { let src_ip = packet.get_sender_proto_addr(); let dst_ip = packet.get_target_proto_addr(); @@ -197,16 +206,54 @@ fn ipv4_packet(time: u64, packet: Ipv4Packet) -> Result<(usize, [u8; 9], FluereR let protocol = packet.get_next_level_protocol().0; let src_ip = packet.get_source(); let dst_ip = packet.get_destination(); - trace!("src_ip: {:?}", src_ip); - trace!("dst_ip: {:?}", dst_ip); - // ports parsing - let (src_port, dst_port) = parse_ports(protocol, packet.payload())?; + // Special handling for DNS over UDP + if packet.get_next_level_protocol() == IpNextHeaderProtocols::Udp { + if let Some(udp) = UdpPacket::new(packet.payload()) { + if udp.get_destination() == 53 || udp.get_source() == 53 { + return Ok(( + packet.packet_size(), + [0; 9], // DNS doesn't use TCP flags + FluereRecord::new( + std::net::IpAddr::V4(src_ip), + std::net::IpAddr::V4(dst_ip), + 0, + 0, + time, + time, + udp.get_source(), + udp.get_destination(), + udp.packet_size() as u32, + udp.packet_size() as u32, + packet.get_ttl(), + packet.get_ttl(), + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 17, // UDP + 0, // No TOS for DNS + ), + )); + } + } + } - // TCP flags Fin Syn Rst Psh Ack Urg Ece Cwr Ns + // Continue with normal packet processing... + let (src_port, dst_port) = parse_ports(protocol, packet.payload()).unwrap_or((0, 0)); + + // TCP flags let flags = parse_flags(protocol, packet.payload()); - // Autonomous system number of the source and destination, either origin or peer let doctets = packet.packet_size(); let tos_convert_result = dscp_to_tos(packet.get_dscp()); let tos = tos_convert_result.unwrap_or(0); @@ -261,7 +308,7 @@ fn ipv6_packet(time: u64, packet: Ipv6Packet) -> Result<(usize, [u8; 9], FluereR //first six bits in the 8-bit Traffic Class field let dscp = packet.get_traffic_class() >> 2; let tos_convert_result = dscp_to_tos(dscp); - let tos = tos_convert_result.unwrap_or(0); + let tos = tos_convert_result.unwrap_or_default(); Ok(( doctets, diff --git a/src/net/parser/keys.rs b/src/net/parser/keys.rs index bb024c8..536da3f 100644 --- a/src/net/parser/keys.rs +++ b/src/net/parser/keys.rs @@ -7,14 +7,11 @@ use crate::net::NetError; use pcap; use pnet::packet::arp::ArpPacket; -use pnet::packet::dns::DnsPacket; use pnet::packet::ethernet::{EtherTypes, EthernetPacket}; use pnet::packet::gre::GrePacket; -use pnet::packet::icmp::IcmpPacket; use pnet::packet::icmpv6::Icmpv6Packet; use pnet::packet::ipv4::Ipv4Packet; use pnet::packet::ipv6::Ipv6Packet; -use pnet::packet::tcp::TcpPacket; use pnet::packet::udp::UdpPacket; use pnet::packet::vlan::VlanPacket; use pnet::packet::Packet; @@ -38,16 +35,15 @@ pub fn parse_keys(packet: pcap::Packet) -> Result<(Key, Key), NetError> { return Err(NetError::EmptyPacket); } trace!("Parsing ethernet packet"); - let ethernet_packet_raw = EthernetPacket::new(packet.data); - let ethernet_packet_unpack = match ethernet_packet_raw { + let ethernet_packet = match EthernetPacket::new(packet.data) { + Some(p) => p, None => return Err(NetError::EmptyPacket), - Some(e) => e, }; trace!("Parsed ethernet packet"); - let is_udp: bool = match ethernet_packet_unpack.get_ethertype() { + let is_udp: bool = match ethernet_packet.get_ethertype() { EtherTypes::Ipv6 => { - let i = Ipv6Packet::new(ethernet_packet_unpack.payload()); + let i = Ipv6Packet::new(ethernet_packet.payload()); if i.is_none() { return Err(NetError::EmptyPacket); } @@ -56,7 +52,7 @@ pub fn parse_keys(packet: pcap::Packet) -> Result<(Key, Key), NetError> { is_udp } EtherTypes::Ipv4 => { - let i = Ipv4Packet::new(ethernet_packet_unpack.payload()); + let i = Ipv4Packet::new(ethernet_packet.payload()); if i.is_none() { return Err(NetError::EmptyPacket); } @@ -66,7 +62,7 @@ pub fn parse_keys(packet: pcap::Packet) -> Result<(Key, Key), NetError> { is_udp } EtherTypes::Arp => { - let i = ArpPacket::new(ethernet_packet_unpack.payload()); + let i = ArpPacket::new(ethernet_packet.payload()); if i.is_none() { return Err(NetError::EmptyPacket); } @@ -83,9 +79,9 @@ pub fn parse_keys(packet: pcap::Packet) -> Result<(Key, Key), NetError> { let mut decapsulated_data: Option> = None; if is_udp { - let udp_payload = match ethernet_packet_unpack.get_ethertype() { + let udp_payload = match ethernet_packet.get_ethertype() { EtherTypes::Ipv6 => { - let i = Ipv6Packet::new(ethernet_packet_unpack.payload()); + let i = Ipv6Packet::new(ethernet_packet.payload()); if i.is_none() { return Err(NetError::EmptyPacket); } @@ -96,7 +92,7 @@ pub fn parse_keys(packet: pcap::Packet) -> Result<(Key, Key), NetError> { .to_vec() } EtherTypes::Ipv4 => { - let i = Ipv4Packet::new(ethernet_packet_unpack.payload()); + let i = Ipv4Packet::new(ethernet_packet.payload()); if i.is_none() { return Err(NetError::EmptyPacket); } @@ -107,7 +103,7 @@ pub fn parse_keys(packet: pcap::Packet) -> Result<(Key, Key), NetError> { .to_vec() } EtherTypes::Arp => { - let i = ArpPacket::new(ethernet_packet_unpack.payload()); + let i = ArpPacket::new(ethernet_packet.payload()); if i.is_none() { return Err(NetError::EmptyPacket); } @@ -135,7 +131,7 @@ pub fn parse_keys(packet: pcap::Packet) -> Result<(Key, Key), NetError> { Some(e) => e, } } else { - ethernet_packet_unpack + ethernet_packet }; let ethernet_packet = ethernet_packet_decapsulated; @@ -191,25 +187,44 @@ pub fn parse_keys(packet: pcap::Packet) -> Result<(Key, Key), NetError> { arp_keys(i.unwrap())? } _ => { - //Check Every possible key parsing method - let parse_test_ipv4 = ipv4_keys(Ipv4Packet::new(ethernet_packet.payload()).unwrap()); - let parse_test_ipv6 = ipv6_keys(Ipv6Packet::new(ethernet_packet.payload()).unwrap()); - let parse_test_arp = arp_keys(ArpPacket::new(ethernet_packet.payload()).unwrap()); - let parse_test_vlan = vlan_keys(VlanPacket::new(ethernet_packet.payload()).unwrap()); - // let parse_test_gre = gre_keys(GrePacket::new(ethernet_packet.payload()).unwrap()); - // let parse_test_icmpv6 = icmpv6_keys(Icmpv6Packet::new(ethernet_packet.payload()).unwrap()); - //let parse_test_tcp = tcp_keys(TcpPacket::new(ethernet_packet.payload()).unwrap()); - //let parse_test_udp = udp_keys(UdpPacket::new(ethernet_packet.payload()).unwrap()); - // let parse_test_dns = dns_keys(DnsPacket::new(ethernet_packet.payload()).unwrap()); + // Try each parsing method without unwrapping + let parse_test_ipv4 = if let Some(packet) = Ipv4Packet::new(ethernet_packet.payload()) { + ipv4_keys(packet) + } else { + Err(NetError::InvalidPacket) + }; + + let parse_test_ipv6 = if let Some(packet) = Ipv6Packet::new(ethernet_packet.payload()) { + ipv6_keys(packet) + } else { + Err(NetError::InvalidPacket) + }; + + let parse_test_arp = if let Some(packet) = ArpPacket::new(ethernet_packet.payload()) { + arp_keys(packet) + } else { + Err(NetError::InvalidPacket) + }; + + let parse_test_vlan = if let Some(packet) = VlanPacket::new(ethernet_packet.payload()) { + vlan_keys(packet) + } else { + Err(NetError::InvalidPacket) + }; trace!("parse_test_ipv4: {:?}", parse_test_ipv4); trace!("parse_test_ipv6: {:?}", parse_test_ipv6); trace!("parse_test_arp: {:?}", parse_test_arp); trace!("parse_test_vlan: {:?}", parse_test_vlan); - return Err(NetError::UnknownEtherType( - ethernet_packet.get_ethertype().to_string(), - )); + // Try to use the first successful parse result + parse_test_ipv4 + .or(parse_test_ipv6) + .or(parse_test_arp) + .or(parse_test_vlan) + .or(Err(NetError::UnknownEtherType( + ethernet_packet.get_ethertype().to_string(), + )))? } }; trace!("Parsed keys"); diff --git a/src/net/parser/ports.rs b/src/net/parser/ports.rs index a778c1c..373ad12 100644 --- a/src/net/parser/ports.rs +++ b/src/net/parser/ports.rs @@ -5,21 +5,36 @@ use pnet::packet::{tcp::TcpPacket, udp::UdpPacket}; pub fn parse_ports(protocol: u8, payload: &[u8]) -> Result<(u16, u16), NetError> { match protocol { - 58 => Ok((0, 0)), - 17 => match UdpPacket::new(payload) { - Some(udp) => Ok((udp.get_source(), udp.get_destination())), - None => Err(NetError::InvalidPacket), - }, + 0 => Ok((0, 0)), // IPv6 Hop-by-Hop Option + 1 => Ok((0, 0)), // ICMP + 2 => Ok((0, 0)), // IGMP + 4 => Ok((0, 0)), // IPv4 encapsulation 6 => match TcpPacket::new(payload) { Some(tcp) => Ok((tcp.get_source(), tcp.get_destination())), None => Err(NetError::InvalidPacket), }, - 2 => Ok((0, 0)), - 1 => Ok((0, 0)), - 0 => Ok((0, 0)), + 17 => match UdpPacket::new(payload) { + Some(udp) => { + if udp.get_source() == 53 || udp.get_destination() == 53 { + Ok((udp.get_source(), udp.get_destination())) + } else { + Ok((udp.get_source(), udp.get_destination())) + } + } + None => Err(NetError::InvalidPacket), + }, + 47 => Ok((0, 0)), // GRE + 50 => Ok((0, 0)), // ESP + 51 => Ok((0, 0)), // AH + 58 => Ok((0, 0)), // ICMPv6 + // DNS typically runs on UDP port 53 + 53 => match UdpPacket::new(payload) { + Some(udp) => Ok((udp.get_source(), udp.get_destination())), + None => Ok((53, 53)), // Default DNS ports if packet parsing fails + }, _ => { - debug!("Unknown protocol: {}", protocol); - Err(NetError::UnknownProtocol(protocol)) + debug!("Unknown protocol: {}, returning default ports", protocol); + Ok((0, 0)) // Return default ports instead of error for unknown protocols } } }