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:
438
envd-rs/src/rpc/process_service.rs
Normal file
438
envd-rs/src/rpc/process_service.rs
Normal file
@ -0,0 +1,438 @@
|
||||
use std::collections::HashMap;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use connectrpc::{ConnectError, Context, ErrorCode};
|
||||
use dashmap::DashMap;
|
||||
use futures::Stream;
|
||||
|
||||
use crate::permissions::path::expand_and_resolve;
|
||||
use crate::permissions::user::lookup_user;
|
||||
use crate::rpc::pb::process::*;
|
||||
use crate::rpc::process_handler::{self, DataEvent, ProcessHandle};
|
||||
use crate::state::AppState;
|
||||
|
||||
pub struct ProcessServiceImpl {
|
||||
state: Arc<AppState>,
|
||||
processes: DashMap<u32, Arc<ProcessHandle>>,
|
||||
}
|
||||
|
||||
impl ProcessServiceImpl {
|
||||
pub fn new(state: Arc<AppState>) -> Self {
|
||||
Self {
|
||||
state,
|
||||
processes: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn get_process_by_selector(
|
||||
&self,
|
||||
selector: &ProcessSelectorView,
|
||||
) -> Result<Arc<ProcessHandle>, ConnectError> {
|
||||
match &selector.selector {
|
||||
Some(process_selector::SelectorView::Pid(pid)) => {
|
||||
let pid_val = *pid;
|
||||
self.processes
|
||||
.get(&pid_val)
|
||||
.map(|entry| Arc::clone(entry.value()))
|
||||
.ok_or_else(|| {
|
||||
ConnectError::new(
|
||||
ErrorCode::NotFound,
|
||||
format!("process with pid {pid_val} not found"),
|
||||
)
|
||||
})
|
||||
}
|
||||
Some(process_selector::SelectorView::Tag(tag)) => {
|
||||
let tag_str: &str = tag;
|
||||
for entry in self.processes.iter() {
|
||||
if let Some(ref t) = entry.value().tag {
|
||||
if t == tag_str {
|
||||
return Ok(Arc::clone(entry.value()));
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(ConnectError::new(
|
||||
ErrorCode::NotFound,
|
||||
format!("process with tag {tag_str} not found"),
|
||||
))
|
||||
}
|
||||
None => Err(ConnectError::new(
|
||||
ErrorCode::InvalidArgument,
|
||||
"process selector required",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn spawn_from_request(
|
||||
&self,
|
||||
request: &StartRequestView<'_>,
|
||||
) -> Result<Arc<ProcessHandle>, ConnectError> {
|
||||
let proc_config = request.process.as_option().ok_or_else(|| {
|
||||
ConnectError::new(ErrorCode::InvalidArgument, "process config required")
|
||||
})?;
|
||||
|
||||
let username = self.state.defaults.user.clone();
|
||||
let user =
|
||||
lookup_user(&username).map_err(|e| ConnectError::new(ErrorCode::Internal, e))?;
|
||||
|
||||
let cmd: &str = proc_config.cmd;
|
||||
let args: Vec<String> = proc_config.args.iter().map(|s| s.to_string()).collect();
|
||||
let envs: HashMap<String, String> = proc_config
|
||||
.envs
|
||||
.iter()
|
||||
.map(|(k, v)| (k.to_string(), v.to_string()))
|
||||
.collect();
|
||||
|
||||
let home_dir = format!("/home/{}", user.name);
|
||||
let cwd_str: &str = proc_config.cwd.unwrap_or("");
|
||||
let cwd = expand_and_resolve(cwd_str, &home_dir, self.state.defaults.workdir.as_deref())
|
||||
.map_err(|e| ConnectError::new(ErrorCode::InvalidArgument, e))?;
|
||||
|
||||
let effective_cwd = if cwd.is_empty() { "/" } else { &cwd };
|
||||
if let Err(_) = std::fs::metadata(effective_cwd) {
|
||||
return Err(ConnectError::new(
|
||||
ErrorCode::InvalidArgument,
|
||||
format!("cwd '{effective_cwd}' does not exist"),
|
||||
));
|
||||
}
|
||||
|
||||
let pty_opts = request.pty.as_option().and_then(|pty| {
|
||||
pty.size
|
||||
.as_option()
|
||||
.map(|sz| (sz.cols as u16, sz.rows as u16))
|
||||
});
|
||||
|
||||
let enable_stdin = request.stdin.unwrap_or(true);
|
||||
let tag = request.tag.map(|s| s.to_string());
|
||||
|
||||
let handle = process_handler::spawn_process(
|
||||
cmd,
|
||||
&args,
|
||||
&envs,
|
||||
effective_cwd,
|
||||
pty_opts,
|
||||
enable_stdin,
|
||||
tag,
|
||||
&user,
|
||||
&self.state.defaults.env_vars,
|
||||
)?;
|
||||
|
||||
self.processes.insert(handle.pid, Arc::clone(&handle));
|
||||
|
||||
let processes = self.processes.clone();
|
||||
let pid = handle.pid;
|
||||
let mut end_rx = handle.subscribe_end();
|
||||
tokio::spawn(async move {
|
||||
let _ = end_rx.recv().await;
|
||||
processes.remove(&pid);
|
||||
});
|
||||
|
||||
Ok(handle)
|
||||
}
|
||||
}
|
||||
|
||||
impl Process for ProcessServiceImpl {
|
||||
async fn list(
|
||||
&self,
|
||||
ctx: Context,
|
||||
_request: buffa::view::OwnedView<ListRequestView<'static>>,
|
||||
) -> Result<(ListResponse, Context), ConnectError> {
|
||||
let processes: Vec<ProcessInfo> = self
|
||||
.processes
|
||||
.iter()
|
||||
.map(|entry| {
|
||||
let h = entry.value();
|
||||
ProcessInfo {
|
||||
config: buffa::MessageField::some(h.config.clone()),
|
||||
pid: h.pid,
|
||||
tag: h.tag.clone(),
|
||||
..Default::default()
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok((
|
||||
ListResponse {
|
||||
processes,
|
||||
..Default::default()
|
||||
},
|
||||
ctx,
|
||||
))
|
||||
}
|
||||
|
||||
async fn start(
|
||||
&self,
|
||||
ctx: Context,
|
||||
request: buffa::view::OwnedView<StartRequestView<'static>>,
|
||||
) -> Result<
|
||||
(
|
||||
Pin<Box<dyn Stream<Item = Result<StartResponse, ConnectError>> + Send>>,
|
||||
Context,
|
||||
),
|
||||
ConnectError,
|
||||
> {
|
||||
let handle = self.spawn_from_request(&request)?;
|
||||
let pid = handle.pid;
|
||||
|
||||
let mut data_rx = handle.subscribe_data();
|
||||
let mut end_rx = handle.subscribe_end();
|
||||
|
||||
let stream = async_stream::stream! {
|
||||
yield Ok(make_start_response(pid));
|
||||
|
||||
loop {
|
||||
match data_rx.recv().await {
|
||||
Ok(ev) => yield Ok(make_data_start_response(ev)),
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
}
|
||||
}
|
||||
|
||||
if let Ok(end) = end_rx.recv().await {
|
||||
yield Ok(make_end_start_response(end));
|
||||
}
|
||||
};
|
||||
|
||||
Ok((Box::pin(stream), ctx))
|
||||
}
|
||||
|
||||
async fn connect(
|
||||
&self,
|
||||
ctx: Context,
|
||||
request: buffa::view::OwnedView<ConnectRequestView<'static>>,
|
||||
) -> Result<
|
||||
(
|
||||
Pin<Box<dyn Stream<Item = Result<ConnectResponse, ConnectError>> + Send>>,
|
||||
Context,
|
||||
),
|
||||
ConnectError,
|
||||
> {
|
||||
let selector = request.process.as_option().ok_or_else(|| {
|
||||
ConnectError::new(ErrorCode::InvalidArgument, "process selector required")
|
||||
})?;
|
||||
let handle = self.get_process_by_selector(selector)?;
|
||||
let pid = handle.pid;
|
||||
|
||||
let mut data_rx = handle.subscribe_data();
|
||||
let mut end_rx = handle.subscribe_end();
|
||||
|
||||
let stream = async_stream::stream! {
|
||||
yield Ok(ConnectResponse {
|
||||
event: buffa::MessageField::some(ProcessEvent {
|
||||
event: Some(process_event::Event::Start(Box::new(
|
||||
process_event::StartEvent { pid, ..Default::default() },
|
||||
))),
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
});
|
||||
|
||||
loop {
|
||||
match data_rx.recv().await {
|
||||
Ok(ev) => {
|
||||
yield Ok(ConnectResponse {
|
||||
event: buffa::MessageField::some(make_data_event(ev)),
|
||||
..Default::default()
|
||||
});
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
}
|
||||
}
|
||||
|
||||
if let Ok(end) = end_rx.recv().await {
|
||||
yield Ok(ConnectResponse {
|
||||
event: buffa::MessageField::some(make_end_event(end)),
|
||||
..Default::default()
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
Ok((Box::pin(stream), ctx))
|
||||
}
|
||||
|
||||
async fn update(
|
||||
&self,
|
||||
ctx: Context,
|
||||
request: buffa::view::OwnedView<UpdateRequestView<'static>>,
|
||||
) -> Result<(UpdateResponse, Context), ConnectError> {
|
||||
let selector = request.process.as_option().ok_or_else(|| {
|
||||
ConnectError::new(ErrorCode::InvalidArgument, "process selector required")
|
||||
})?;
|
||||
let handle = self.get_process_by_selector(selector)?;
|
||||
|
||||
if let Some(pty) = request.pty.as_option() {
|
||||
if let Some(size) = pty.size.as_option() {
|
||||
handle.resize_pty(size.cols as u16, size.rows as u16)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok((UpdateResponse { ..Default::default() }, ctx))
|
||||
}
|
||||
|
||||
async fn stream_input(
|
||||
&self,
|
||||
ctx: Context,
|
||||
mut requests: Pin<
|
||||
Box<
|
||||
dyn Stream<
|
||||
Item = Result<
|
||||
buffa::view::OwnedView<StreamInputRequestView<'static>>,
|
||||
ConnectError,
|
||||
>,
|
||||
> + Send,
|
||||
>,
|
||||
>,
|
||||
) -> Result<(StreamInputResponse, Context), ConnectError> {
|
||||
use futures::StreamExt;
|
||||
|
||||
let mut handle: Option<Arc<ProcessHandle>> = None;
|
||||
|
||||
while let Some(result) = requests.next().await {
|
||||
let req = result?;
|
||||
match &req.event {
|
||||
Some(stream_input_request::EventView::Start(start)) => {
|
||||
if let Some(selector) = start.process.as_option() {
|
||||
handle = Some(self.get_process_by_selector(selector)?);
|
||||
}
|
||||
}
|
||||
Some(stream_input_request::EventView::Data(data)) => {
|
||||
let h = handle.as_ref().ok_or_else(|| {
|
||||
ConnectError::new(ErrorCode::FailedPrecondition, "no start event received")
|
||||
})?;
|
||||
if let Some(input) = data.input.as_option() {
|
||||
write_input(h, input)?;
|
||||
}
|
||||
}
|
||||
Some(stream_input_request::EventView::Keepalive(_)) => {}
|
||||
None => {}
|
||||
}
|
||||
}
|
||||
|
||||
Ok((StreamInputResponse { ..Default::default() }, ctx))
|
||||
}
|
||||
|
||||
async fn send_input(
|
||||
&self,
|
||||
ctx: Context,
|
||||
request: buffa::view::OwnedView<SendInputRequestView<'static>>,
|
||||
) -> Result<(SendInputResponse, Context), ConnectError> {
|
||||
let selector = request.process.as_option().ok_or_else(|| {
|
||||
ConnectError::new(ErrorCode::InvalidArgument, "process selector required")
|
||||
})?;
|
||||
let handle = self.get_process_by_selector(selector)?;
|
||||
|
||||
if let Some(input) = request.input.as_option() {
|
||||
write_input(&handle, input)?;
|
||||
}
|
||||
|
||||
Ok((SendInputResponse { ..Default::default() }, ctx))
|
||||
}
|
||||
|
||||
async fn send_signal(
|
||||
&self,
|
||||
ctx: Context,
|
||||
request: buffa::view::OwnedView<SendSignalRequestView<'static>>,
|
||||
) -> Result<(SendSignalResponse, Context), ConnectError> {
|
||||
let selector = request.process.as_option().ok_or_else(|| {
|
||||
ConnectError::new(ErrorCode::InvalidArgument, "process selector required")
|
||||
})?;
|
||||
let handle = self.get_process_by_selector(selector)?;
|
||||
|
||||
let sig = match request.signal.as_known() {
|
||||
Some(Signal::SIGNAL_SIGKILL) => nix::sys::signal::Signal::SIGKILL,
|
||||
Some(Signal::SIGNAL_SIGTERM) => nix::sys::signal::Signal::SIGTERM,
|
||||
_ => {
|
||||
return Err(ConnectError::new(
|
||||
ErrorCode::InvalidArgument,
|
||||
"invalid or unspecified signal",
|
||||
))
|
||||
}
|
||||
};
|
||||
|
||||
handle.send_signal(sig)?;
|
||||
Ok((SendSignalResponse { ..Default::default() }, ctx))
|
||||
}
|
||||
|
||||
async fn close_stdin(
|
||||
&self,
|
||||
ctx: Context,
|
||||
request: buffa::view::OwnedView<CloseStdinRequestView<'static>>,
|
||||
) -> Result<(CloseStdinResponse, Context), ConnectError> {
|
||||
let selector = request.process.as_option().ok_or_else(|| {
|
||||
ConnectError::new(ErrorCode::InvalidArgument, "process selector required")
|
||||
})?;
|
||||
let handle = self.get_process_by_selector(selector)?;
|
||||
handle.close_stdin()?;
|
||||
Ok((CloseStdinResponse { ..Default::default() }, ctx))
|
||||
}
|
||||
}
|
||||
|
||||
fn write_input(handle: &ProcessHandle, input: &ProcessInputView) -> Result<(), ConnectError> {
|
||||
match &input.input {
|
||||
Some(process_input::InputView::Pty(d)) => handle.write_pty(d),
|
||||
Some(process_input::InputView::Stdin(d)) => handle.write_stdin(d),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
fn make_start_response(pid: u32) -> StartResponse {
|
||||
StartResponse {
|
||||
event: buffa::MessageField::some(ProcessEvent {
|
||||
event: Some(process_event::Event::Start(Box::new(
|
||||
process_event::StartEvent {
|
||||
pid,
|
||||
..Default::default()
|
||||
},
|
||||
))),
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn make_data_event(ev: DataEvent) -> ProcessEvent {
|
||||
let output = match ev {
|
||||
DataEvent::Stdout(d) => Some(process_event::data_event::Output::Stdout(d.into())),
|
||||
DataEvent::Stderr(d) => Some(process_event::data_event::Output::Stderr(d.into())),
|
||||
DataEvent::Pty(d) => Some(process_event::data_event::Output::Pty(d.into())),
|
||||
};
|
||||
ProcessEvent {
|
||||
event: Some(process_event::Event::Data(Box::new(
|
||||
process_event::DataEvent {
|
||||
output,
|
||||
..Default::default()
|
||||
},
|
||||
))),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn make_data_start_response(ev: DataEvent) -> StartResponse {
|
||||
StartResponse {
|
||||
event: buffa::MessageField::some(make_data_event(ev)),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn make_end_event(end: process_handler::EndEvent) -> ProcessEvent {
|
||||
ProcessEvent {
|
||||
event: Some(process_event::Event::End(Box::new(
|
||||
process_event::EndEvent {
|
||||
exit_code: end.exit_code,
|
||||
exited: end.exited,
|
||||
status: end.status,
|
||||
error: end.error,
|
||||
..Default::default()
|
||||
},
|
||||
))),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn make_end_start_response(end: process_handler::EndEvent) -> StartResponse {
|
||||
StartResponse {
|
||||
event: buffa::MessageField::some(make_end_event(end)),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user