Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add sockaddr
Browse files Browse the repository at this point in the history
adgaultier committed Nov 8, 2024
1 parent 570b1d1 commit 8c2686a
Showing 16 changed files with 1,284 additions and 0 deletions.
9 changes: 9 additions & 0 deletions test-sock-addr/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
### https://raw.github.com/github/gitignore/master/Rust.gitignore

# Generated by Cargo
# will have compiled files and executables
debug/
target/

# These are backup files generated by rustfmt
**/*.rs.bk
813 changes: 813 additions & 0 deletions test-sock-addr/Cargo.lock

Large diffs are not rendered by default.

32 changes: 32 additions & 0 deletions test-sock-addr/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
[workspace]
resolver = "2"
members = ["test-sockaddr", "test-sockaddr-common", "test-sockaddr-ebpf"]
default-members = ["test-sockaddr", "test-sockaddr-common"]

[workspace.dependencies]
aya = { version = "0.13.0", default-features = false }
aya-ebpf = { version = "0.1.1", default-features = false }
aya-log = { version = "0.2.1", default-features = false }
aya-log-ebpf = { version = "0.1.1", default-features = false }

anyhow = { version = "1", default-features = false }
cargo_metadata = { version = "0.18.0", default-features = false }
# `std` feature is currently required to build `clap`.
#
# See https://github.com/clap-rs/clap/blob/61f5ee5/clap_builder/src/lib.rs#L15.
clap = { version = "4.5.20", default-features = false, features = ["std"] }
env_logger = { version = "0.11.5", default-features = false }
libc = { version = "0.2.159", default-features = false }
log = { version = "0.4.22", default-features = false }
tokio = { version = "1.40.0", default-features = false }
which = { version = "6.0.0", default-features = false }

[profile.dev]
panic = "abort"

[profile.release]
panic = "abort"

[profile.release.package.test-sockaddr-ebpf]
debug = 2
codegen-units = 1
33 changes: 33 additions & 0 deletions test-sock-addr/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
# test-sockaddr

## Prerequisites

1. stable rust toolchains: `rustup toolchain install stable`
1. nightly rust toolchains: `rustup toolchain install nightly --component rust-src`
1. (if cross-compiling) rustup target: `rustup target add ${ARCH}-unknown-linux-musl`
1. (if cross-compiling) LLVM: (e.g.) `brew install llvm` (on macOS)
1. (if cross-compiling) C toolchain: (e.g.) [`brew install filosottile/musl-cross/musl-cross`](https://github.com/FiloSottile/homebrew-musl-cross) (on macOS)
1. bpf-linker: `cargo install bpf-linker` (`--no-default-features` on macOS)

## Build & Run

Use `cargo build`, `cargo check`, etc. as normal. Run your program with:

```shell
cargo run --release --config 'target."cfg(all())".runner="sudo -E"'
```

Cargo build scripts are used to automatically build the eBPF correctly and include it in the
program.

## Cross-compiling on macOS

Cross compilation should work on both Intel and Apple Silicon Macs.

```shell
CC=${ARCH}-linux-musl-gcc cargo build --package test-sockaddr --release \
--target=${ARCH}-unknown-linux-musl \
--config=target.${ARCH}-unknown-linux-musl.linker=\"${ARCH}-linux-musl-gcc\"
```
The cross-compiled program `target/${ARCH}-unknown-linux-musl/release/test-sockaddr` can be
copied to a Linux server or VM and run there.
4 changes: 4 additions & 0 deletions test-sock-addr/rustfmt.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
group_imports = "StdExternalCrate"
imports_granularity = "Crate"
reorder_imports = true
unstable_features = true
14 changes: 14 additions & 0 deletions test-sock-addr/test-sockaddr-common/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
[package]
name = "test-sockaddr-common"
version = "0.1.0"
edition = "2021"

[features]
default = []
user = ["aya"]

[dependencies]
aya = { workspace = true, optional = true }

[lib]
path = "src/lib.rs"
1 change: 1 addition & 0 deletions test-sock-addr/test-sockaddr-common/src/lib.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
#![no_std]
12 changes: 12 additions & 0 deletions test-sock-addr/test-sockaddr-ebpf/.cargo/config.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
# We have this so that one doesn't need to manually pass
# --target=bpfel-unknown-none -Z build-std=core when running cargo
# check/build/doc etc.
#
# NB: this file gets loaded only if you run cargo from this directory, it's
# ignored if you run from the workspace root. See
# https://doc.rust-lang.org/cargo/reference/config.html#hierarchical-structure
[build]
target = ["bpfeb-unknown-none", "bpfel-unknown-none"]

[unstable]
build-std = ["core"]
17 changes: 17 additions & 0 deletions test-sock-addr/test-sockaddr-ebpf/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
[package]
name = "test-sockaddr-ebpf"
version = "0.1.0"
edition = "2021"

[dependencies]
test-sockaddr-common = { path = "../test-sockaddr-common" }

aya-ebpf = { workspace = true }
aya-log-ebpf = { workspace = true }

[build-dependencies]
which = { workspace = true }

[[bin]]
name = "test-sockaddr"
path = "src/main.rs"
17 changes: 17 additions & 0 deletions test-sock-addr/test-sockaddr-ebpf/build.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
use which::which;

/// Building this crate has an undeclared dependency on the `bpf-linker` binary. This would be
/// better expressed by [artifact-dependencies][bindeps] but issues such as
/// https://github.com/rust-lang/cargo/issues/12385 make their use impractical for the time being.
///
/// This file implements an imperfect solution: it causes cargo to rebuild the crate whenever the
/// mtime of `which bpf-linker` changes. Note that possibility that a new bpf-linker is added to
/// $PATH ahead of the one used as the cache key still exists. Solving this in the general case
/// would require rebuild-if-changed-env=PATH *and* rebuild-if-changed={every-directory-in-PATH}
/// which would likely mean far too much cache invalidation.
///
/// [bindeps]: https://doc.rust-lang.org/nightly/cargo/reference/unstable.html?highlight=feature#artifact-dependencies
fn main() {
let bpf_linker = which("bpf-linker").unwrap();
println!("cargo:rerun-if-changed={}", bpf_linker.to_str().unwrap());
}
3 changes: 3 additions & 0 deletions test-sock-addr/test-sockaddr-ebpf/rust-toolchain.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
[toolchain]
channel = "nightly"
components = ["rust-src"]
3 changes: 3 additions & 0 deletions test-sock-addr/test-sockaddr-ebpf/src/lib.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
#![no_std]

// This file exists to enable the library target.
81 changes: 81 additions & 0 deletions test-sock-addr/test-sockaddr-ebpf/src/main.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
#![no_std]
#![no_main]

use core::net::Ipv4Addr;

use aya_ebpf::{
bindings::{bpf_sock, bpf_sock_addr},
helpers::bpf_get_current_pid_tgid,
macros::{cgroup_sock, cgroup_sock_addr},
programs::{SockAddrContext, SockContext},
};
use aya_log_ebpf::info;

#[cgroup_sock_addr(connect4)]
pub fn socket_connect(ctx: SockAddrContext) -> i32 {
info!(&ctx, "connect");
match sock_connect(ctx) {
Ok(ret) => ret,
Err(ret) => ret,
}
}

#[cgroup_sock(post_bind4)]
pub fn socket_create(ctx: SockContext) -> i32 {
info!(&ctx, "create");
match sock_create(ctx) {
Ok(ret) => ret,
Err(ret) => ret,
}
}
fn sock_create(ctx: SockContext) -> Result<i32, i32> {
let sock = unsafe { *(ctx.sock as *const bpf_sock) };
let proto = sock.protocol;
let family = sock.family;
let type_ = sock.type_;

info!(
&ctx,
"create sock with proto :{} fam {} type {} ", proto, family, type_,
);
Ok(1)
}

// pub struct bpf_sock_addr {
// pub user_family: __u32,
// pub user_ip4: __u32,
// pub user_ip6: [__u32; 4usize],
// pub user_port: __u32,
// pub family: __u32,
// pub type_: __u32,
// pub protocol: __u32,
// pub msg_src_ip4: __u32,
// pub msg_src_ip6: [__u32; 4usize],
// pub __bindgen_anon_1: bpf_sock_addr__bindgen_ty_1,
// }
fn sock_connect(ctx: SockAddrContext) -> Result<i32, i32> {
let pid = (bpf_get_current_pid_tgid() >> 32) as u32;
let sock_addr = unsafe { *(ctx.sock_addr as *const bpf_sock_addr) };

let family = sock_addr.family;
let proto = sock_addr.protocol as u8;
let ipv4 = sock_addr.user_ip4;

let port = sock_addr.user_port;
info!(
&ctx,
"pid:{} family:{} proto:{} ip:{} port:{} ",
pid,
family,
proto,
Ipv4Addr::from_bits(u32::from_be(ipv4)),
u16::from_be(port as u16)
);
Ok(1)
}

#[cfg(not(test))]
#[panic_handler]
fn panic(_info: &core::panic::PanicInfo) -> ! {
loop {}
}
36 changes: 36 additions & 0 deletions test-sock-addr/test-sockaddr/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
[package]
name = "test-sockaddr"
version = "0.1.0"
edition = "2021"

[dependencies]
test-sockaddr-common = { path = "../test-sockaddr-common", features = ["user"] }

anyhow = { workspace = true, default-features = true }
aya = { workspace = true }
aya-log = { workspace = true }
env_logger = { workspace = true }
libc = { workspace = true }
log = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread", "net", "signal"] }

clap = { workspace = true, features = ["derive"] }
[build-dependencies]
cargo_metadata = { workspace = true }
# TODO(https://github.com/rust-lang/cargo/issues/12375): this should be an artifact dependency, but
# it's not possible to tell cargo to use `-Z build-std` to build it. We cargo-in-cargo in the build
# script to build this, but we want to teach cargo about the dependecy so that cache invalidation
# works properly.
#
# Note also that https://github.com/rust-lang/cargo/issues/10593 occurs when `target = ...` is added
# to an artifact dependency; it seems possible to work around that by setting `resolver = "1"` in
# Cargo.toml in the workspace root.
#
# Finally note that *any* usage of `artifact = ...` in *any* Cargo.toml in the workspace breaks
# workflows with stable cargo; stable cargo outright refuses to load manifests that use unstable
# features.
test-sockaddr-ebpf = { path = "../test-sockaddr-ebpf" }

[[bin]]
name = "test-sockaddr"
path = "src/main.rs"
150 changes: 150 additions & 0 deletions test-sock-addr/test-sockaddr/build.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,150 @@
use std::{
env, fs,
io::{BufRead as _, BufReader},
path::PathBuf,
process::{Child, Command, Stdio},
};

use cargo_metadata::{
Artifact, CompilerMessage, Message, Metadata, MetadataCommand, Package, Target,
};

/// This crate has a runtime dependency on artifacts produced by the `test-sockaddr-ebpf` crate.
/// This would be better expressed as one or more [artifact-dependencies][bindeps] but issues such
/// as:
///
/// * https://github.com/rust-lang/cargo/issues/12374
/// * https://github.com/rust-lang/cargo/issues/12375
/// * https://github.com/rust-lang/cargo/issues/12385
///
/// prevent their use for the time being.
///
/// [bindeps]: https://doc.rust-lang.org/nightly/cargo/reference/unstable.html?highlight=feature#artifact-dependencies
fn main() {
let Metadata { packages, .. } = MetadataCommand::new().no_deps().exec().unwrap();
let ebpf_package = packages
.into_iter()
.find(|Package { name, .. }| name == "test-sockaddr-ebpf")
.unwrap();

let out_dir = env::var_os("OUT_DIR").unwrap();
let out_dir = PathBuf::from(out_dir);

let endian = env::var_os("CARGO_CFG_TARGET_ENDIAN").unwrap();
let target = if endian == "big" {
"bpfeb"
} else if endian == "little" {
"bpfel"
} else {
panic!("unsupported endian={:?}", endian)
};

// TODO(https://github.com/rust-lang/cargo/issues/4001): Make this `false` if we can determine
// we're in a check build.
let build_ebpf = true;
if build_ebpf {
let arch = env::var_os("CARGO_CFG_TARGET_ARCH").unwrap();

let target = format!("{target}-unknown-none");

let Package { manifest_path, .. } = ebpf_package;
let ebpf_dir = manifest_path.parent().unwrap();

// We have a build-dependency on `test-sockaddr-ebpf`, so cargo will automatically rebuild us
// if `test-sockaddr-ebpf`'s *library* target or any of its dependencies change. Since we
// depend on `test-sockaddr-ebpf`'s *binary* targets, that only gets us half of the way. This
// stanza ensures cargo will rebuild us on changes to the binaries too, which gets us the
// rest of the way.
println!("cargo:rerun-if-changed={}", ebpf_dir.as_str());

let mut cmd = Command::new("cargo");
cmd.args([
"build",
"-Z",
"build-std=core",
"--bins",
"--message-format=json",
"--release",
"--target",
&target,
]);

cmd.env("CARGO_CFG_BPF_TARGET_ARCH", arch);

// Workaround to make sure that the rust-toolchain.toml is respected.
for key in ["RUSTUP_TOOLCHAIN", "RUSTC", "RUSTC_WORKSPACE_WRAPPER"] {
cmd.env_remove(key);
}
cmd.current_dir(ebpf_dir);

// Workaround for https://github.com/rust-lang/cargo/issues/6412 where cargo flocks itself.
let ebpf_target_dir = out_dir.join("test-sockaddr-ebpf");
cmd.arg("--target-dir").arg(&ebpf_target_dir);

let mut child = cmd
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.unwrap_or_else(|err| panic!("failed to spawn {cmd:?}: {err}"));
let Child { stdout, stderr, .. } = &mut child;

// Trampoline stdout to cargo warnings.
let stderr = stderr.take().unwrap();
let stderr = BufReader::new(stderr);
let stderr = std::thread::spawn(move || {
for line in stderr.lines() {
let line = line.unwrap();
println!("cargo:warning={line}");
}
});

let stdout = stdout.take().unwrap();
let stdout = BufReader::new(stdout);
let mut executables = Vec::new();
for message in Message::parse_stream(stdout) {
#[allow(clippy::collapsible_match)]
match message.expect("valid JSON") {
Message::CompilerArtifact(Artifact {
executable,
target: Target { name, .. },
..
}) => {
if let Some(executable) = executable {
executables.push((name, executable.into_std_path_buf()));
}
}
Message::CompilerMessage(CompilerMessage { message, .. }) => {
for line in message.rendered.unwrap_or_default().split('\n') {
println!("cargo:warning={line}");
}
}
Message::TextLine(line) => {
println!("cargo:warning={line}");
}
_ => {}
}
}

let status = child
.wait()
.unwrap_or_else(|err| panic!("failed to wait for {cmd:?}: {err}"));
assert_eq!(status.code(), Some(0), "{cmd:?} failed: {status:?}");

stderr.join().map_err(std::panic::resume_unwind).unwrap();

for (name, binary) in executables {
let dst = out_dir.join(name);
let _: u64 = fs::copy(&binary, &dst)
.unwrap_or_else(|err| panic!("failed to copy {binary:?} to {dst:?}: {err}"));
}
} else {
let Package { targets, .. } = ebpf_package;
for Target { name, kind, .. } in targets {
if *kind != ["bin"] {
continue;
}
let dst = out_dir.join(name);
fs::write(&dst, []).unwrap_or_else(|err| panic!("failed to create {dst:?}: {err}"));
}
}
}
59 changes: 59 additions & 0 deletions test-sock-addr/test-sockaddr/src/main.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
use std::fs::File;

use aya::programs::{links::CgroupAttachMode, CgroupSock, CgroupSockAddr};
use clap::Parser;
#[rustfmt::skip]
use log::{debug, warn};

use tokio::signal;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
env_logger::init();

// Bump the memlock rlimit. This is needed for older kernels that don't use the
// new memcg based accounting, see https://lwn.net/Articles/837122/
let rlim = libc::rlimit {
rlim_cur: libc::RLIM_INFINITY,
rlim_max: libc::RLIM_INFINITY,
};
let ret = unsafe { libc::setrlimit(libc::RLIMIT_MEMLOCK, &rlim) };
if ret != 0 {
debug!("remove limit on locked memory failed, ret is: {}", ret);
}

let mut ebpf = aya::Ebpf::load(aya::include_bytes_aligned!(concat!(
env!("OUT_DIR"),
"/test-sockaddr"
)))?;
if let Err(e) = aya_log::EbpfLogger::init(&mut ebpf) {
// This can happen if you remove all log statements from your eBPF program.
warn!("failed to initialize eBPF logger: {}", e);
}
let file = File::open("/sys/fs/cgroup/user.slice")?;
let sock_create: &mut CgroupSock = ebpf.program_mut("socket_create").unwrap().try_into()?;
sock_create.load()?;
sock_create.attach(file, CgroupAttachMode::Single)?;

let mut ebpf = aya::Ebpf::load(aya::include_bytes_aligned!(concat!(
env!("OUT_DIR"),
"/test-sockaddr"
)))?;
if let Err(e) = aya_log::EbpfLogger::init(&mut ebpf) {
// This can happen if you remove all log statements from your eBPF program.
warn!("failed to initialize eBPF logger: {}", e);
}
let sock_connect: &mut CgroupSockAddr =
ebpf.program_mut("socket_connect").unwrap().try_into()?;
sock_connect.load()?;
let file = File::open("/sys/fs/cgroup/user.slice")?;

sock_connect.attach(file, CgroupAttachMode::Single)?;

let ctrl_c = signal::ctrl_c();
println!("Waiting for Ctrl-C...");
ctrl_c.await?;
println!("Exiting...");

Ok(())
}

0 comments on commit 8c2686a

Please sign in to comment.