From 7c20b32b44db1e9e0a297b3722593bb0b98fbf04 Mon Sep 17 00:00:00 2001
From: Nikita Kalyazin <kalyazin@amazon.com>
Date: Fri, 4 Oct 2024 16:30:59 +0000
Subject: [PATCH] test(net): add test for tap offload features

This is a regression test for the following fix:

commit a9e5f1381bb30f9ad9fca11f2eade83a642711b5
Author: Nikita Kalyazin <kalyazin@amazon.co.uk>
Date:   Mon Sep 30 13:05:28 2024 +0000

    fix(net): set tap offload features on restore

The test verifies that tap offload features are configured for both
booted and restored VMs.

Signed-off-by: Nikita Kalyazin <kalyazin@amazon.com>
---
 tests/host_tools/udp_offload.py               | 58 +++++++++++++++++
 .../integration_tests/functional/test_net.py  | 62 +++++++++++++++++++
 2 files changed, 120 insertions(+)
 create mode 100644 tests/host_tools/udp_offload.py

diff --git a/tests/host_tools/udp_offload.py b/tests/host_tools/udp_offload.py
new file mode 100644
index 00000000000..e9ab6a93966
--- /dev/null
+++ b/tests/host_tools/udp_offload.py
@@ -0,0 +1,58 @@
+# Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+A utility for sending a UDP message with UDP oflload enabled.
+
+Inspired by the "TUN_F_CSUM is a must" chapter
+in https://blog.cloudflare.com/fr-fr/virtual-networking-101-understanding-tap/
+by Cloudflare.
+"""
+
+import socket
+import sys
+
+
+def eprint(*args, **kwargs):
+    """Print to stderr"""
+    print(*args, file=sys.stderr, **kwargs)
+
+
+# Define SOL_UDP and UDP_SEGMENT if not defined in the system headers
+try:
+    from socket import SOL_UDP, UDP_SEGMENT
+except ImportError:
+    SOL_UDP = 17  # Protocol number for UDP
+    UDP_SEGMENT = 103  # Option code for UDP segmentation (non-standard)
+
+# Get the IP and port from command-line arguments
+if len(sys.argv) != 3:
+    eprint("Usage: python3 udp_offload.py <ip_address> <port>")
+    sys.exit(1)
+
+ip_address = sys.argv[1]
+port = int(sys.argv[2])
+
+# Create a UDP socket
+sockfd = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+
+# Set the UDP segmentation option (UDP_SEGMENT) to 1400 bytes
+OPTVAL = 1400
+try:
+    sockfd.setsockopt(SOL_UDP, UDP_SEGMENT, OPTVAL)
+except (AttributeError, PermissionError):
+    eprint("Unable to set UDP_SEGMENT option")
+    sys.exit(1)
+
+# Set the destination address and port
+servaddr = (ip_address, port)
+
+# Send the message to the destination address
+MESSAGE = b"x"
+try:
+    sockfd.sendto(MESSAGE, servaddr)
+    print("Message sent successfully")
+except socket.error as e:
+    eprint(f"Error sending message: {e}")
+    sys.exit(1)
+
+sockfd.close()
diff --git a/tests/integration_tests/functional/test_net.py b/tests/integration_tests/functional/test_net.py
index c177b82267d..12980c727b2 100644
--- a/tests/integration_tests/functional/test_net.py
+++ b/tests/integration_tests/functional/test_net.py
@@ -78,3 +78,65 @@ def test_multi_queue_unsupported(uvm_plain):
             host_dev_name=tapname,
             guest_mac="AA:FC:00:00:00:01",
         )
+
+
+def run_udp_offload_test(vm):
+    """
+    - Start a socat UDP server in the guest.
+    - Try to send a UDP message with UDP offload enabled.
+
+    If tap offload features are not configured, an attempt to send a message will fail with EIO "Input/output error".
+    More info (search for "TUN_F_CSUM is a must"): https://blog.cloudflare.com/fr-fr/virtual-networking-101-understanding-tap/
+    """
+    port = "81"
+    out_filename = "/tmp/out.txt"
+    message = "x"
+
+    # Start a UDP server in the guest
+    # vm.ssh.check_output(f"nohup socat UDP-LISTEN:{port} - > {out_filename} &")
+    vm.ssh.check_output(
+        f"nohup socat UDP-LISTEN:{port} OPEN:{out_filename},creat > /dev/null 2>&1 &"
+    )
+
+    # Try to send a UDP message from host with UDP offload enabled
+    cmd = f"ip netns exec {vm.ssh_iface().netns} python3 ./host_tools/udp_offload.py {vm.ssh_iface().host} {port}"
+    ret = utils.run_cmd(cmd)
+
+    # Check that the transmission was successful
+    assert ret.returncode == 0, f"{ret.stdout=} {ret.stderr=}"
+
+    # Check that the server received the message
+    ret = vm.ssh.run(f"cat {out_filename}")
+    assert ret.stdout == message, f"{ret.stdout=} {ret.stderr=}"
+
+
+def test_tap_offload_booted(uvm_plain_any):
+    """
+    Verify that tap offload features are configured for a booted VM.
+    """
+    vm = uvm_plain_any
+    vm.spawn()
+    vm.basic_config()
+    vm.add_net_iface()
+    vm.start()
+
+    run_udp_offload_test(vm)
+
+
+def test_tap_offload_restored(microvm_factory, guest_kernel, rootfs_ubuntu_22):
+    """
+    Verify that tap offload features are configured for a restored VM.
+    """
+    src = microvm_factory.build(guest_kernel, rootfs_ubuntu_22, monitor_memory=False)
+    src.spawn()
+    src.basic_config()
+    src.add_net_iface()
+    src.start()
+    snapshot = src.snapshot_full()
+    src.kill()
+
+    dst = microvm_factory.build()
+    dst.spawn()
+    dst.restore_from_snapshot(snapshot, resume=True)
+
+    run_udp_offload_test(dst)