use std::path::PathBuf;
use anyhow::Result;
#[derive(Debug, Clone)]
pub struct ModuleBConfig {
pub memfd_bytes: u64,
pub socket_path: Option<PathBuf>,
pub marker: String,
}
#[derive(Debug, Clone)]
pub struct ModuleBReceiverConfig {
pub socket_path: PathBuf,
pub marker: String,
pub memfd_bytes: u64,
}
#[derive(Debug, Clone)]
pub struct ModuleBStats {
pub memfd_bytes: u64,
pub handoff_latency_ms: f64,
pub marker_ok: bool,
pub socket_path: PathBuf,
}
impl ModuleBStats {
pub fn to_json(&self) -> String {
format!(
"{{\"module\":\"B\",\"memfd_bytes\":{},\"handoff_latency_ms\":{:.3},\"marker_ok\":{},\"socket_path\":\"{}\"}}",
self.memfd_bytes,
self.handoff_latency_ms,
self.marker_ok,
self.socket_path.display()
)
}
}
#[cfg(target_os = "linux")]
mod linux {
use std::ffi::CString;
use std::fs;
use std::io::{IoSlice, IoSliceMut, Read, Write};
use std::num::NonZeroUsize;
use std::os::fd::{AsFd, AsRawFd, FromRawFd, OwnedFd, RawFd};
use std::os::unix::net::{UnixListener, UnixStream};
use std::path::{Path, PathBuf};
use std::process::Command;
use std::ptr::NonNull;
use std::thread;
use std::time::{Duration, Instant};
use anyhow::{Context, Result};
use nix::cmsg_space;
use nix::sys::memfd::{memfd_create, MFdFlags};
use nix::sys::mman::{mmap, munmap, MapFlags, ProtFlags};
use nix::sys::socket::{recvmsg, sendmsg, ControlMessage, ControlMessageOwned, MsgFlags};
use nix::unistd::ftruncate;
use crate::module_b::{ModuleBConfig, ModuleBReceiverConfig, ModuleBStats};
const ACK_SUCCESS: u8 = 1;
pub fn run_parent(config: ModuleBConfig) -> Result<ModuleBStats> {
let memfd_len = usize::try_from(config.memfd_bytes)
.context("memfd_bytes exceeds usize on this platform")?;
if memfd_len == 0 {
anyhow::bail!("memfd_bytes must be > 0");
}
if config.marker.is_empty() {
anyhow::bail!("marker must not be empty");
}
if config.marker.len() > memfd_len {
anyhow::bail!("marker length must be <= memfd_bytes");
}
let socket_path = config
.socket_path
.unwrap_or_else(|| default_socket_path(std::process::id()));
if socket_path.exists() {
fs::remove_file(&socket_path).with_context(|| {
format!(
"failed to remove pre-existing socket path {}",
socket_path.display()
)
})?;
}
let listener = UnixListener::bind(&socket_path)
.with_context(|| format!("failed to bind UDS listener at {}", socket_path.display()))?;
let memfd_name = CString::new("tracer-bullet-memfd").context("invalid memfd name")?;
let memfd = memfd_create(memfd_name.as_c_str(), MFdFlags::MFD_CLOEXEC)
.context("memfd_create failed")?;
ftruncate(&memfd, config.memfd_bytes as i64).context("ftruncate(memfd) failed")?;
let mut writable_map = MmapGuard::map(&memfd, memfd_len, true)
.context("failed to mmap writable sender memory")?;
write_marker(writable_map.as_mut_slice(), config.marker.as_bytes())
.context("failed to write validation marker into memfd")?;
let exe = std::env::current_exe().context("failed to locate current executable")?;
let mut child = Command::new(exe)
.arg("module-b-receiver-internal")
.arg("--socket-path")
.arg(&socket_path)
.arg("--marker")
.arg(&config.marker)
.arg("--memfd-bytes")
.arg(config.memfd_bytes.to_string())
.spawn()
.context("failed to spawn module-b receiver child process")?;
let (mut stream, _) = listener
.accept()
.context("receiver failed to connect to UDS")?;
let start = Instant::now();
send_fd(stream.as_raw_fd(), memfd.as_raw_fd()).context("sendmsg SCM_RIGHTS failed")?;
let mut ack = [0_u8; 1];
stream
.read_exact(&mut ack)
.context("failed to read receiver ACK")?;
let handoff_latency_ms = start.elapsed().as_secs_f64() * 1_000.0;
let status = child
.wait()
.context("failed waiting for receiver child process")?;
if !status.success() {
anyhow::bail!("receiver child exited unsuccessfully: {status}");
}
drop(writable_map);
if socket_path.exists() {
let _ = fs::remove_file(&socket_path);
}
if ack[0] != ACK_SUCCESS {
anyhow::bail!("receiver returned non-success ACK byte: {}", ack[0]);
}
Ok(ModuleBStats {
memfd_bytes: config.memfd_bytes,
handoff_latency_ms,
marker_ok: true,
socket_path,
})
}
pub fn run_receiver_internal(config: ModuleBReceiverConfig) -> Result<()> {
let memfd_len = usize::try_from(config.memfd_bytes)
.context("memfd_bytes exceeds usize on this platform")?;
if memfd_len == 0 {
anyhow::bail!("memfd_bytes must be > 0");
}
let mut stream = connect_with_retry(&config.socket_path, 200, Duration::from_millis(10))
.with_context(|| {
format!(
"failed to connect to sender socket {}",
config.socket_path.display()
)
})?;
let received_fd = recv_fd(stream.as_raw_fd()).context("recvmsg SCM_RIGHTS failed")?;
let mapped = MmapGuard::map(&received_fd, memfd_len, false)
.context("receiver failed to mmap read-only memfd")?;
validate_marker(mapped.as_slice(), config.marker.as_bytes())
.context("receiver marker validation failed")?;
stream
.write_all(&[ACK_SUCCESS])
.context("failed to write ACK to sender")?;
stream.flush().context("failed to flush ACK to sender")?;
Ok(())
}
fn default_socket_path(pid: u32) -> PathBuf {
PathBuf::from(format!("/tmp/tracer-bullet-{}.sock", pid))
}
fn send_fd(socket_fd: RawFd, fd_to_send: RawFd) -> Result<()> {
let payload = [0xAB_u8];
let iov = [IoSlice::new(&payload)];
let fds = [fd_to_send];
let cmsg = [ControlMessage::ScmRights(&fds)];
let bytes = sendmsg::<()>(socket_fd, &iov, &cmsg, MsgFlags::empty(), None)
.context("sendmsg failed")?;
if bytes != payload.len() {
anyhow::bail!("sendmsg wrote {} bytes instead of {}", bytes, payload.len());
}
Ok(())
}
fn recv_fd(socket_fd: RawFd) -> Result<OwnedFd> {
let mut payload = [0_u8; 1];
let mut iov = [IoSliceMut::new(&mut payload)];
let mut cmsgspace = cmsg_space!([RawFd; 1]);
let msg = recvmsg::<()>(socket_fd, &mut iov, Some(&mut cmsgspace), MsgFlags::empty())
.context("recvmsg failed")?;
if msg.bytes == 0 {
anyhow::bail!("recvmsg received no payload bytes");
}
if msg
.flags
.intersects(MsgFlags::MSG_CTRUNC | MsgFlags::MSG_TRUNC)
{
anyhow::bail!("recvmsg control/payload data was truncated");
}
for cmsg in msg
.cmsgs()
.context("failed to parse recvmsg control messages")?
{
if let ControlMessageOwned::ScmRights(fds) = cmsg {
if let Some(fd) = fds.first().copied() {
let owned = unsafe { OwnedFd::from_raw_fd(fd) };
return Ok(owned);
}
}
}
anyhow::bail!("no SCM_RIGHTS file descriptor received")
}
fn connect_with_retry(path: &Path, attempts: usize, sleep: Duration) -> Result<UnixStream> {
let mut last_err = None;
for _ in 0..attempts {
match UnixStream::connect(path) {
Ok(stream) => return Ok(stream),
Err(err) => {
last_err = Some(err);
thread::sleep(sleep);
}
}
}
match last_err {
Some(err) => {
Err(err).with_context(|| format!("unable to connect to {}", path.display()))
}
None => anyhow::bail!("unable to connect to {}", path.display()),
}
}
fn write_marker(mapped: &mut [u8], marker: &[u8]) -> Result<()> {
if marker.len() > mapped.len() {
anyhow::bail!("marker longer than mapped memory");
}
let start = mapped.len() - marker.len();
mapped[start..].copy_from_slice(marker);
Ok(())
}
fn validate_marker(mapped: &[u8], marker: &[u8]) -> Result<()> {
if marker.len() > mapped.len() {
anyhow::bail!("marker longer than mapped memory");
}
let start = mapped.len() - marker.len();
if mapped[start..] != *marker {
anyhow::bail!("marker mismatch at memfd tail");
}
Ok(())
}
struct MmapGuard {
ptr: NonNull<std::ffi::c_void>,
len: usize,
}
impl MmapGuard {
fn map<Fd: AsFd>(fd: &Fd, len: usize, writable: bool) -> Result<Self> {
let nz_len = NonZeroUsize::new(len).context("mmap length must be > 0")?;
let prot = if writable {
ProtFlags::PROT_READ | ProtFlags::PROT_WRITE
} else {
ProtFlags::PROT_READ
};
let ptr = unsafe { mmap(None, nz_len, prot, MapFlags::MAP_SHARED, fd, 0) }
.context("mmap syscall failed")?;
Ok(Self { ptr, len })
}
fn as_slice(&self) -> &[u8] {
unsafe { std::slice::from_raw_parts(self.ptr.as_ptr() as *const u8, self.len) }
}
fn as_mut_slice(&mut self) -> &mut [u8] {
unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr() as *mut u8, self.len) }
}
}
impl Drop for MmapGuard {
fn drop(&mut self) {
let _ = unsafe { munmap(self.ptr, self.len) };
}
}
#[cfg(test)]
mod tests {
use std::io::{Read, Write};
use std::os::fd::AsRawFd;
use std::os::unix::net::UnixStream;
use std::thread;
use super::*;
#[test]
fn transfers_fd_and_validates_marker() {
let memfd_name = CString::new("module-b-test").expect("valid test memfd name");
let memfd = memfd_create(memfd_name.as_c_str(), MFdFlags::MFD_CLOEXEC)
.expect("memfd_create should succeed in test");
let len = 8 * 1024 * 1024;
let marker = b"TRACER_BULLET_SUCCESS";
ftruncate(&memfd, len as i64).expect("ftruncate should succeed in test");
let mut map = MmapGuard::map(&memfd, len, true).expect("mmap should succeed in test");
write_marker(map.as_mut_slice(), marker).expect("writing marker should succeed");
let (mut send_stream, recv_stream) =
UnixStream::pair().expect("socketpair should succeed");
let marker_vec = marker.to_vec();
let join = thread::spawn(move || {
let fd =
recv_fd(recv_stream.as_raw_fd()).expect("recv_fd should succeed in thread");
let mapped = MmapGuard::map(&fd, len, false).expect("receiver mmap should succeed");
validate_marker(mapped.as_slice(), &marker_vec)
.expect("marker should validate in receiver");
let mut stream = recv_stream;
stream
.write_all(&[ACK_SUCCESS])
.expect("thread ACK write should succeed");
});
send_fd(send_stream.as_raw_fd(), memfd.as_raw_fd()).expect("send_fd should succeed");
let mut ack = [0_u8; 1];
send_stream
.read_exact(&mut ack)
.expect("ACK read should succeed in sender");
assert_eq!(ack, [ACK_SUCCESS]);
join.join()
.expect("receiver thread should join successfully");
}
#[test]
#[ignore = "heavy memory mapping validation"]
fn validates_full_5gb_mapping_without_copy() {
let memfd_name = CString::new("module-b-test-5gb").expect("valid test memfd name");
let memfd = memfd_create(memfd_name.as_c_str(), MFdFlags::MFD_CLOEXEC)
.expect("memfd_create should succeed in test");
let len = 5 * 1024 * 1024 * 1024usize;
let marker = b"TRACER_BULLET_SUCCESS";
ftruncate(&memfd, len as i64).expect("ftruncate should succeed in test");
let mut map = MmapGuard::map(&memfd, len, true).expect("mmap should succeed in test");
write_marker(map.as_mut_slice(), marker).expect("write marker should succeed");
validate_marker(map.as_slice(), marker).expect("validate marker should succeed");
}
}
}
#[cfg(target_os = "linux")]
pub fn run_parent(config: ModuleBConfig) -> Result<ModuleBStats> {
linux::run_parent(config)
}
#[cfg(target_os = "linux")]
pub fn run_receiver_internal(config: ModuleBReceiverConfig) -> Result<()> {
linux::run_receiver_internal(config)
}
#[cfg(not(target_os = "linux"))]
pub fn run_parent(_config: ModuleBConfig) -> Result<ModuleBStats> {
anyhow::bail!("module-b requires Linux for memfd + SCM_RIGHTS")
}
#[cfg(not(target_os = "linux"))]
pub fn run_receiver_internal(_config: ModuleBReceiverConfig) -> Result<()> {
anyhow::bail!("module-b receiver requires Linux")
}