From 2f227e74fe97ea45ab04397762cb16128971b186 Mon Sep 17 00:00:00 2001
From: Tuetuopay <tuetuopay@me.com>
Date: Tue, 21 Feb 2023 22:33:53 +0100
Subject: [PATCH] test: add AF_XDP integration test

---
 test/integration-ebpf/Cargo.toml       |  4 ++
 test/integration-ebpf/src/redirect.rs  | 43 ++++++++++++++++++
 test/integration-test/Cargo.toml       |  1 +
 test/integration-test/src/tests/mod.rs |  1 +
 test/integration-test/src/tests/xdp.rs | 60 ++++++++++++++++++++++++++
 5 files changed, 109 insertions(+)
 create mode 100644 test/integration-ebpf/src/redirect.rs
 create mode 100644 test/integration-test/src/tests/xdp.rs

diff --git a/test/integration-ebpf/Cargo.toml b/test/integration-ebpf/Cargo.toml
index 8668b91d1..e31cd26cc 100644
--- a/test/integration-ebpf/Cargo.toml
+++ b/test/integration-ebpf/Cargo.toml
@@ -22,3 +22,7 @@ path = "src/pass.rs"
 [[bin]]
 name = "test"
 path = "src/test.rs"
+
+[[bin]]
+name = "redirect"
+path = "src/redirect.rs"
diff --git a/test/integration-ebpf/src/redirect.rs b/test/integration-ebpf/src/redirect.rs
new file mode 100644
index 000000000..a077224df
--- /dev/null
+++ b/test/integration-ebpf/src/redirect.rs
@@ -0,0 +1,43 @@
+#![no_std]
+#![no_main]
+
+use aya_bpf::{
+    bindings::xdp_action,
+    macros::{map, xdp},
+    maps::{CpuMap, DevMap, DevMapHash, XskMap},
+    programs::XdpContext,
+};
+
+#[map]
+static SOCKS: XskMap = XskMap::with_max_entries(1, 0);
+#[map]
+static DEVS: DevMap = DevMap::with_max_entries(1, 0);
+#[map]
+static DEVS_HASH: DevMapHash = DevMapHash::with_max_entries(1, 0);
+#[map]
+static CPUS: CpuMap = CpuMap::with_max_entries(1, 0);
+
+#[xdp(name = "redirect_sock")]
+pub fn redirect_sock(_ctx: XdpContext) -> u32 {
+    SOCKS.redirect(0, xdp_action::XDP_ABORTED as u64)
+}
+
+#[xdp(name = "redirect_dev")]
+pub fn redirect_dev(_ctx: XdpContext) -> u32 {
+    DEVS.redirect(0, xdp_action::XDP_ABORTED as u64)
+}
+
+#[xdp(name = "redirect_dev_hash")]
+pub fn redirect_dev_hash(_ctx: XdpContext) -> u32 {
+    DEVS_HASH.redirect(10, xdp_action::XDP_ABORTED as u64)
+}
+
+#[xdp(name = "redirect_cpu")]
+pub fn redirect_cpu(_ctx: XdpContext) -> u32 {
+    CPUS.redirect(0, xdp_action::XDP_ABORTED as u64)
+}
+
+#[panic_handler]
+fn panic(_info: &core::panic::PanicInfo) -> ! {
+    unsafe { core::hint::unreachable_unchecked() }
+}
diff --git a/test/integration-test/Cargo.toml b/test/integration-test/Cargo.toml
index bfcc5d0b5..45dc43d88 100644
--- a/test/integration-test/Cargo.toml
+++ b/test/integration-test/Cargo.toml
@@ -19,4 +19,5 @@ object = { version = "0.30", default-features = false, features = ["std", "read_
 rbpf = "0.1.0"
 regex = "1"
 tempfile = "3.3.0"
+xsk-rs = "0.4"
 libtest-mimic = "0.6.0"
diff --git a/test/integration-test/src/tests/mod.rs b/test/integration-test/src/tests/mod.rs
index ddb1b5045..08b323c1f 100644
--- a/test/integration-test/src/tests/mod.rs
+++ b/test/integration-test/src/tests/mod.rs
@@ -9,6 +9,7 @@ pub mod load;
 pub mod rbpf;
 pub mod relocations;
 pub mod smoke;
+pub mod xdp;
 
 pub use integration_test_macros::integration_test;
 #[derive(Debug)]
diff --git a/test/integration-test/src/tests/xdp.rs b/test/integration-test/src/tests/xdp.rs
new file mode 100644
index 000000000..a79f6bb8c
--- /dev/null
+++ b/test/integration-test/src/tests/xdp.rs
@@ -0,0 +1,60 @@
+use std::{os::fd::AsRawFd, process::Command};
+
+use aya::{
+    include_bytes_aligned,
+    maps::XskMap,
+    programs::{Xdp, XdpFlags},
+    Bpf,
+};
+use xsk_rs::{
+    config::{LibbpfFlags, SocketConfigBuilder},
+    Socket, Umem,
+};
+
+use super::{integration_test, IntegrationTest};
+
+#[integration_test]
+fn af_xdp() {
+    let bytes = include_bytes_aligned!("../../../../target/bpfel-unknown-none/debug/redirect");
+    let mut bpf = Bpf::load(bytes).unwrap();
+    let xdp: &mut Xdp = bpf
+        .program_mut("redirect_sock")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    xdp.load().unwrap();
+    xdp.attach("lo", XdpFlags::default()).unwrap();
+
+    let (umem, mut descs) = Umem::new(Default::default(), 32.try_into().unwrap(), false).unwrap();
+    let sk_cfg = SocketConfigBuilder::new()
+        .libbpf_flags(LibbpfFlags::XSK_LIBBPF_FLAGS_INHIBIT_PROG_LOAD)
+        .build();
+    let (_tx, mut rx, fq_cq) = Socket::new(sk_cfg, &umem, &"lo".parse().unwrap(), 0).unwrap();
+    let (mut fq, _cq) = fq_cq.unwrap();
+
+    let mut socks: XskMap<_> = bpf.map_mut("SOCKS").unwrap().try_into().unwrap();
+    socks.set(0, rx.fd().as_raw_fd(), 0).unwrap();
+
+    // SAFETY: descs are from the same umem as the socket is tied to
+    // (valid for all further unsafe)
+    unsafe { fq.produce(&descs) };
+
+    // Expected to fail as our probe redirects all packets to the socket
+    let out = Command::new("ping")
+        .args(["-c", "1", "-w", "1", "127.0.0.1"])
+        .output()
+        .unwrap();
+    assert!(!out.status.success());
+
+    let n = unsafe { rx.consume(&mut descs) };
+    assert_eq!(n, 1);
+
+    let data = unsafe { umem.data(&descs[0]) };
+    let buf = data.contents();
+    let (eth, buf) = buf.split_at(14);
+    assert_eq!(eth[12..14], [0x08, 0x00]); // IP
+    let (ip, buf) = buf.split_at(20);
+    assert_eq!(ip[9], 1); // ICMP
+    let (icmp, _buf) = buf.split_at(8);
+    assert_eq!(icmp[0], 8); // Echo request
+}