Adjust tcp_ca example to use struct_ops shadow data

This change adjusts the tcp_ca example to include the usage of
struct_ops shadow data. Specifically, we modify the name under which we
register the algorithm before attaching the program. We then make sure
that we can use the algorithm under the new name. In addition, we also
illustrate setting of a callback/program before load.

Signed-off-by: Daniel Müller <[email protected]>

Author: d-e-s-o

examples/tcp_ca/src/bpf/tcp_ca.bpf.c

@@ -8,6 +8,7 @@
 char _license[] SEC("license") = "GPL";
 
 int ca_cnt = 0;
+bool cong_control = false;
 
 static inline struct tcp_sock *tcp_sk(const struct sock *sk)
 {
@@ -20,12 +21,24 @@ void BPF_PROG(ca_update_init, struct sock *sk)
 	ca_cnt++;
 }
 
+SEC("struct_ops/ca_update_2_init")
+void BPF_PROG(ca_update_2_init, struct sock *sk)
+{
+}
+
 SEC("struct_ops/ca_update_cong_control")
 void BPF_PROG(ca_update_cong_control, struct sock *sk,
 	      const struct rate_sample *rs)
 {
 }
 
+SEC("struct_ops/ca_update_cong_control2")
+void BPF_PROG(ca_update_cong_control2, struct sock *sk,
+	      const struct rate_sample *rs)
+{
+  cong_control = true;
+}
+
 SEC("struct_ops/ca_update_ssthresh")
 __u32 BPF_PROG(ca_update_ssthresh, struct sock *sk)
 {
@@ -38,11 +51,25 @@ __u32 BPF_PROG(ca_update_undo_cwnd, struct sock *sk)
 	return tcp_sk(sk)->snd_cwnd;
 }
 
-SEC(".struct_ops")
+SEC(".struct_ops.link")
 struct tcp_congestion_ops ca_update = {
 	.init = (void *)ca_update_init,
 	.cong_control = (void *)ca_update_cong_control,
 	.ssthresh = (void *)ca_update_ssthresh,
 	.undo_cwnd = (void *)ca_update_undo_cwnd,
 	.name = "tcp_ca_update",
 };
+
+SEC(".struct_ops")
+struct tcp_congestion_ops ca_update_2 = {
+	.init = (void *)ca_update_2_init,
+	.cong_control = (void *)ca_update_cong_control2,
+	.ssthresh = (void *)ca_update_ssthresh,
+	.undo_cwnd = (void *)ca_update_undo_cwnd,
+	.name = "tcp_ca_update_2",
+};
+
+SEC(".struct_ops.link")
+struct tcp_congestion_ops ca_wrong = {
+	.cong_control = (void *)ca_update_cong_control,
+};

examples/tcp_ca/src/main.rs

@@ -4,6 +4,7 @@
 
 use std::ffi::c_int;
 use std::ffi::c_void;
+use std::ffi::CStr;
 use std::io;
 use std::io::Read as _;
 use std::io::Write as _;
@@ -12,6 +13,7 @@ use std::net::TcpStream;
 use std::os::fd::AsFd as _;
 use std::os::fd::AsRawFd as _;
 use std::os::fd::BorrowedFd;
+use std::ptr::copy_nonoverlapping;
 use std::thread;
 
 use clap::Parser;
@@ -22,13 +24,18 @@ use libc::TCP_CONGESTION;
 
 use libbpf_rs::skel::OpenSkel;
 use libbpf_rs::skel::SkelBuilder;
+use libbpf_rs::AsRawLibbpf as _;
 use libbpf_rs::ErrorExt as _;
+use libbpf_rs::ErrorKind;
 use libbpf_rs::Result;
 
 use crate::tcp_ca::TcpCaSkelBuilder;
 
 mod tcp_ca {
-    include!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/bpf/tcp_ca.skel.rs"));
+    include!(concat!(
+        env!("CARGO_MANIFEST_DIR"),
+        "/src/bpf/tcp_ca.skel.rs"
+    ));
 }
 
 const TCP_CA_UPDATE: &[u8] = b"tcp_ca_update\0";
@@ -58,24 +65,29 @@ fn set_sock_opt(
 
 /// Set the `tcp_ca_update` congestion algorithm on the socket represented by
 /// the provided file descriptor.
-fn set_tcp_ca(fd: BorrowedFd<'_>) -> Result<()> {
+fn set_tcp_ca(fd: BorrowedFd<'_>, tcp_ca: &CStr) -> Result<()> {
     let () = set_sock_opt(
         fd,
         IPPROTO_TCP,
         TCP_CONGESTION,
-        TCP_CA_UPDATE.as_ptr().cast(),
-        (TCP_CA_UPDATE.len() - 1) as _,
+        tcp_ca.as_ptr().cast(),
+        tcp_ca.to_bytes().len() as _,
     )
-    .context("failed to set TCP_CONGESTION")?;
+    .with_context(|| {
+        format!(
+            "failed to set TCP_CONGESTION algorithm `{}`",
+            tcp_ca.to_str().unwrap()
+        )
+    })?;
     Ok(())
 }
 
 /// Send and receive a bunch of data over TCP sockets using the `tcp_ca_update`
 /// congestion algorithm.
-fn send_recv() -> Result<()> {
+fn send_recv(tcp_ca: &CStr) -> Result<()> {
     let num_bytes = 8 * 1024 * 1024;
     let listener = TcpListener::bind("[::1]:0")?;
-    let () = set_tcp_ca(listener.as_fd())?;
+    let () = set_tcp_ca(listener.as_fd(), tcp_ca)?;
     let addr = listener.local_addr()?;
 
     let send_handle = thread::spawn(move || {
@@ -86,38 +98,96 @@ fn send_recv() -> Result<()> {
 
     let mut received = Vec::new();
     let mut stream = TcpStream::connect(addr)?;
-    let () = set_tcp_ca(stream.as_fd())?;
+    let () = set_tcp_ca(stream.as_fd(), tcp_ca)?;
     let _count = stream.read_to_end(&mut received)?;
     let () = send_handle.join().unwrap();
 
     assert_eq!(received.len(), num_bytes);
     Ok(())
 }
 
-fn main() -> Result<()> {
-    let args = Args::parse();
-
+fn test(name_to_register: Option<&CStr>, name_to_use: &CStr, verbose: bool) -> Result<()> {
     let mut skel_builder = TcpCaSkelBuilder::default();
-    if args.verbose {
+    if verbose {
         skel_builder.obj_builder.debug(true);
     }
 
-    let open_skel = skel_builder.open()?;
+    let mut open_skel = skel_builder.open()?;
+
+    if let Some(name) = name_to_register {
+        // Here we illustrate the possibility of updating `struct_ops` data before
+        // load. That can be used to communicate data to the kernel, e.g., for
+        // initialization purposes.
+        let ca_update = open_skel.struct_ops.ca_update_mut();
+        if name.to_bytes_with_nul().len() > ca_update.name.len() {
+            panic!(
+                "TCP CA name `{}` exceeds maximum length {}",
+                name.to_str().unwrap(),
+                ca_update.name.len()
+            );
+        }
+        let len = name.to_bytes_with_nul().len();
+        let () = unsafe { copy_nonoverlapping(name.as_ptr(), ca_update.name.as_mut_ptr(), len) };
+        let () = ca_update.name[len..].fill(0);
+    }
+
+    let ca_update_cong_control2 = open_skel
+        .progs()
+        .ca_update_cong_control2()
+        .as_libbpf_object()
+        .as_ptr();
+    let ca_update = open_skel.struct_ops.ca_update_mut();
+    ca_update.cong_control = ca_update_cong_control2;
+
     let mut skel = open_skel.load()?;
     let mut maps = skel.maps_mut();
     let map = maps.ca_update();
     let _link = map.attach_struct_ops()?;
 
-    println!("Registered `tcp_ca_update` congestion algorithm; using it for loopback based data exchange...");
+    println!(
+        "Registered `{}` congestion algorithm; using `{}` for loopback based data exchange...",
+        name_to_register.unwrap_or(name_to_use).to_str().unwrap(),
+        name_to_use.to_str().unwrap()
+    );
+
+    // NB: At this point `/proc/sys/net/ipv4/tcp_available_congestion_control`
+    //     would list the registered congestion algorithm.
 
     assert_eq!(skel.bss().ca_cnt, 0);
+    assert!(!skel.bss().cong_control);
 
     // Use our registered TCP congestion algorithm while sending a bunch of data
     // over the loopback device.
-    let () = send_recv()?;
+    let () = send_recv(name_to_use)?;
     println!("Done.");
 
-    let saved_ca1_cnt = skel.bss().ca_cnt;
-    assert_ne!(saved_ca1_cnt, 0);
+    let saved_ca_cnt = skel.bss().ca_cnt;
+    assert_ne!(saved_ca_cnt, 0);
+    // With `ca_update_cong_control2` active, we should have seen the
+    // `cong_control` value changed as well.
+    assert!(skel.bss().cong_control);
+    Ok(())
+}
+
+fn main() -> Result<()> {
+    let args = Args::parse();
+
+    let tcp_ca = CStr::from_bytes_until_nul(TCP_CA_UPDATE).unwrap();
+    let () = test(None, tcp_ca, args.verbose)?;
+
+    // Use a different name under which the algorithm is registered; just for
+    // illustration purposes of how to change `struct_ops` related data before
+    // load/attachment.
+    let new_ca = CStr::from_bytes_until_nul(b"anotherca\0").unwrap();
+    let () = test(Some(new_ca), new_ca, args.verbose)?;
+
+    // Just to be sure we are not bullshitting with the above, use a different
+    // congestion algorithm than what we register. This is expected to fail,
+    // because said algorithm to use cannot be found.
+    let to_register = CStr::from_bytes_until_nul(b"holycowca\0").unwrap();
+    let err = test(Some(to_register), tcp_ca, args.verbose).unwrap_err();
+    assert_eq!(err.kind(), ErrorKind::NotFound);
+    println!("Expected failure: {err:#}");
+
     Ok(())
 }