use std::ffi::{OsStr, OsString};
use std::io::{self, Read};
use std::os::fd::AsRawFd;
use std::os::unix::net::UnixStream;
use std::os::unix::process::CommandExt;
use std::path::{Path, PathBuf};
use std::time::Duration;
pub const OWNER_SOCKET_ENV: &str = "VTCODE_COMPANION_OWNER_SOCKET";
pub const HELPER_MODE: &str = "--vtcode-owned-exec";
pub const OWNER_DIRECTORY_PREFIX: &str = "vtc-";
pub const OWNER_SOCKET_FILE: &str = "owner";
const READY: u8 = 1;
pub struct OwnedCommand {
pub executable: PathBuf,
pub arguments: Vec<OsString>,
pub owner_socket: PathBuf,
}
impl OwnedCommand {
pub fn new(
executable: &Path,
owner_socket: &Path,
program: &OsStr,
args: impl IntoIterator<Item = impl AsRef<OsStr>>,
arg0: Option<&OsStr>,
) -> Self {
let mut arguments = vec![
OsString::from(HELPER_MODE),
arg0.unwrap_or(program).to_owned(),
program.to_owned(),
];
arguments.extend(args.into_iter().map(|arg| arg.as_ref().to_owned()));
Self {
executable: executable.to_owned(),
arguments,
owner_socket: owner_socket.to_owned(),
}
}
pub fn from_environment(
program: &OsStr,
args: impl IntoIterator<Item = impl AsRef<OsStr>>,
arg0: Option<&OsStr>,
) -> io::Result<Option<Self>> {
let Some(socket) = std::env::var_os(OWNER_SOCKET_ENV) else {
return Ok(None);
};
let socket = PathBuf::from(socket);
if !is_owned_endpoint(&socket) {
return Ok(None);
}
Ok(Some(Self::new(&std::env::current_exe()?, &socket, program, args, arg0)))
}
}
pub fn is_owned_endpoint(socket: &Path) -> bool {
is_owned_endpoint_under(socket, &std::env::temp_dir())
}
fn is_owned_endpoint_under(socket: &Path, temp_root: &Path) -> bool {
use std::os::unix::fs::{FileTypeExt, MetadataExt, PermissionsExt};
if !socket.is_absolute() || socket.file_name() != Some(OsStr::new(OWNER_SOCKET_FILE)) {
return false;
}
let Ok(endpoint) = std::fs::symlink_metadata(socket) else {
return false;
};
if !endpoint.file_type().is_socket() {
return false;
}
let Some(directory) = socket.parent() else {
return false;
};
let Ok(directory_metadata) = std::fs::metadata(directory) else {
return false;
};
let owner_only = directory_metadata.permissions().mode() & 0o077 == 0;
let owned_by_self = directory_metadata.is_dir() && directory_metadata.uid() == nix::unistd::geteuid().as_raw();
if !owner_only || !owned_by_self {
return false;
}
if !directory
.file_name()
.and_then(OsStr::to_str)
.is_some_and(|name| name.starts_with(OWNER_DIRECTORY_PREFIX))
{
return false;
}
let Ok(directory) = vtcode_commons::canonicalize(directory) else {
return false;
};
let temp_root = vtcode_commons::canonicalize(temp_root).unwrap_or_else(|_| temp_root.to_path_buf());
directory.starts_with(temp_root)
}
pub fn maybe_run_owned_command() -> io::Result<bool> {
let args: Vec<_> = std::env::args_os().skip(1).collect();
if args.first().is_none_or(|arg| arg != HELPER_MODE) {
return Ok(false);
}
let arg0 = args
.get(1)
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "missing target argv[0]"))?;
let program = args
.get(2)
.filter(|p| !p.is_empty())
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "missing target program"))?;
if nix::unistd::getpgrp() != nix::unistd::getpid() {
return Err(io::Error::new(io::ErrorKind::InvalidInput, "owned command requires an isolated process group"));
}
let socket = std::env::var_os(OWNER_SOCKET_ENV)
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "missing task owner endpoint"))?;
let mut owner = UnixStream::connect(socket)?;
owner.set_read_timeout(Some(Duration::from_secs(2)))?;
let mut acknowledgement = [0];
owner.read_exact(&mut acknowledgement)?;
if acknowledgement != [READY] {
return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid task owner acknowledgement"));
}
owner.set_read_timeout(None)?;
spawn_watchdog(&owner)?;
drop(owner);
let error = std::process::Command::new(program).arg0(arg0).args(args.iter().skip(3)).exec();
Err(error)
}
#[expect(
unsafe_code,
reason = "Fork children execute only async-signal-safe syscalls and exit without Rust destructors"
)]
fn spawn_watchdog(owner: &UnixStream) -> io::Result<()> {
let (mut ready_reader, ready_writer) = UnixStream::pair()?;
ready_reader.set_read_timeout(Some(Duration::from_secs(2)))?;
let owner_fd = owner.as_raw_fd();
let ready_fd = ready_writer.as_raw_fd();
if owner_fd <= 2 || ready_fd <= 2 || ready_reader.as_raw_fd() <= 2 {
return Err(io::Error::new(io::ErrorKind::InvalidInput, "owned command requires open standard streams"));
}
let group = nix::unistd::getpgrp().as_raw();
let intermediate = unsafe { libc::fork() };
if intermediate == -1 {
return Err(io::Error::last_os_error());
}
if intermediate == 0 {
let monitor = unsafe { libc::fork() };
if monitor == -1 {
unsafe { libc::_exit(1) };
}
if monitor != 0 {
unsafe { libc::_exit(0) };
}
unsafe { libc::close(ready_reader.as_raw_fd()) };
unsafe { run_watchdog(owner_fd, ready_fd, group) };
}
drop(ready_writer);
let mut status = 0;
loop {
let waited = unsafe { libc::waitpid(intermediate, &mut status, 0) };
if waited == intermediate {
break;
}
let error = io::Error::last_os_error();
if error.kind() != io::ErrorKind::Interrupted {
return Err(error);
}
}
if !libc::WIFEXITED(status) || libc::WEXITSTATUS(status) != 0 {
return Err(io::Error::other("failed to start command lifetime watchdog"));
}
let mut ready = [0];
ready_reader.read_exact(&mut ready)?;
if ready != [READY] {
return Err(io::Error::other("command lifetime watchdog did not initialize"));
}
Ok(())
}
#[expect(
unsafe_code,
reason = "Minimal fork-child watchdog uses only async-signal-safe POSIX calls"
)]
unsafe fn run_watchdog(owner_fd: libc::c_int, ready_fd: libc::c_int, group: libc::pid_t) -> ! {
let mut ignore: libc::sigaction = unsafe { std::mem::zeroed() };
ignore.sa_sigaction = libc::SIG_IGN;
unsafe { libc::sigemptyset(&mut ignore.sa_mask) };
for signal in [libc::SIGINT, libc::SIGTERM, libc::SIGHUP, libc::SIGPIPE] {
if unsafe { libc::sigaction(signal, &ignore, std::ptr::null_mut()) } == -1 {
unsafe { libc::_exit(1) };
}
}
for fd in 0..=2 {
unsafe { libc::close(fd) };
}
let ready = READY;
if unsafe { libc::write(ready_fd, std::ptr::addr_of!(ready).cast(), 1) } != 1 {
unsafe { libc::_exit(1) };
}
unsafe { libc::close(ready_fd) };
let mut byte = 0u8;
loop {
let read = unsafe { libc::read(owner_fd, std::ptr::addr_of_mut!(byte).cast(), 1) };
if read == -1 && nix::errno::Errno::last() == nix::errno::Errno::EINTR {
continue;
}
unsafe { libc::kill(-group, libc::SIGKILL) };
unsafe { libc::_exit(1) };
}
}
#[cfg(test)]
mod tests {
use super::*;
use assert_fs::TempDir;
use std::os::unix::fs::{PermissionsExt, symlink};
use std::os::unix::net::UnixListener;
fn private_directory(root: &Path, name: &str) -> PathBuf {
let path = root.join(name);
std::fs::create_dir(&path).expect("create private directory");
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o700)).expect("owner-only directory");
path
}
fn bind_endpoint(directory: &Path, name: &str) -> PathBuf {
let socket = directory.join(name);
drop(UnixListener::bind(&socket).expect("bind endpoint"));
socket
}
#[test]
fn accepts_a_supervisor_shaped_endpoint() {
let root = TempDir::new().expect("temp root");
let directory = private_directory(root.path(), "vtc-accepted");
let socket = bind_endpoint(&directory, OWNER_SOCKET_FILE);
assert!(is_owned_endpoint(&socket));
assert!(is_owned_endpoint_under(&socket, &std::env::temp_dir()));
}
#[test]
fn ignores_a_stale_or_invalid_endpoint() {
let root = TempDir::new().expect("temp root");
let missing = root.path().join("vtc-stale").join(OWNER_SOCKET_FILE);
assert!(!is_owned_endpoint(&missing));
let directory = private_directory(root.path(), "vtc-stale");
std::fs::write(&missing, b"not a socket").expect("regular file");
assert!(!is_owned_endpoint(&missing));
let _foreign = bind_endpoint(&directory, "endpoint");
assert!(!is_owned_endpoint(&directory.join("endpoint")));
assert!(!is_owned_endpoint(Path::new(OWNER_SOCKET_FILE)));
}
#[test]
fn ignores_an_endpoint_outside_a_private_directory() {
let root = TempDir::new().expect("temp root");
let directory = root.path().join("vtc-public");
std::fs::create_dir(&directory).expect("public directory");
std::fs::set_permissions(&directory, std::fs::Permissions::from_mode(0o755)).expect("world-accessible");
let socket = bind_endpoint(&directory, OWNER_SOCKET_FILE);
assert!(!is_owned_endpoint(&socket));
}
#[test]
fn ignores_a_directory_the_supervisor_did_not_name() {
let root = TempDir::new().expect("temp root");
let directory = private_directory(root.path(), "other-owner");
let socket = bind_endpoint(&directory, OWNER_SOCKET_FILE);
assert!(!is_owned_endpoint(&socket));
}
#[test]
fn ignores_a_symlinked_endpoint() {
let root = TempDir::new().expect("temp root");
let target = bind_endpoint(root.path(), "target.sock");
let directory = private_directory(root.path(), "vtc-symlinked");
let socket = directory.join(OWNER_SOCKET_FILE);
symlink(&target, &socket).expect("symlink endpoint");
assert!(!is_owned_endpoint(&socket));
}
#[test]
fn ignores_an_endpoint_outside_the_temp_root() {
let root = TempDir::new().expect("temp root");
let other_root = TempDir::new().expect("other root");
let directory = private_directory(root.path(), "vtc-containment");
let socket = bind_endpoint(&directory, OWNER_SOCKET_FILE);
assert!(is_owned_endpoint(&socket));
assert!(!is_owned_endpoint_under(&socket, other_root.path()));
}
}