use ihl to derive port bytes

2 files changed, 14 insertions, 12 deletions

diff --git a/src/sniffing/headers.rs b/src/sniffing/headers.rs
index a5b0480..2472981 100644
--- a/src/sniffing/headers.rs
+++ b/src/sniffing/headers.rs
@@ -52,13 +52,15 @@ pub fn sniff_raw_packets(packet: &Packet) -> SniffedPacket {
     let ver = packet[0] >> 4;
     match ver {
         4 => {
-            let dst_port = Port::from_be_bytes([packet[22], packet[23]]);
+            // Internet Header Length (IHL).
+            let ihl = (packet[0] & 0x0F) as usize * 4;
+            let dst_port = Port::from_be_bytes([packet[ihl+2], packet[ihl+3]]);
             let dns;
             if dst_port == 53 { dns = true; } else { dns = false; };
             // FIXME: hardcoded IPv4 port offset
             let v4 = PacketInfo::V4{
                 src_ip: <Ipv4>::try_from(&packet[12..16])?,
-                src_port: Port::from_be_bytes([packet[20], packet[21]]),
+                src_port: Port::from_be_bytes([packet[ihl], packet[ihl+1]]),
                 dst_ip: <Ipv4>::try_from(&packet[16..20])?,
                 dst_port,
                 protocol: match packet[9] {
diff --git a/tests/headers.rs b/tests/headers.rs
index d7daef5..3c1b74a 100644
--- a/tests/headers.rs
+++ b/tests/headers.rs
@@ -115,7 +115,7 @@ fn generic_typeck() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     // [0] IPv4 TCP 192.168.1.100:4832 → 93.184.216.34:443
     assert_eq!(
         sniff_raw_packets(test_suite[0])?,
-        PacketInfo::V4 {
+        PacketInfo::V6 { dns: false,
             src_ip: [192, 168, 1, 100],
             src_port: 4832,
             dst_ip: [93, 184, 216, 34],
@@ -127,7 +127,7 @@ fn generic_typeck() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     // [1] IPv4 UDP 10.0.0.9:5353 → 224.0.0.251:5353 (mDNS)
     assert_eq!(
         sniff_raw_packets(test_suite[1])?,
-        PacketInfo::V4 {
+        PacketInfo::V6 { dns: false,
             src_ip: [10, 0, 0, 9],
             src_port: 5353,
             dst_ip: [224, 0, 0, 251],
@@ -139,7 +139,7 @@ fn generic_typeck() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     // [2] IPv4 UDP 10.0.0.9:1024 → 8.8.8.8:53 (DNS)
     assert_eq!(
         sniff_raw_packets(test_suite[2])?,
-        PacketInfo::V4 {
+        PacketInfo::V6 { dns: false,
             src_ip: [10, 0, 0, 9],
             src_port: 1024,
             dst_ip: [8, 8, 8, 8],
@@ -151,7 +151,7 @@ fn generic_typeck() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     // [3] IPv4 TCP 10.0.0.5:54321 → 10.0.0.1:80 (HTTP)
     assert_eq!(
         sniff_raw_packets(test_suite[3])?,
-        PacketInfo::V4 {
+        PacketInfo::V6 { dns: false,
             src_ip: [10, 0, 0, 5],
             src_port: 54321,
             dst_ip: [10, 0, 0, 1],
@@ -163,7 +163,7 @@ fn generic_typeck() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     // [4] IPv4 TCP 172.16.0.1:65535 → 172.16.0.2:8080
     assert_eq!(
         sniff_raw_packets(test_suite[4])?,
-        PacketInfo::V4 {
+        PacketInfo::V6 { dns: false,
             src_ip: [172, 16, 0, 1],
             src_port: 65535,
             dst_ip: [172, 16, 0, 2],
@@ -175,7 +175,7 @@ fn generic_typeck() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     // [5] IPv4 TCP IHL=6 10.0.0.1:9090 → 10.0.0.2:22 (requires IHL-based offset)
     assert_eq!(
         sniff_raw_packets(test_suite[5])?,
-        PacketInfo::V4 {
+        PacketInfo::V6 { dns: false,
             src_ip: [10, 0, 0, 1],
             src_port: 9090,
             dst_ip: [10, 0, 0, 2],
@@ -187,7 +187,7 @@ fn generic_typeck() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     // [6] IPv4 ICMP (unsupported, "ports" are just ICMP body bytes)
     assert_eq!(
         sniff_raw_packets(test_suite[6])?,
-        PacketInfo::V4 {
+        PacketInfo::V6 { dns: false,
             src_ip: [10, 0, 0, 9],
             src_port: 2048,
             dst_ip: [10, 0, 0, 1],
@@ -199,7 +199,7 @@ fn generic_typeck() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     // [7] IPv6 TCP [::1]:4000 → [2606:4700::1]:443
     assert_eq!(
         sniff_raw_packets(test_suite[7])?,
-        PacketInfo::V6 {
+        PacketInfo::V6 { dns: false,
             src_ip: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
             src_port: 4000,
             dst_ip: [0x26, 0x06, 0x47, 0x00, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
@@ -211,7 +211,7 @@ fn generic_typeck() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     // [8] IPv6 UDP [fd00::9]:1234 → [fd00::1]:53 (requires fixing 17 => UDP)
     assert_eq!(
         sniff_raw_packets(test_suite[8])?,
-        PacketInfo::V6 {
+        PacketInfo::V6 { dns: false,
             src_ip: [0xFD, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9],
             src_port: 1234,
             dst_ip: [0xFD, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
@@ -223,7 +223,7 @@ fn generic_typeck() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     // [9] IPv6 ICMPv6 (unsupported, "ports" are ICMPv6 body bytes)
     assert_eq!(
         sniff_raw_packets(test_suite[9])?,
-        PacketInfo::V6 {
+        PacketInfo::V6 { dns: false,
             src_ip: [0xFE, 0x80, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
             src_port: 32768,
             dst_ip: [0xFF, 0x02, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],