forked from wrenn/wrenn
feat: rewrite envd guest agent in Rust (envd-rs)
Complete Rust rewrite of the Go envd guest daemon that runs as PID 1 inside Firecracker microVMs. Feature-complete across all 8 phases: - Health, metrics, and env var endpoints - Crypto (SHA-256/512, HMAC), auth (secure token, signing), init/snapshot - Connect RPC via connectrpc + buffa (process + filesystem services) - File transfer (GET/POST /files) with gzip, multipart, chown, ENOSPC - Port subsystem (/proc/net/tcp scanner, socat forwarder) - Cgroup2 manager with noop fallback - Snapshot/restore lifecycle (conntracker, port subsystem stop/restart) - SIGTERM graceful shutdown, --cmd initial process spawn - MMDS metadata polling for Firecracker mode 42 source files, ~4200 LOC, 4.1MB stripped release binary. Makefile updated: build-envd now targets Rust (musl static), build-envd-go preserved for Go builds.
This commit is contained in:
400
envd-rs/src/rpc/process_handler.rs
Normal file
400
envd-rs/src/rpc/process_handler.rs
Normal file
@ -0,0 +1,400 @@
|
||||
use std::io::Read;
|
||||
use std::os::unix::process::CommandExt;
|
||||
use std::process::Stdio;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use connectrpc::{ConnectError, ErrorCode};
|
||||
use nix::pty::{openpty, Winsize};
|
||||
use nix::sys::signal::{self, Signal};
|
||||
use nix::unistd::Pid;
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
use crate::rpc::pb::process::*;
|
||||
|
||||
const STD_CHUNK_SIZE: usize = 32768;
|
||||
const PTY_CHUNK_SIZE: usize = 16384;
|
||||
const BROADCAST_CAPACITY: usize = 4096;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub enum DataEvent {
|
||||
Stdout(Vec<u8>),
|
||||
Stderr(Vec<u8>),
|
||||
Pty(Vec<u8>),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct EndEvent {
|
||||
pub exit_code: i32,
|
||||
pub exited: bool,
|
||||
pub status: String,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
pub struct ProcessHandle {
|
||||
pub config: ProcessConfig,
|
||||
pub tag: Option<String>,
|
||||
pub pid: u32,
|
||||
|
||||
data_tx: broadcast::Sender<DataEvent>,
|
||||
end_tx: broadcast::Sender<EndEvent>,
|
||||
|
||||
stdin: Mutex<Option<std::process::ChildStdin>>,
|
||||
pty_master: Mutex<Option<std::fs::File>>,
|
||||
}
|
||||
|
||||
impl ProcessHandle {
|
||||
pub fn subscribe_data(&self) -> broadcast::Receiver<DataEvent> {
|
||||
self.data_tx.subscribe()
|
||||
}
|
||||
|
||||
pub fn subscribe_end(&self) -> broadcast::Receiver<EndEvent> {
|
||||
self.end_tx.subscribe()
|
||||
}
|
||||
|
||||
pub fn send_signal(&self, sig: Signal) -> Result<(), ConnectError> {
|
||||
signal::kill(Pid::from_raw(self.pid as i32), sig).map_err(|e| {
|
||||
ConnectError::new(ErrorCode::Internal, format!("error sending signal: {e}"))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn write_stdin(&self, data: &[u8]) -> Result<(), ConnectError> {
|
||||
use std::io::Write;
|
||||
let mut guard = self.stdin.lock().unwrap();
|
||||
match guard.as_mut() {
|
||||
Some(stdin) => stdin.write_all(data).map_err(|e| {
|
||||
ConnectError::new(ErrorCode::Internal, format!("error writing to stdin: {e}"))
|
||||
}),
|
||||
None => Err(ConnectError::new(
|
||||
ErrorCode::FailedPrecondition,
|
||||
"stdin not enabled or closed",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn write_pty(&self, data: &[u8]) -> Result<(), ConnectError> {
|
||||
use std::io::Write;
|
||||
let mut guard = self.pty_master.lock().unwrap();
|
||||
match guard.as_mut() {
|
||||
Some(master) => master.write_all(data).map_err(|e| {
|
||||
ConnectError::new(ErrorCode::Internal, format!("error writing to pty: {e}"))
|
||||
}),
|
||||
None => Err(ConnectError::new(
|
||||
ErrorCode::FailedPrecondition,
|
||||
"pty not assigned to process",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn close_stdin(&self) -> Result<(), ConnectError> {
|
||||
if self.pty_master.lock().unwrap().is_some() {
|
||||
return Err(ConnectError::new(
|
||||
ErrorCode::FailedPrecondition,
|
||||
"cannot close stdin for PTY process — send Ctrl+D (0x04) instead",
|
||||
));
|
||||
}
|
||||
let mut guard = self.stdin.lock().unwrap();
|
||||
*guard = None;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn resize_pty(&self, cols: u16, rows: u16) -> Result<(), ConnectError> {
|
||||
let guard = self.pty_master.lock().unwrap();
|
||||
match guard.as_ref() {
|
||||
Some(master) => {
|
||||
use std::os::unix::io::AsRawFd;
|
||||
let ws = libc::winsize {
|
||||
ws_row: rows,
|
||||
ws_col: cols,
|
||||
ws_xpixel: 0,
|
||||
ws_ypixel: 0,
|
||||
};
|
||||
let ret = unsafe { libc::ioctl(master.as_raw_fd(), libc::TIOCSWINSZ, &ws) };
|
||||
if ret != 0 {
|
||||
return Err(ConnectError::new(
|
||||
ErrorCode::Internal,
|
||||
format!(
|
||||
"ioctl TIOCSWINSZ failed: {}",
|
||||
std::io::Error::last_os_error()
|
||||
),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
None => Err(ConnectError::new(
|
||||
ErrorCode::FailedPrecondition,
|
||||
"tty not assigned to process",
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn spawn_process(
|
||||
cmd_str: &str,
|
||||
args: &[String],
|
||||
envs: &std::collections::HashMap<String, String>,
|
||||
cwd: &str,
|
||||
pty_opts: Option<(u16, u16)>,
|
||||
enable_stdin: bool,
|
||||
tag: Option<String>,
|
||||
user: &nix::unistd::User,
|
||||
default_env_vars: &dashmap::DashMap<String, String>,
|
||||
) -> Result<Arc<ProcessHandle>, ConnectError> {
|
||||
let mut env: Vec<(String, String)> = Vec::new();
|
||||
env.push(("PATH".into(), std::env::var("PATH").unwrap_or_default()));
|
||||
let home = format!("/home/{}", user.name);
|
||||
env.push(("HOME".into(), home));
|
||||
env.push(("USER".into(), user.name.clone()));
|
||||
env.push(("LOGNAME".into(), user.name.clone()));
|
||||
|
||||
default_env_vars.iter().for_each(|entry| {
|
||||
env.push((entry.key().clone(), entry.value().clone()));
|
||||
});
|
||||
|
||||
for (k, v) in envs {
|
||||
env.push((k.clone(), v.clone()));
|
||||
}
|
||||
|
||||
let nice_delta = 0 - current_nice();
|
||||
let oom_script = format!(
|
||||
r#"echo 100 > /proc/$$/oom_score_adj && exec /usr/bin/nice -n {} "${{@}}""#,
|
||||
nice_delta
|
||||
);
|
||||
let mut wrapper_args = vec![
|
||||
"-c".to_string(),
|
||||
oom_script,
|
||||
"--".to_string(),
|
||||
cmd_str.to_string(),
|
||||
];
|
||||
wrapper_args.extend_from_slice(args);
|
||||
|
||||
let uid = user.uid.as_raw();
|
||||
let gid = user.gid.as_raw();
|
||||
|
||||
let (data_tx, _) = broadcast::channel(BROADCAST_CAPACITY);
|
||||
let (end_tx, _) = broadcast::channel(16);
|
||||
|
||||
let config = ProcessConfig {
|
||||
cmd: cmd_str.to_string(),
|
||||
args: args.to_vec(),
|
||||
envs: envs.clone(),
|
||||
cwd: Some(cwd.to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
if let Some((cols, rows)) = pty_opts {
|
||||
let pty_result = openpty(
|
||||
Some(&Winsize {
|
||||
ws_row: rows,
|
||||
ws_col: cols,
|
||||
ws_xpixel: 0,
|
||||
ws_ypixel: 0,
|
||||
}),
|
||||
None,
|
||||
)
|
||||
.map_err(|e| ConnectError::new(ErrorCode::Internal, format!("openpty failed: {e}")))?;
|
||||
|
||||
let master_fd = pty_result.master;
|
||||
let slave_fd = pty_result.slave;
|
||||
|
||||
let mut command = std::process::Command::new("/bin/sh");
|
||||
command
|
||||
.args(&wrapper_args)
|
||||
.env_clear()
|
||||
.envs(env.iter().map(|(k, v)| (k.as_str(), v.as_str())))
|
||||
.current_dir(cwd);
|
||||
|
||||
unsafe {
|
||||
use std::os::unix::io::AsRawFd;
|
||||
let slave_raw = slave_fd.as_raw_fd();
|
||||
command.pre_exec(move || {
|
||||
nix::unistd::setsid()
|
||||
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
|
||||
libc::ioctl(slave_raw, libc::TIOCSCTTY, 0);
|
||||
libc::dup2(slave_raw, 0);
|
||||
libc::dup2(slave_raw, 1);
|
||||
libc::dup2(slave_raw, 2);
|
||||
if slave_raw > 2 {
|
||||
libc::close(slave_raw);
|
||||
}
|
||||
libc::setgid(gid);
|
||||
libc::setuid(uid);
|
||||
Ok(())
|
||||
});
|
||||
}
|
||||
|
||||
command.stdin(Stdio::null());
|
||||
command.stdout(Stdio::null());
|
||||
command.stderr(Stdio::null());
|
||||
|
||||
let child = command.spawn().map_err(|e| {
|
||||
ConnectError::new(ErrorCode::Internal, format!("error starting pty process: {e}"))
|
||||
})?;
|
||||
|
||||
drop(slave_fd);
|
||||
|
||||
let pid = child.id();
|
||||
let master_file: std::fs::File = master_fd.into();
|
||||
let master_clone = master_file.try_clone().unwrap();
|
||||
|
||||
let handle = Arc::new(ProcessHandle {
|
||||
config,
|
||||
tag,
|
||||
pid,
|
||||
data_tx: data_tx.clone(),
|
||||
end_tx: end_tx.clone(),
|
||||
stdin: Mutex::new(None),
|
||||
pty_master: Mutex::new(Some(master_file)),
|
||||
});
|
||||
|
||||
let data_tx_clone = data_tx.clone();
|
||||
std::thread::spawn(move || {
|
||||
let mut master = master_clone;
|
||||
let mut buf = vec![0u8; PTY_CHUNK_SIZE];
|
||||
loop {
|
||||
match master.read(&mut buf) {
|
||||
Ok(0) => break,
|
||||
Ok(n) => {
|
||||
let _ = data_tx_clone.send(DataEvent::Pty(buf[..n].to_vec()));
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let end_tx_clone = end_tx.clone();
|
||||
std::thread::spawn(move || {
|
||||
let mut child = child;
|
||||
match child.wait() {
|
||||
Ok(s) => {
|
||||
let _ = end_tx_clone.send(EndEvent {
|
||||
exit_code: s.code().unwrap_or(-1),
|
||||
exited: s.code().is_some(),
|
||||
status: format!("{s}"),
|
||||
error: None,
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = end_tx_clone.send(EndEvent {
|
||||
exit_code: -1,
|
||||
exited: false,
|
||||
status: "error".into(),
|
||||
error: Some(e.to_string()),
|
||||
});
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
tracing::info!(pid, cmd = cmd_str, "process started (pty)");
|
||||
Ok(handle)
|
||||
} else {
|
||||
let mut command = std::process::Command::new("/bin/sh");
|
||||
command
|
||||
.args(&wrapper_args)
|
||||
.env_clear()
|
||||
.envs(env.iter().map(|(k, v)| (k.as_str(), v.as_str())))
|
||||
.current_dir(cwd)
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped());
|
||||
|
||||
if enable_stdin {
|
||||
command.stdin(Stdio::piped());
|
||||
} else {
|
||||
command.stdin(Stdio::null());
|
||||
}
|
||||
|
||||
unsafe {
|
||||
command.pre_exec(move || {
|
||||
libc::setgid(gid);
|
||||
libc::setuid(uid);
|
||||
Ok(())
|
||||
});
|
||||
}
|
||||
|
||||
let mut child = command.spawn().map_err(|e| {
|
||||
ConnectError::new(ErrorCode::Internal, format!("error starting process: {e}"))
|
||||
})?;
|
||||
|
||||
let pid = child.id();
|
||||
let stdin = child.stdin.take();
|
||||
let stdout = child.stdout.take();
|
||||
let stderr = child.stderr.take();
|
||||
|
||||
let handle = Arc::new(ProcessHandle {
|
||||
config,
|
||||
tag,
|
||||
pid,
|
||||
data_tx: data_tx.clone(),
|
||||
end_tx: end_tx.clone(),
|
||||
stdin: Mutex::new(stdin),
|
||||
pty_master: Mutex::new(None),
|
||||
});
|
||||
|
||||
if let Some(mut out) = stdout {
|
||||
let tx = data_tx.clone();
|
||||
std::thread::spawn(move || {
|
||||
let mut buf = vec![0u8; STD_CHUNK_SIZE];
|
||||
loop {
|
||||
match out.read(&mut buf) {
|
||||
Ok(0) => break,
|
||||
Ok(n) => {
|
||||
let _ = tx.send(DataEvent::Stdout(buf[..n].to_vec()));
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(mut err_pipe) = stderr {
|
||||
let tx = data_tx.clone();
|
||||
std::thread::spawn(move || {
|
||||
let mut buf = vec![0u8; STD_CHUNK_SIZE];
|
||||
loop {
|
||||
match err_pipe.read(&mut buf) {
|
||||
Ok(0) => break,
|
||||
Ok(n) => {
|
||||
let _ = tx.send(DataEvent::Stderr(buf[..n].to_vec()));
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
let end_tx_clone = end_tx.clone();
|
||||
std::thread::spawn(move || {
|
||||
match child.wait() {
|
||||
Ok(s) => {
|
||||
let _ = end_tx_clone.send(EndEvent {
|
||||
exit_code: s.code().unwrap_or(-1),
|
||||
exited: s.code().is_some(),
|
||||
status: format!("{s}"),
|
||||
error: None,
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = end_tx_clone.send(EndEvent {
|
||||
exit_code: -1,
|
||||
exited: false,
|
||||
status: "error".into(),
|
||||
error: Some(e.to_string()),
|
||||
});
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
tracing::info!(pid, cmd = cmd_str, "process started (pipe)");
|
||||
Ok(handle)
|
||||
}
|
||||
}
|
||||
|
||||
fn current_nice() -> i32 {
|
||||
unsafe {
|
||||
*libc::__errno_location() = 0;
|
||||
let prio = libc::getpriority(libc::PRIO_PROCESS, 0);
|
||||
if *libc::__errno_location() != 0 {
|
||||
return 0;
|
||||
}
|
||||
20 - prio
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user