xxx fix snp test: make it idempotent and remove magic values

Currently the packet is not received back. I know that the existing
packet (application payload in udp packet in ipv4 packet) works.
Therefore, anything must be wrong with the creation of the packet.

Author: phip1611

uefi-test-runner/src/proto/network/mod.rs

@@ -1,11 +1,109 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
+use alloc::vec::Vec;
+use core::net::Ipv4Addr;
+
 pub fn test() {
     info!("Testing Network protocols");
 
-    pxe::test();
-    snp::test();
+    // Test idempotence of network stuff
+    for _ in 0..2 {
+        #[cfg(feature = "pxe")]
+        //pxe::test();
+        snp::test();
+    }
 }
 
+#[cfg(feature = "pxe")]
 mod pxe;
 mod snp;
+
+/// Computes the standard IPv4 or UDP checksum (one's complement)
+fn compute_ipv4_checksum(data: &[u8]) -> u16 {
+    let mut sum = 0u32;
+
+    for chunk in data.chunks(2) {
+        let word = if chunk.len() == 2 {
+            u16::from_be_bytes([chunk[0], chunk[1]]) as u32
+        } else {
+            (chunk[0] as u32) << 8
+        };
+        sum = sum.wrapping_add(word);
+    }
+
+    // Fold carry
+    while (sum >> 16) != 0 {
+        sum = (sum & 0xFFFF) + (sum >> 16);
+    }
+
+    !sum as u16
+}
+
+/// Builds the UDP pseudo-header used for UDP checksum calculation
+fn build_udp_pseudo_header(src_ip: &Ipv4Addr, dst_ip: &Ipv4Addr, udp_segment: &[u8]) -> Vec<u8> {
+    let mut pseudo = Vec::with_capacity(12 + udp_segment.len());
+
+    // Pseudo-header: src IP, dst IP, zero, protocol, UDP length
+    pseudo.extend_from_slice(&src_ip.octets());
+    pseudo.extend_from_slice(&dst_ip.octets());
+    pseudo.push(0);
+    pseudo.push(0x11); // Protocol = UDP
+    pseudo.extend_from_slice(&(udp_segment.len() as u16).to_be_bytes());
+
+    pseudo.extend_from_slice(udp_segment);
+
+    if pseudo.len() % 2 != 0 {
+        pseudo.push(0); // Pad to even length
+    }
+
+    pseudo
+}
+
+fn build_ipv4_udp_packet(
+    src_ip: Ipv4Addr,
+    dst_ip: Ipv4Addr,
+    src_port: u16,
+    dst_port: u16,
+    payload: &[u8],
+) -> Vec<u8> {
+    let ip_header_len = 20;
+    let udp_header_len = 8;
+    let total_len = ip_header_len + udp_header_len + payload.len();
+
+    let mut packet = Vec::with_capacity(total_len);
+
+    // === IPv4 Header ===
+    packet.push(0x45); // Version (4) + IHL (5)
+    packet.push(0x00); // DSCP/ECN
+    packet.extend_from_slice(&(total_len as u16).to_be_bytes()); // Total Length
+    packet.extend_from_slice(&0x0001u16.to_be_bytes()); // Identification
+    packet.extend_from_slice(&0x0000u16.to_be_bytes()); // Flags/Fragment offset
+    packet.push(0x40); // TTL
+    packet.push(0x11); // Protocol = UDP (17)
+    packet.extend_from_slice(&[0x00, 0x00]); // Checksum placeholder
+    packet.extend_from_slice(&src_ip.octets()); // Source IP
+    packet.extend_from_slice(&dst_ip.octets()); // Destination IP
+
+    // Compute IPv4 header checksum
+    let ip_checksum = compute_ipv4_checksum(&packet[..20]);
+    packet[10..12].copy_from_slice(&ip_checksum.to_be_bytes());
+
+    // === UDP Header ===
+    packet.extend_from_slice(&src_port.to_be_bytes()); // Source port
+    packet.extend_from_slice(&dst_port.to_be_bytes()); // Destination port
+    let udp_len = (udp_header_len + payload.len()) as u16;
+    packet.extend_from_slice(&udp_len.to_be_bytes()); // UDP length
+    packet.extend_from_slice(&[0x00, 0x00]); // UDP checksum placeholder
+
+    // === Payload ===
+    packet.extend_from_slice(payload);
+
+    // === UDP Checksum ===
+    let udp_offset = 20;
+    let udp_packet = &packet[udp_offset..];
+    let pseudo_header = build_udp_pseudo_header(&src_ip, &dst_ip, udp_packet);
+    let udp_checksum = compute_ipv4_checksum(&pseudo_header);
+    packet[udp_offset + 6..udp_offset + 8].copy_from_slice(&udp_checksum.to_be_bytes());
+
+    packet
+}

uefi-test-runner/src/proto/network/snp.rs

@@ -1,30 +1,94 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
-use uefi::proto::network::snp::{InterruptStatus, ReceiveFlags, SimpleNetwork};
-use uefi::proto::network::MacAddress;
+use crate::proto::network::build_ipv4_udp_packet;
+use alloc::string::ToString;
+use core::net::Ipv4Addr;
+use uefi::boot::{OpenProtocolAttributes, OpenProtocolParams};
+use uefi::proto::device_path::text::{AllowShortcuts, DisplayOnly};
+use uefi::proto::device_path::DevicePath;
+use uefi::proto::network::snp::{InterruptStatus, NetworkState, ReceiveFlags, SimpleNetwork};
 use uefi::{boot, Status};
 
+fn compute_ipv4_checksum(header: &[u8]) -> u16 {
+    assert_eq!(header.len() % 2, 0);
+    let sum = header
+        .chunks(2)
+        .map(|chunk| u16::from_be_bytes([chunk[0], chunk[1]]) as u32)
+        .sum::<u32>();
+
+    let carry_add = (sum & 0xFFFF) + (sum >> 16);
+    !(carry_add as u16)
+}
+
+fn build_ipv4_packet_with_payload(
+    src_ip: Ipv4Addr,
+    dest_ip: Ipv4Addr,
+    payload: [u8; 2],
+) -> [u8; 22] {
+    let mut packet = [0u8; 22];
+    let len = packet.len() as u16;
+
+    // IPv4 header
+    // Version = 4, IHL = 5
+    packet[0] = 0x45;
+    // DSCP/ECN
+    packet[1] = 0x00;
+    // Total length
+    packet[2..4].copy_from_slice(&(len.to_be_bytes()));
+    // Identification
+    packet[4..6].copy_from_slice(&0u16.to_be_bytes());
+    // Flags (DF), Fragment offset
+    packet[6..8].copy_from_slice(&0x4000u16.to_be_bytes());
+    // TTL
+    packet[8] = 0x40;
+    // Protocol (UDP)
+    packet[9] = 0x11;
+    // Checksum placeholder at [10..12]
+    packet[12..16].copy_from_slice(&src_ip.octets()); // Source IP
+    packet[16..20].copy_from_slice(&dest_ip.octets()); // Destination IP
+
+    // Calculate checksum
+    let checksum = compute_ipv4_checksum(&packet[0..20]);
+    packet[10..12].copy_from_slice(&checksum.to_be_bytes());
+
+    // Payload
+    packet[20] = payload[0];
+    packet[21] = payload[1];
+
+    packet
+}
+
 pub fn test() {
     info!("Testing the simple network protocol");
 
-    let handles = boot::find_handles::<SimpleNetwork>().unwrap_or_default();
+    let handles = boot::find_handles::<SimpleNetwork>().unwrap();
 
     for handle in handles {
-        let simple_network = boot::open_protocol_exclusive::<SimpleNetwork>(handle);
-        if simple_network.is_err() {
-            continue;
-        }
-        let simple_network = simple_network.unwrap();
+        // Buggy firmware; although it should be there, we have to test if the
+        // protocol is actually installed.
+        let simple_network = match boot::open_protocol_exclusive::<SimpleNetwork>(handle) {
+            Ok(snp) => snp,
+            Err(e) => {
+                log::debug!("Handle {handle:?} doesn't actually support SNP; skipping");
+                continue;
+            }
+        };
+        let simple_network_dvp = boot::open_protocol_exclusive::<DevicePath>(handle).unwrap();
+        debug!(
+            "Testing network device: {}",
+            simple_network_dvp
+                .to_string(DisplayOnly(false), AllowShortcuts(false))
+                .unwrap()
+        );
 
-        // Check shutdown
-        simple_network
-            .shutdown()
-            .expect("Failed to shutdown Simple Network");
+        // Check media
+        assert!(
+            simple_network.mode().media_present && simple_network.mode().media_present_supported
+        );
 
-        // Check stop
-        simple_network
-            .stop()
-            .expect("Failed to stop Simple Network");
+        // Ensure network is not started yet. If it is started, another test
+        // didn't clean up properly.
+        assert_eq!(simple_network.mode().state, NetworkState::STOPPED);
 
         // Check start
         simple_network
@@ -36,56 +100,50 @@ pub fn test() {
             .initialize(0, 0)
             .expect("Failed to initialize Simple Network");
 
-        simple_network.reset_statistics().unwrap();
+        // edk2 virtio-net driver does not support statistics, so
+        // allow UNSUPPORTED (same for collect_statistics below).
+        let res = simple_network.reset_statistics();
+        assert!(res == Ok(()) || res == Err(Status::UNSUPPORTED.into()));
 
         // Reading the interrupt status clears it
         simple_network.get_interrupt_status().unwrap();
 
         // Set receive filters
         simple_network
             .receive_filters(
-                ReceiveFlags::UNICAST | ReceiveFlags::MULTICAST | ReceiveFlags::BROADCAST,
+                ReceiveFlags::UNICAST | ReceiveFlags::BROADCAST,
                 ReceiveFlags::empty(),
                 false,
                 None,
             )
             .expect("Failed to set receive filters");
 
-        // Check media
-        if !simple_network.mode().media_present_supported || !simple_network.mode().media_present {
-            continue;
-        }
-
-        let payload = b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\
-            \x45\x00\
-            \x00\x21\
-            \x00\x01\
-            \x00\x00\
-            \x10\
-            \x11\
-            \x07\x6a\
-            \xc0\xa8\x11\x0f\
-            \xc0\xa8\x11\x02\
-            \x54\x45\
-            \x54\x44\
-            \x00\x0d\
-            \xa9\xe4\
-            \x04\x01\x02\x03\x04";
-
-        let dest_addr = MacAddress([0xffu8; 32]);
         assert!(!simple_network
             .get_interrupt_status()
             .unwrap()
             .contains(InterruptStatus::TRANSMIT));
 
-        // Send the frame
+        // Broadcast address (send to self)
+        let src_addr = simple_network.mode().current_address;
+        let dest_addr = simple_network.mode().broadcast_address;
+        const ETHER_TYPE_IPV4: u16 = 0x0800;
+
+        // IPv4 can be arbitrary, not used for routing here.
+        let src_ip = Ipv4Addr::new(192, 168, 17, 15);
+        let dst_ip = Ipv4Addr::new(192, 168, 17, 2);
+        let src_port = 0x5445; // "TE"
+        let dst_port = 0x5444; // "TD"
+        let payload = 0x1337_u16.to_ne_bytes();
+        let packet = build_ipv4_udp_packet(src_ip, dst_ip, src_port, dst_port, &payload);
+
+        // Send the frame to ourselves
         simple_network
             .transmit(
                 simple_network.mode().media_header_size as usize,
-                payload,
-                None,
+                &packet,
+                Some(src_addr),
                 Some(dest_addr),
-                Some(0x0800),
+                Some(ETHER_TYPE_IPV4),
             )
             .expect("Failed to transmit frame");
 
@@ -97,29 +155,36 @@ pub fn test() {
         {}
 
         // Attempt to receive a frame
-        let mut buffer = [0u8; 1500];
+        let mut recv_buffer = [0u8; 1500];
 
         info!("Waiting for the reception");
-        if simple_network.receive(&mut buffer, None, None, None, None)
+        if simple_network.receive(&mut recv_buffer, None, None, None, None)
             == Err(Status::NOT_READY.into())
         {
             boot::stall(1_000_000);
 
             simple_network
-                .receive(&mut buffer, None, None, None, None)
+                .receive(&mut recv_buffer, None, None, None, None)
                 .unwrap();
         }
 
-        assert_eq!(buffer[42..47], [4, 4, 3, 2, 1]);
-
         // Get stats
-        let stats = simple_network
-            .collect_statistics()
-            .expect("Failed to collect statistics");
-        info!("Stats: {:?}", stats);
-
-        // One frame should have been transmitted and one received
-        assert_eq!(stats.tx_total_frames().unwrap(), 1);
-        assert_eq!(stats.rx_total_frames().unwrap(), 1);
+        let res = simple_network.collect_statistics();
+        match res {
+            Ok(stats) => {
+                info!("Stats: {:?}", stats);
+
+                // One frame should have been transmitted and one received
+                assert_eq!(stats.tx_total_frames().unwrap(), 1);
+                assert_eq!(stats.rx_total_frames().unwrap(), 1);
+            }
+            Err(e) => {
+                if e == Status::UNSUPPORTED.into() {
+                    info!("Stats: unsupported.");
+                } else {
+                    panic!("{e}");
+                }
+            }
+        }
     }
 }