add basic tcx integration tests

Add basic tcx integration tests to ensure the new
Mprog API is working as expected.

Signed-off-by: astoycos <astoycos@gmail.com>

Cargo.toml

@@ -98,6 +98,7 @@ tokio = { version = "1.24.0", default-features = false }
 which = { version = "6.0.0", default-features = false }
 xdpilone = { version = "1.0.5", default-features = false }
 xtask = { path = "xtask", default-features = false }
+chrono = { version = "0.4" }
 
 [profile.dev]
 panic = "abort"

test/integration-ebpf/Cargo.toml

@@ -37,6 +37,10 @@ path = "src/pass.rs"
 name = "test"
 path = "src/test.rs"
 
+[[bin]]
+name = "tcx"
+path = "src/tcx.rs"
+
 [[bin]]
 name = "relocations"
 path = "src/relocations.rs"

test/integration-ebpf/src/tcx.rs

@@ -0,0 +1,79 @@
+#![no_std]
+#![no_main]
+
+use core::mem;
+
+use aya_ebpf::{
+    bindings::tcx_action_base::{TCX_NEXT, TCX_PASS},
+    macros::classifier,
+    programs::TcContext,
+};
+use aya_log_ebpf::info;
+use network_types::{
+    eth::{EthHdr, EtherType},
+    ip::{IpProto, Ipv4Hdr},
+    udp::UdpHdr,
+};
+
+#[no_mangle]
+static ORDER: i32 = 0;
+
+// Gives us raw pointers to a specific offset in the packet
+#[inline(always)]
+unsafe fn ptr_at<T>(ctx: &TcContext, offset: usize) -> Result<*mut T, i64> {
+    let start = ctx.data();
+    let end = ctx.data_end();
+    let len = mem::size_of::<T>();
+
+    if start + offset + len > end {
+        return Err(TCX_PASS.into());
+    }
+    Ok((start + offset) as *mut T)
+}
+
+#[classifier]
+pub fn tcx_order(ctx: TcContext) -> i32 {
+    match try_tcxtest(ctx) {
+        Ok(ret) => ret,
+        Err(_ret) => TCX_PASS,
+    }
+}
+
+fn try_tcxtest(ctx: TcContext) -> Result<i32, i64> {
+    let eth_hdr: *const EthHdr = unsafe { ptr_at(&ctx, 0) }?;
+    let order = unsafe { core::ptr::read_volatile(&ORDER) };
+    match unsafe { *eth_hdr }.ether_type {
+        EtherType::Ipv4 => {
+            let ipv4_hdr: *const Ipv4Hdr = unsafe { ptr_at(&ctx, EthHdr::LEN)? };
+            let saddr = u32::from_be(unsafe { (*ipv4_hdr).src_addr });
+            let daddr = u32::from_be(unsafe { (*ipv4_hdr).dst_addr });
+            match unsafe { (*ipv4_hdr).proto } {
+                IpProto::Udp => {
+                    let udphdr: *const UdpHdr =
+                        unsafe { ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN) }?;
+                    let dport = u16::from_be(unsafe { (*udphdr).dest });
+                    let sport = u16::from_be(unsafe { (*udphdr).source });
+                    info!(
+                        &ctx,
+                        "order: {}, cookie: ({:i}, {:i}, {}, {})",
+                        order,
+                        daddr,
+                        saddr,
+                        dport,
+                        sport
+                    );
+
+                    Ok(TCX_NEXT)
+                }
+                _ => Ok(TCX_PASS),
+            }
+        }
+        _ => Ok(TCX_PASS),
+    }
+}
+
+#[cfg(not(test))]
+#[panic_handler]
+fn panic(_info: &core::panic::PanicInfo) -> ! {
+    loop {}
+}

test/integration-test/Cargo.toml

@@ -27,6 +27,7 @@ test-case = { workspace = true }
 test-log = { workspace = true, features = ["log"] }
 tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] }
 xdpilone = { workspace = true }
+chrono = { workspace = true }
 
 [build-dependencies]
 cargo_metadata = { workspace = true }

test/integration-test/src/lib.rs

@@ -14,6 +14,7 @@ pub const LOG: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/log"));
 pub const MAP_TEST: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/map_test"));
 pub const NAME_TEST: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/name_test"));
 pub const PASS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/pass"));
+pub const TCX: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/tcx"));
 pub const TEST: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/test"));
 pub const RELOCATIONS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/relocations"));
 pub const TWO_PROGS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/two_progs"));

test/integration-test/src/tests.rs

@@ -8,4 +8,5 @@ mod rbpf;
 mod relocations;
 mod ring_buf;
 mod smoke;
+mod tcx;
 mod xdp;

test/integration-test/src/tests/log.rs

@@ -5,6 +5,7 @@ use std::{
 
 use aya::{programs::UProbe, Ebpf};
 use aya_log::EbpfLogger;
+use chrono::{DateTime, Utc};
 use log::{Level, Log, Record};
 use test_log::test;
 
@@ -14,8 +15,8 @@ pub extern "C" fn trigger_ebpf_program() {
     core::hint::black_box(trigger_ebpf_program);
 }
 
-struct TestingLogger<F> {
+pub(crate) struct TestingLogger<F> {
-    log: F,
+    pub(crate) log: F,
 }
 
 impl<F: Send + Sync + Fn(&Record)> Log for TestingLogger<F> {
@@ -32,10 +33,11 @@ impl<F: Send + Sync + Fn(&Record)> Log for TestingLogger<F> {
 }
 
 #[derive(Debug, PartialEq)]
-struct CapturedLog<'a> {
+pub(crate) struct CapturedLog<'a> {
-    pub body: Cow<'a, str>,
+    pub(crate) body: Cow<'a, str>,
-    pub level: Level,
+    pub(crate) level: Level,
-    pub target: Cow<'a, str>,
+    pub(crate) target: Cow<'a, str>,
+    pub(crate) timestamp: Option<DateTime<Utc>>,
 }
 
 #[test(tokio::test)]
@@ -54,6 +56,7 @@ async fn log() {
                         body: format!("{}", record.args()).into(),
                         level: record.level(),
                         target: record.target().to_string().into(),
+                        timestamp: None,
                     });
                 },
             },
@@ -91,6 +94,7 @@ async fn log() {
             body: "Hello from eBPF!".into(),
             level: Level::Debug,
             target: "log".into(),
+            timestamp: None,
         }),
     );
 
@@ -100,6 +104,7 @@ async fn log() {
             body: "69, 420, wao, 77616f".into(),
             level: Level::Error,
             target: "log".into(),
+            timestamp: None,
         })
     );
 
@@ -109,6 +114,7 @@ async fn log() {
             body: "ip structs, without format hint: ipv4: 10.0.0.1, ipv6: 2001:db8::1".into(),
             level: Level::Info,
             target: "log".into(),
+            timestamp: None,
         })
     );
 
@@ -118,6 +124,7 @@ async fn log() {
             body: "ip structs, with format hint: ipv4: 10.0.0.1, ipv6: 2001:db8::1".into(),
             level: Level::Info,
             target: "log".into(),
+            timestamp: None,
         })
     );
 
@@ -127,6 +134,7 @@ async fn log() {
             body: "ip enums, without format hint: ipv4: 10.0.0.1, ipv6: 2001:db8::1".into(),
             level: Level::Info,
             target: "log".into(),
+            timestamp: None,
         })
     );
 
@@ -136,6 +144,7 @@ async fn log() {
             body: "ip enums, with format hint: ipv4: 10.0.0.1, ipv6: 2001:db8::1".into(),
             level: Level::Info,
             target: "log".into(),
+            timestamp: None,
         })
     );
 
@@ -145,6 +154,7 @@ async fn log() {
             body: "ip as bits: ipv4: 10.0.0.1".into(),
             level: Level::Info,
             target: "log".into(),
+            timestamp: None,
         })
     );
 
@@ -154,6 +164,7 @@ async fn log() {
             body: "ip as octets: ipv4: 10.0.0.1, ipv6: 2001:db8::1".into(),
             level: Level::Info,
             target: "log".into(),
+            timestamp: None,
         })
     );
 
@@ -163,6 +174,7 @@ async fn log() {
             body: "mac lc: 04:20:06:09:00:40, mac uc: 04:20:06:09:00:40".into(),
             level: Level::Trace,
             target: "log".into(),
+            timestamp: None,
         })
     );
 
@@ -172,6 +184,7 @@ async fn log() {
             body: "hex lc: 2f, hex uc: 2F".into(),
             level: Level::Warn,
             target: "log".into(),
+            timestamp: None,
         })
     );
 
@@ -181,6 +194,7 @@ async fn log() {
             body: "hex lc: deadbeef, hex uc: DEADBEEF".into(),
             level: Level::Debug,
             target: "log".into(),
+            timestamp: None,
         })
     );
 
@@ -190,6 +204,7 @@ async fn log() {
             body: "42 43 44 45".into(),
             level: Level::Debug,
             target: "log".into(),
+            timestamp: None,
         })
     );
 

test/integration-test/src/tests/tcx.rs

@@ -0,0 +1,171 @@
+use std::{
+    net::UdpSocket,
+    sync::{Arc, Mutex},
+    time::Duration,
+};
+
+use aya::{
+    programs::{tc::TcAttachOptions, LinkOrder, SchedClassifier, TcAttachType},
+    util::KernelVersion,
+    Ebpf, EbpfLoader,
+};
+use aya_log::EbpfLogger;
+use chrono::Utc;
+use log::{debug, Record};
+use test_log::test;
+
+use crate::{
+    tests::log::{CapturedLog, TestingLogger},
+    utils::NetNsGuard,
+};
+
+fn setup_logs(loader: &mut Ebpf, logs: &Arc<Mutex<Vec<CapturedLog<'static>>>>) {
+    let captured_logs = logs.clone();
+    EbpfLogger::init_with_logger(
+        loader,
+        TestingLogger {
+            log: move |record: &Record| {
+                let mut logs = captured_logs.lock().unwrap();
+                logs.push(CapturedLog {
+                    body: format!("{}", record.args()).into(),
+                    level: record.level(),
+                    target: record.target().to_string().into(),
+                    timestamp: Some(Utc::now()),
+                });
+            },
+        },
+    )
+    .unwrap();
+}
+
+#[test(tokio::test)]
+async fn tcx_ordering() {
+    let kernel_version = KernelVersion::current().unwrap();
+    if kernel_version < KernelVersion::new(6, 6, 0) {
+        eprintln!("skipping tcx_ordering test on kernel {kernel_version:?}");
+        return;
+    }
+
+    let _netns = NetNsGuard::new();
+
+    let mut program0 = EbpfLoader::new()
+        .set_global("ORDER", &0, true)
+        .load(crate::TCX)
+        .unwrap();
+    let mut program1 = EbpfLoader::new()
+        .set_global("ORDER", &1, true)
+        .load(crate::TCX)
+        .unwrap();
+    let mut program2 = EbpfLoader::new()
+        .set_global("ORDER", &2, true)
+        .load(crate::TCX)
+        .unwrap();
+    let mut program3 = EbpfLoader::new()
+        .set_global("ORDER", &3, true)
+        .load(crate::TCX)
+        .unwrap();
+
+    let logs0: Arc<Mutex<Vec<CapturedLog>>> = Arc::new(Mutex::new(Vec::new()));
+    setup_logs(&mut program0, &logs0);
+
+    let logs1 = Arc::new(Mutex::new(Vec::new()));
+    setup_logs(&mut program1, &logs1);
+
+    let logs2 = Arc::new(Mutex::new(Vec::new()));
+    setup_logs(&mut program2, &logs2);
+
+    let logs3 = Arc::new(Mutex::new(Vec::new()));
+    setup_logs(&mut program3, &logs3);
+
+    let prog0: &mut SchedClassifier = program0
+        .program_mut("tcx_order")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog0.load().unwrap();
+
+    let prog1: &mut SchedClassifier = program1
+        .program_mut("tcx_order")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog1.load().unwrap();
+
+    let prog2: &mut SchedClassifier = program2
+        .program_mut("tcx_order")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog2.load().unwrap();
+
+    let prog3: &mut SchedClassifier = program3
+        .program_mut("tcx_order")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog3.load().unwrap();
+
+    // Test LinkOrder::first()
+    let options = TcAttachOptions::tcxoptions(LinkOrder::first());
+    prog0
+        .attach_with_options("lo", TcAttachType::Ingress, options)
+        .unwrap();
+
+    // Test LinkOrder::after_program()
+    let order = LinkOrder::after_program(prog0).unwrap();
+    let options = TcAttachOptions::tcxoptions(order);
+    let prog1_link_id = prog1
+        .attach_with_options("lo", TcAttachType::Ingress, options)
+        .unwrap();
+
+    let prog1_link = prog1.take_link(prog1_link_id).unwrap();
+
+    // Test LinkOrder::after_link()
+    let order = LinkOrder::after_link(&prog1_link).unwrap();
+    let options = TcAttachOptions::tcxoptions(order);
+    prog2
+        .attach_with_options("lo", TcAttachType::Ingress, options)
+        .unwrap();
+
+    // Test LinkOrder::last()
+    let options = TcAttachOptions::tcxoptions(LinkOrder::last());
+    prog3
+        .attach_with_options("lo", TcAttachType::Ingress, options)
+        .unwrap();
+
+    const PAYLOAD: &str = "hello tcx";
+
+    let sock = UdpSocket::bind("127.0.0.1:1778").unwrap();
+    let addr = sock.local_addr().unwrap();
+    sock.set_read_timeout(Some(Duration::from_secs(60)))
+        .unwrap();
+    // We only need to send data since we're attaching tx programs to the ingress hook
+    sock.send_to(PAYLOAD.as_bytes(), addr).unwrap();
+
+    // Allow logs to populate
+    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+
+    let log0 = logs0.lock().unwrap();
+    let log1 = logs1.lock().unwrap();
+    let log2 = logs2.lock().unwrap();
+    let log3 = logs3.lock().unwrap();
+
+    debug!("log0: {:?}", log0.first());
+    debug!("log1: {:?}", log1.first());
+    debug!("log2: {:?}", log2.first());
+    debug!("log3: {:?}", log3.first());
+
+    // sort logs by timestamp
+    let mut sorted_logs = [
+        log0.first().unwrap(),
+        log1.first().unwrap(),
+        log2.first().unwrap(),
+        log3.first().unwrap(),
+    ];
+    sorted_logs.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));
+
+    assert!(sorted_logs[0].body.contains("order: 0"));
+    assert!(sorted_logs[1].body.contains("order: 1"));
+    assert!(sorted_logs[2].body.contains("order: 2"));
+    assert!(sorted_logs[3].body.contains("order: 3"));
+}