Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 70 additions & 38 deletions host/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1038,12 +1038,12 @@ use std::net::SocketAddr;
use std::sync::Mutex;

enum HostSocket {
Socket(Socket),
Socket(Socket, i32),
}

struct SocketTable {
sockets: HashMap<u32, HostSocket>,
next_id: u32,
sockets: HashMap<u64, HostSocket>,
next_id: u64,
}

impl SocketTable {
Expand All @@ -1054,26 +1054,32 @@ impl SocketTable {
}
}

fn insert(&mut self, sock: HostSocket) -> u32 {
fn insert(&mut self, sock: HostSocket) -> u64 {
let id = self.next_id;
self.next_id += 1;
self.sockets.insert(id, sock);
id
}

fn get(&self, fd: u32) -> Result<&HostSocket> {
fn get(&self, fd: u64) -> Result<&HostSocket> {
self.sockets
.get(&fd)
.ok_or_else(|| anyhow!("bad_fd: {}", fd))
}

fn get_socket(&self, fd: u32) -> Result<&Socket> {
fn get_socket(&self, fd: u64) -> Result<&Socket> {
match self.get(fd)? {
HostSocket::Socket(s) => Ok(s),
HostSocket::Socket(s, _) => Ok(s),
}
}

fn remove(&mut self, fd: u32) -> Result<()> {
fn get_sock_type(&self, fd: u64) -> Result<i32> {
match self.get(fd)? {
HostSocket::Socket(_, t) => Ok(*t),
}
}

fn remove(&mut self, fd: u64) -> Result<()> {
self.sockets
.remove(&fd)
.map(|_| ())
Expand Down Expand Up @@ -1136,15 +1142,18 @@ fn register_net_tools(
Some(Protocol::from(protocol))
};
let sock = Socket::new(domain, stype, proto)?;
let fd = t.lock().unwrap().insert(HostSocket::Socket(sock));
let fd = t
.lock()
.unwrap()
.insert(HostSocket::Socket(sock, sock_type));
Ok(json!({ "fd": fd }))
});

// net_connect
let t = table.clone();
let pol = policy.clone();
tools.register("net_connect", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let addr = parse_sockaddr(&args)?;
pol.check(&addr)?;
let sa: SockAddr = addr.into();
Expand All @@ -1158,7 +1167,7 @@ fn register_net_tools(
let t = table.clone();
let lp = listen_ports.cloned().map(Arc::new);
tools.register("net_bind", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let addr = parse_sockaddr(&args)?;
match lp.as_ref() {
Some(ports) => ports.check(addr.port())?,
Expand All @@ -1174,7 +1183,7 @@ fn register_net_tools(
// net_listen
let t = table.clone();
tools.register("net_listen", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let backlog = args["backlog"].as_i64().unwrap_or(128) as i32;
let tbl = t.lock().unwrap();
let sock = tbl.get_socket(fd)?;
Expand All @@ -1185,14 +1194,19 @@ fn register_net_tools(
// net_accept
let t = table.clone();
tools.register("net_accept", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let (new_sock, peer) = {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let (new_sock, peer, parent_type) = {
let tbl = t.lock().unwrap();
let sock = tbl.get_socket(fd)?;
sock.accept()?
let (s, p) = sock.accept()?;
let st = tbl.get_sock_type(fd)?;
(s, p, st)
};
let peer_addr: Option<SocketAddr> = peer.as_socket();
let new_fd = t.lock().unwrap().insert(HostSocket::Socket(new_sock));
let new_fd = t
.lock()
.unwrap()
.insert(HostSocket::Socket(new_sock, parent_type));
let mut resp = json!({ "fd": new_fd });
if let Some(pa) = peer_addr {
resp["addr"] = json!(pa.ip().to_string());
Expand All @@ -1204,7 +1218,7 @@ fn register_net_tools(
// net_send
let t = table.clone();
tools.register("net_send", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let data_b64 = args["data"]
.as_str()
.ok_or_else(|| anyhow!("missing 'data'"))?;
Expand All @@ -1221,7 +1235,7 @@ fn register_net_tools(
let t = table.clone();
let pol = policy.clone();
tools.register("net_sendto", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let data_b64 = args["data"]
.as_str()
.ok_or_else(|| anyhow!("missing 'data'"))?;
Expand All @@ -1240,7 +1254,7 @@ fn register_net_tools(
// net_recv (alias for net_recvfrom with no addr returned for stream)
let t = table.clone();
tools.register("net_recv", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let len = args["len"].as_u64().unwrap_or(4096) as usize;
let mut buf = vec![std::mem::MaybeUninit::uninit(); len.min(65536)];
let tbl = t.lock().unwrap();
Expand All @@ -1257,7 +1271,7 @@ fn register_net_tools(
// net_recvfrom
let t = table.clone();
tools.register("net_recvfrom", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let len = args["len"].as_u64().unwrap_or(4096) as usize;
let mut buf = vec![0u8; len.min(65536)];

Expand All @@ -1282,15 +1296,15 @@ fn register_net_tools(
// net_close
let t = table.clone();
tools.register("net_close", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
t.lock().unwrap().remove(fd)?;
Ok(json!({}))
});

// net_shutdown
let t = table.clone();
tools.register("net_shutdown", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let how = args["how"].as_i64().unwrap_or(2) as i32;
let shutdown = match how {
0 => std::net::Shutdown::Read,
Expand All @@ -1306,7 +1320,7 @@ fn register_net_tools(
// net_setsockopt
let t = table.clone();
tools.register("net_setsockopt", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let level = args["level"].as_i64().unwrap_or(0) as i32;
let optname = args["optname"].as_i64().unwrap_or(0) as i32;
let value = args["value"].as_i64().unwrap_or(0) as i32;
Expand All @@ -1330,14 +1344,14 @@ fn register_net_tools(
// net_getsockopt
let t = table.clone();
tools.register("net_getsockopt", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let level = args["level"].as_i64().unwrap_or(0) as i32;
let optname = args["optname"].as_i64().unwrap_or(0) as i32;
Comment on lines 1345 to 1349
let tbl = t.lock().unwrap();
let sock = tbl.get_socket(fd)?;
let val: i32 = if level == 1 && optname == 3 {
// SOL_SOCKET + SO_TYPE — all our sockets are SOCK_STREAM
1
// SOL_SOCKET + SO_TYPE — return the actual socket type
tbl.get_sock_type(fd)?
} else if level == 1 && optname == 2 {
sock.reuse_address()? as i32
} else if level == 6 && optname == 1 {
Expand All @@ -1351,7 +1365,7 @@ fn register_net_tools(
// net_getpeername
let t = table.clone();
tools.register("net_getpeername", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let tbl = t.lock().unwrap();
let sock = tbl.get_socket(fd)?;
let peer = sock.peer_addr()?;
Expand All @@ -1365,7 +1379,7 @@ fn register_net_tools(
// net_getsockname
let t = table.clone();
tools.register("net_getsockname", move |args| {
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))? as u32;
let fd = args["fd"].as_u64().ok_or_else(|| anyhow!("missing 'fd'"))?;
let tbl = t.lock().unwrap();
let sock = tbl.get_socket(fd)?;
let local = sock.local_addr()?;
Expand Down Expand Up @@ -2207,7 +2221,11 @@ pub fn run_vm_capture_output(
let setup_time = setup_start.elapsed();

// Redirect stderr to a temp file before the call phase
let capture_file = std::env::temp_dir().join(format!("hl-capture-{}", std::process::id()));
let capture_file = std::env::temp_dir().join(format!(
"hl-capture-{}-{:?}",
std::process::id(),
std::thread::current().id()
));
let capture = stderr_capture::Capture::redirect_to_file(&capture_file)?;

// Phase 2: restore + call — application runs and produces output
Expand Down Expand Up @@ -2800,14 +2818,11 @@ mod tests {

#[test]
fn test_fs_read_bytes_capped() {
// Request a huge len (well above MAX_FS_READ) on a small file.
// Without the cap this would try to allocate terabytes and OOM.
let root = tmpdir("readcap");
fs::write(root.join("small.bin"), b"hello").unwrap();
let mut reg = ToolRegistry::new();
FsSandbox::new(&root).unwrap().register(&mut reg);

// Ask for 1 TiB — the cap should silently clamp to MAX_FS_READ.
let req = br#"{"name":"fs_read_bytes","args":{"path":"small.bin","len":1099511627776}}"#;
let resp = reg.dispatch(req);
let s = std::str::from_utf8(&resp).unwrap();
Expand All @@ -2817,14 +2832,8 @@ mod tests {

#[test]
fn test_sleep_capped() {
// Verify the cap constant and that sleeping with a huge value
// completes quickly (the cap brings it down to 60 s max, but we
// pass 0 to keep the test instant — the important thing is
// confirming the cap constant exists and has the right value).
assert_eq!(MAX_SLEEP_NS, 60_000_000_000);

// Dispatch a sleep with ns=0 through the real handler to confirm
// the code path works.
let mut tools = ToolRegistry::new();
let exit_code = Arc::new(AtomicI32::new(0));
register_internal_tools(&mut tools, &exit_code, None, None);
Expand All @@ -2834,4 +2843,27 @@ mod tests {
let s = std::str::from_utf8(&resp).unwrap();
assert!(!s.contains("\"error\""), "sleep(0) should succeed: {s}");
}

#[test]
fn net_getsockopt_returns_correct_type_for_dgram() {
let mut reg = ToolRegistry::new();
let policy = NetworkPolicy::AllowAll;
register_net_tools(&mut reg, &policy, None);

let req = br#"{"name":"net_socket","args":{"family":2,"type":2}}"#;
let resp = std::str::from_utf8(&reg.dispatch(req)).unwrap().to_string();
let v: serde_json::Value = serde_json::from_str(&resp).unwrap();
let fd = v["result"]["fd"].as_u64().unwrap();

let req =
format!(r#"{{"name":"net_getsockopt","args":{{"fd":{fd},"level":1,"optname":3}}}}"#);
let resp = std::str::from_utf8(&reg.dispatch(req.as_bytes()))
.unwrap()
.to_string();
let v: serde_json::Value = serde_json::from_str(&resp).unwrap();
assert_eq!(
v["result"]["value"], 2,
"SO_TYPE should return 2 (DGRAM), got: {resp}"
);
}
}
42 changes: 33 additions & 9 deletions host/src/stderr_capture.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,28 +8,52 @@
mod imp {
use anyhow::Result;
use nix::unistd;
use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd};
use std::os::fd::{AsRawFd, FromRawFd, OwnedFd};
use std::path::Path;
use std::sync::{Mutex, MutexGuard};

/// Process-wide lock protecting fd 2 (stderr) manipulation.
/// Held from `redirect_to_file` until `restore` to prevent concurrent
/// VMs from racing on dup2.
static STDERR_LOCK: Mutex<()> = Mutex::new(());

pub struct Capture {
original_stderr: OwnedFd,
original_stderr: Option<OwnedFd>,
_guard: MutexGuard<'static, ()>,
}

impl Capture {
pub fn redirect_to_file(path: &Path) -> Result<Self> {
let capture_fd = std::fs::File::create(path)?.into_raw_fd();
let original_stderr_raw = unistd::dup(2)?;
unistd::dup2(capture_fd, 2)?;
unistd::close(capture_fd)?;
let original_stderr = unsafe { OwnedFd::from_raw_fd(original_stderr_raw) };
Ok(Self { original_stderr })
let guard = STDERR_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let capture_file = std::fs::File::create(path)?;
let original_stderr = unistd::dup(2).map(|fd| unsafe { OwnedFd::from_raw_fd(fd) })?;
unistd::dup2(capture_file.as_raw_fd(), 2)?;
// capture_file dropped here — its OwnedFd closes the fd via RAII
Ok(Self {
original_stderr: Some(original_stderr),
_guard: guard,
})
}

pub fn restore(self) -> Result<()> {
unistd::dup2(self.original_stderr.as_raw_fd(), 2)?;
if let Some(ref fd) = self.original_stderr {
unistd::dup2(fd.as_raw_fd(), 2)?;
}
// Drop restores via Drop impl if restore() wasn't called,
// but explicit call lets caller handle errors.
Ok(())
Comment on lines 40 to 46
}
}

impl Drop for Capture {
fn drop(&mut self) {
if let Some(ref fd) = self.original_stderr.take() {
let _ = unistd::dup2(fd.as_raw_fd(), 2);
}
}
}
}

#[cfg(windows)]
Expand Down
Loading