From 169153cec7cd5b625ba5e38945ffdc68a1f03d02 Mon Sep 17 00:00:00 2001
From: roblabla <unfiltered@roblab.la>
Date: Thu, 10 Apr 2025 16:10:08 +0200
Subject: [PATCH] Add integration test for split BPF

---
 test/integration-test/bpf/split.bpf.c        | 29 ++++++++++++++++
 test/integration-test/build.rs               |  1 +
 test/integration-test/src/lib.rs             |  1 +
 test/integration-test/src/tests.rs           |  1 +
 test/integration-test/src/tests/btf_split.rs | 36 ++++++++++++++++++++
 5 files changed, 68 insertions(+)
 create mode 100644 test/integration-test/bpf/split.bpf.c
 create mode 100644 test/integration-test/src/tests/btf_split.rs

diff --git a/test/integration-test/bpf/split.bpf.c b/test/integration-test/bpf/split.bpf.c
new file mode 100644
index 00000000..5697bb3b
--- /dev/null
+++ b/test/integration-test/bpf/split.bpf.c
@@ -0,0 +1,29 @@
+// clang-format off
+#include <vmlinux.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_core_read.h>
+// clang-format on
+
+char _license[] SEC("license") = "GPL";
+
+struct {
+  __uint(type, BPF_MAP_TYPE_ARRAY);
+  __type(key, __u32);
+  __type(value, __u64);
+  __uint(max_entries, 1);
+} output_map SEC(".maps");
+
+long set_output(__u64 value) {
+  __u32 key = 0;
+  return bpf_map_update_elem(&output_map, &key, &value, BPF_ANY);
+}
+
+// Try to access ip_tables structures. In most distros, ip_tables is compiled
+// and loaded as a separate module, making it a pretty good target.
+SEC("uprobe") int check_can_access_module(void *ctx) {
+  int is_successful =
+      bpf_core_type_exists(struct ipt_entry) &&
+      bpf_core_field_offset(struct ipt_entry, target_offset) != 0;
+  set_output(is_successful);
+  return is_successful;
+}
diff --git a/test/integration-test/build.rs b/test/integration-test/build.rs
index 47bfce48..694fd34a 100644
--- a/test/integration-test/build.rs
+++ b/test/integration-test/build.rs
@@ -67,6 +67,7 @@ fn main() -> Result<()> {
         ("main.bpf.c", false),
         ("multimap-btf.bpf.c", false),
         ("reloc.bpf.c", true),
+        ("split.bpf.c", false),
         ("text_64_64_reloc.c", false),
         ("variables_reloc.bpf.c", false),
     ];
diff --git a/test/integration-test/src/lib.rs b/test/integration-test/src/lib.rs
index 90277bb4..6b857b4c 100644
--- a/test/integration-test/src/lib.rs
+++ b/test/integration-test/src/lib.rs
@@ -8,6 +8,7 @@ pub const MULTIMAP_BTF: &[u8] =
 pub const RELOC_BPF: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/reloc.bpf.o"));
 pub const RELOC_BTF: &[u8] =
     include_bytes_aligned!(concat!(env!("OUT_DIR"), "/reloc.bpf.target.o"));
+pub const SPLIT_BPF: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/split.bpf.o"));
 pub const TEXT_64_64_RELOC: &[u8] =
     include_bytes_aligned!(concat!(env!("OUT_DIR"), "/text_64_64_reloc.o"));
 pub const VARIABLES_RELOC: &[u8] =
diff --git a/test/integration-test/src/tests.rs b/test/integration-test/src/tests.rs
index 9ca83669..4af52c83 100644
--- a/test/integration-test/src/tests.rs
+++ b/test/integration-test/src/tests.rs
@@ -1,5 +1,6 @@
 mod bpf_probe_read;
 mod btf_relocations;
+mod btf_split;
 mod elf;
 mod info;
 mod iter;
diff --git a/test/integration-test/src/tests/btf_split.rs b/test/integration-test/src/tests/btf_split.rs
new file mode 100644
index 00000000..ec777f92
--- /dev/null
+++ b/test/integration-test/src/tests/btf_split.rs
@@ -0,0 +1,36 @@
+//! Test to make sure loading split BTF (kernel module BTF) works properly.
+
+use aya::{EbpfLoader, maps::Array, programs::UProbe};
+
+#[test]
+fn rebase_tests() {
+    // First, check that we have ip_tables in the split btf.
+    if !matches!(std::fs::exists("/sys/kernel/btf/ip_tables"), Ok(true)) {
+        eprintln!(
+            "skipping test on kernel, as ip_tables is not loaded as an external kernel module."
+        );
+        return;
+    }
+    let mut bpf = EbpfLoader::new().load(crate::SPLIT_BPF).unwrap();
+    let program: &mut UProbe = bpf
+        .program_mut("check_can_access_module")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    program.load().unwrap();
+    program
+        .attach("trigger_btf_split_program", "/proc/self/exe", None, None)
+        .unwrap();
+
+    trigger_btf_split_program();
+
+    let output_map: Array<_, u64> = bpf.take_map("output_map").unwrap().try_into().unwrap();
+    let key = 0;
+    assert_eq!(output_map.get(&key, 0).unwrap(), 1)
+}
+
+#[unsafe(no_mangle)]
+#[inline(never)]
+pub extern "C" fn trigger_btf_split_program() {
+    core::hint::black_box(trigger_btf_split_program);
+}