1
0
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:
2026-05-03 02:47:15 +06:00
parent 3deecbff89
commit 0b53d34417
54 changed files with 7089 additions and 17 deletions

View File

@ -0,0 +1,2 @@
[target.x86_64-unknown-linux-musl]
linker = "musl-gcc"

2622
envd-rs/Cargo.lock generated Normal file

File diff suppressed because it is too large Load Diff

83
envd-rs/Cargo.toml Normal file
View File

@ -0,0 +1,83 @@
[package]
name = "envd"
version = "0.1.2"
edition = "2024"
rust-version = "1.88"
[dependencies]
# Async runtime
tokio = { version = "1", features = ["full"] }
# HTTP framework
axum = { version = "0.8", features = ["multipart"] }
tower = { version = "0.5", features = ["util"] }
tower-http = { version = "0.6", features = ["cors", "fs"] }
tower-service = "0.3"
# RPC (Connect protocol — serves Connect + gRPC + gRPC-Web on same port)
connectrpc = { version = "0.3", features = ["axum"] }
buffa-types = { path = "buffa-types-shim" }
# CLI
clap = { version = "4", features = ["derive"] }
# Serialization
serde = { version = "1", features = ["derive"] }
serde_json = "1"
# Logging
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["json", "env-filter"] }
# System metrics
sysinfo = "0.33"
# Unix syscalls
nix = { version = "0.30", features = ["fs", "process", "signal", "user", "term", "mount", "ioctl"] }
# Concurrent map
dashmap = "6"
# Crypto
sha2 = "0.10"
hmac = "0.12"
hex = "0.4"
base64 = "0.22"
# Secure memory
zeroize = { version = "1", features = ["derive"] }
# File watching
notify = "7"
# Compression
flate2 = "1"
# HTTP client (MMDS polling)
reqwest = { version = "0.12", default-features = false, features = ["json"] }
# Directory walking
walkdir = "2"
# Misc
libc = "0.2"
bytes = "1"
http = "1"
http-body-util = "0.1"
futures = "0.3"
tokio-util = { version = "0.7", features = ["io"] }
subtle = "2"
http-body = "1.0.1"
buffa = "0.3"
async-stream = "0.3.6"
mime_guess = "2"
[build-dependencies]
connectrpc-build = "0.3"
[profile.release]
strip = true
lto = true
opt-level = "z"
codegen-units = 1
panic = "abort"

141
envd-rs/README.md Normal file
View File

@ -0,0 +1,141 @@
# envd (Rust)
Wrenn guest agent daemon — runs as PID 1 inside Firecracker microVMs. Provides process management, filesystem operations, file transfer, port forwarding, and VM lifecycle control over Connect RPC and HTTP.
Rust rewrite of `envd/` (Go). Drop-in replacement — same wire protocol, same endpoints, same CLI flags.
## Prerequisites
- Rust 1.88+ (required by `connectrpc` 0.3.3)
- `protoc` (protobuf compiler, for proto codegen at build time)
- `musl-tools` (for static linking)
```bash
# Ubuntu/Debian
sudo apt install musl-tools protobuf-compiler
# Rust musl target
rustup target add x86_64-unknown-linux-musl
```
## Building
### Static binary (production — what goes into the rootfs)
```bash
cd envd-rs
ENVD_COMMIT=$(git rev-parse --short HEAD) \
cargo build --release --target x86_64-unknown-linux-musl
```
Output: `target/x86_64-unknown-linux-musl/release/envd`
Verify static linking:
```bash
file target/x86_64-unknown-linux-musl/release/envd
# should say: "statically linked"
ldd target/x86_64-unknown-linux-musl/release/envd
# should say: "not a dynamic executable"
```
### Debug binary (dev machine, dynamically linked)
```bash
cd envd-rs
cargo build
```
Run locally (outside a VM):
```bash
./target/debug/envd --isnotfc --port 49983
```
### Via Makefile (from repo root)
```bash
make build-envd # static musl release build
make build-envd-go # Go version (for comparison)
```
## CLI Flags
```
--port <PORT> Listen port [default: 49983]
--isnotfc Not running inside Firecracker (disables MMDS, cgroups)
--version Print version and exit
--commit Print git commit and exit
--cmd <CMD> Spawn a process at startup (e.g. --cmd "/bin/bash")
--cgroup-root <PATH> Cgroup v2 root [default: /sys/fs/cgroup]
```
## Endpoints
### HTTP
| Method | Path | Description |
|--------|---------------------|--------------------------------------|
| GET | `/health` | Health check, triggers post-restore |
| GET | `/metrics` | System metrics (CPU, memory, disk) |
| GET | `/envs` | Current environment variables |
| POST | `/init` | Host agent init (token, env, mounts) |
| POST | `/snapshot/prepare` | Quiesce before Firecracker snapshot |
| GET | `/files` | Download file (gzip, range support) |
| POST | `/files` | Upload file(s) via multipart |
### Connect RPC (same port)
| Service | RPCs |
|------------|-------------------------------------------------------------------------|
| Process | List, Start, Connect, Update, StreamInput, SendInput, SendSignal, CloseStdin |
| Filesystem | Stat, MakeDir, Move, ListDir, Remove, WatchDir, CreateWatcher, GetWatcherEvents, RemoveWatcher |
## Architecture
```
42 files, ~4200 LOC Rust
Binary: ~4 MB (stripped, LTO, musl static)
src/
├── main.rs # Entry point, CLI, server setup
├── state.rs # Shared AppState
├── config.rs # Constants
├── conntracker.rs # TCP connection tracking for snapshot/restore
├─<EFBFBD><EFBFBD><EFBFBD> execcontext.rs # Default user/workdir/env
├── logging.rs # tracing-subscriber (JSON or pretty)
├── util.rs # AtomicMax
├── auth/ # Token, signing, middleware
├── crypto/ # SHA-256, SHA-512, HMAC
├── host/ # MMDS polling, system metrics
├── http/ # Axum handlers (health, init, snapshot, files, encoding)
├── permissions/ # Path resolution, user lookup, chown
├── rpc/ # Connect RPC services
│ ├── pb.rs # Generated proto types
│ ├── process_*.rs # Process service + handler (PTY, pipe, broadcast)
│ ├── filesystem_*.rs # Filesystem service (stat, list, watch, mkdir, move, remove)
│ └── entry.rs # EntryInfo builder
├── port/ # Port subsystem
│ ├── conn.rs # /proc/net/tcp parser
│ ├── scanner.rs # Periodic TCP port scanner
│ ├── forwarder.rs # socat-based port forwarding
│ └── subsystem.rs # Lifecycle (start/stop/restart)
└── cgroups/ # Cgroup v2 manager (pty/user/socat groups)
```
## Updating the rootfs
After building the static binary, copy it into the rootfs:
```bash
bash scripts/update-debug-rootfs.sh [rootfs_path]
```
Or manually:
```bash
sudo mount -o loop /var/lib/wrenn/images/minimal.ext4 /mnt
sudo cp target/x86_64-unknown-linux-musl/release/envd /mnt/usr/bin/envd
sudo umount /mnt
```

View File

@ -0,0 +1,12 @@
[package]
name = "buffa-types"
version = "0.3.0"
edition = "2024"
publish = false
[dependencies]
buffa = "0.3"
serde = { version = "1", features = ["derive"] }
[build-dependencies]
connectrpc-build = "0.3"

View File

@ -0,0 +1,9 @@
fn main() {
connectrpc_build::Config::new()
.files(&["/usr/include/google/protobuf/timestamp.proto"])
.includes(&["/usr/include"])
.include_file("_types.rs")
.emit_register_fn(false)
.compile()
.unwrap();
}

View File

@ -0,0 +1,6 @@
#![allow(dead_code, non_camel_case_types, unused_imports, clippy::derivable_impls)]
use ::buffa;
use ::serde;
include!(concat!(env!("OUT_DIR"), "/_types.rs"));

11
envd-rs/build.rs Normal file
View File

@ -0,0 +1,11 @@
fn main() {
connectrpc_build::Config::new()
.files(&[
"../proto/envd/process.proto",
"../proto/envd/filesystem.proto",
])
.includes(&["../proto/envd", "/usr/include"])
.include_file("_connectrpc.rs")
.compile()
.unwrap();
}

View File

@ -0,0 +1,3 @@
[toolchain]
channel = "stable"
targets = ["x86_64-unknown-linux-gnu", "x86_64-unknown-linux-musl"]

View File

@ -0,0 +1,56 @@
use std::sync::Arc;
use axum::extract::Request;
use axum::http::StatusCode;
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use serde_json::json;
use crate::auth::token::SecureToken;
const ACCESS_TOKEN_HEADER: &str = "x-access-token";
/// Paths excluded from general token auth.
/// Format: "METHOD/path"
const AUTH_EXCLUDED: &[&str] = &[
"GET/health",
"GET/files",
"POST/files",
"POST/init",
"POST/snapshot/prepare",
];
/// Axum middleware that checks X-Access-Token header.
pub async fn auth_layer(
request: Request,
next: Next,
access_token: Arc<SecureToken>,
) -> Response {
if access_token.is_set() {
let method = request.method().as_str();
let path = request.uri().path();
let key = format!("{method}{path}");
let is_excluded = AUTH_EXCLUDED.iter().any(|p| *p == key);
let header_val = request
.headers()
.get(ACCESS_TOKEN_HEADER)
.and_then(|v| v.to_str().ok())
.unwrap_or("");
if !access_token.equals(header_val) && !is_excluded {
tracing::error!("unauthorized access attempt");
return (
StatusCode::UNAUTHORIZED,
axum::Json(json!({
"code": 401,
"message": "unauthorized access, please provide a valid access token or method signing if supported"
})),
)
.into_response();
}
}
next.run(request).await
}

3
envd-rs/src/auth/mod.rs Normal file
View File

@ -0,0 +1,3 @@
pub mod token;
pub mod signing;
pub mod middleware;

View File

@ -0,0 +1,85 @@
use crate::auth::token::SecureToken;
use crate::crypto;
use zeroize::Zeroize;
pub const READ_OPERATION: &str = "read";
pub const WRITE_OPERATION: &str = "write";
/// Generate a v1 signature: `v1_{sha256_base64(path:operation:username:token[:expiration])}`
pub fn generate_signature(
token: &SecureToken,
path: &str,
username: &str,
operation: &str,
expiration: Option<i64>,
) -> Result<String, &'static str> {
let mut token_bytes = token.bytes().ok_or("access token is not set")?;
let payload = match expiration {
Some(exp) => format!(
"{}:{}:{}:{}:{}",
path,
operation,
username,
String::from_utf8_lossy(&token_bytes),
exp
),
None => format!(
"{}:{}:{}:{}",
path,
operation,
username,
String::from_utf8_lossy(&token_bytes),
),
};
token_bytes.zeroize();
let hash = crypto::sha256::hash_without_prefix(payload.as_bytes());
Ok(format!("v1_{hash}"))
}
/// Validate a request's signing. Returns Ok(()) if valid.
pub fn validate_signing(
token: &SecureToken,
header_token: Option<&str>,
signature: Option<&str>,
signature_expiration: Option<i64>,
username: &str,
path: &str,
operation: &str,
) -> Result<(), String> {
if !token.is_set() {
return Ok(());
}
if let Some(ht) = header_token {
if !ht.is_empty() {
if token.equals(ht) {
return Ok(());
}
return Err("access token present in header but does not match".into());
}
}
let sig = signature.ok_or("missing signature query parameter")?;
let expected = generate_signature(token, path, username, operation, signature_expiration)
.map_err(|e| format!("error generating signing key: {e}"))?;
if expected != sig {
return Err("invalid signature".into());
}
if let Some(exp) = signature_expiration {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs() as i64;
if exp < now {
return Err("signature is already expired".into());
}
}
Ok(())
}

127
envd-rs/src/auth/token.rs Normal file
View File

@ -0,0 +1,127 @@
use std::sync::RwLock;
use subtle::ConstantTimeEq;
use zeroize::Zeroize;
/// Secure token storage with constant-time comparison and zeroize-on-drop.
///
/// Mirrors Go's SecureToken backed by memguard.LockedBuffer.
/// In Rust we rely on `zeroize` for Drop-based zeroing.
pub struct SecureToken {
inner: RwLock<Option<Vec<u8>>>,
}
impl SecureToken {
pub fn new() -> Self {
Self {
inner: RwLock::new(None),
}
}
pub fn set(&self, token: &[u8]) -> Result<(), &'static str> {
if token.is_empty() {
return Err("empty token not allowed");
}
let mut guard = self.inner.write().unwrap();
if let Some(ref mut old) = *guard {
old.zeroize();
}
*guard = Some(token.to_vec());
Ok(())
}
pub fn is_set(&self) -> bool {
let guard = self.inner.read().unwrap();
guard.is_some()
}
/// Constant-time comparison.
pub fn equals(&self, other: &str) -> bool {
let guard = self.inner.read().unwrap();
match guard.as_ref() {
Some(buf) => buf.as_slice().ct_eq(other.as_bytes()).into(),
None => false,
}
}
/// Constant-time comparison with another SecureToken.
pub fn equals_secure(&self, other: &SecureToken) -> bool {
let other_bytes = match other.bytes() {
Some(b) => b,
None => return false,
};
let guard = self.inner.read().unwrap();
let result = match guard.as_ref() {
Some(buf) => buf.as_slice().ct_eq(&other_bytes).into(),
None => false,
};
// other_bytes dropped here, Vec<u8> doesn't auto-zeroize but
// we accept this — same as Go's `defer memguard.WipeBytes(otherBytes)`
result
}
/// Returns a copy of the token bytes (for signature generation).
pub fn bytes(&self) -> Option<Vec<u8>> {
let guard = self.inner.read().unwrap();
guard.as_ref().map(|b| b.clone())
}
/// Transfer token from another SecureToken, clearing the source.
pub fn take_from(&self, src: &SecureToken) {
let taken = {
let mut src_guard = src.inner.write().unwrap();
src_guard.take()
};
let mut guard = self.inner.write().unwrap();
if let Some(ref mut old) = *guard {
old.zeroize();
}
*guard = taken;
}
pub fn destroy(&self) {
let mut guard = self.inner.write().unwrap();
if let Some(ref mut buf) = *guard {
buf.zeroize();
}
*guard = None;
}
}
impl Drop for SecureToken {
fn drop(&mut self) {
if let Ok(mut guard) = self.inner.write() {
if let Some(ref mut buf) = *guard {
buf.zeroize();
}
}
}
}
/// Deserialize from JSON string, matching Go's UnmarshalJSON behavior.
/// Expects a quoted JSON string. Rejects escape sequences.
impl SecureToken {
pub fn from_json_bytes(data: &mut [u8]) -> Result<Self, &'static str> {
if data.len() < 2 || data[0] != b'"' || data[data.len() - 1] != b'"' {
data.zeroize();
return Err("invalid secure token JSON string");
}
let content = &data[1..data.len() - 1];
if content.contains(&b'\\') {
data.zeroize();
return Err("invalid secure token: unexpected escape sequence");
}
if content.is_empty() {
data.zeroize();
return Err("empty token not allowed");
}
let token = Self::new();
token.set(content).map_err(|_| "failed to set token")?;
data.zeroize();
Ok(token)
}
}

View File

@ -0,0 +1,66 @@
use std::collections::HashMap;
use std::fs;
use std::os::unix::io::{OwnedFd, RawFd};
use std::path::PathBuf;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ProcessType {
Pty,
User,
Socat,
}
pub trait CgroupManager: Send + Sync {
fn get_fd(&self, proc_type: ProcessType) -> Option<RawFd>;
}
pub struct Cgroup2Manager {
fds: HashMap<ProcessType, OwnedFd>,
}
impl Cgroup2Manager {
pub fn new(root: &str, configs: &[(ProcessType, &str, &[(&str, &str)])]) -> Result<Self, String> {
let mut fds = HashMap::new();
for (proc_type, sub_path, properties) in configs {
let full_path = PathBuf::from(root).join(sub_path);
fs::create_dir_all(&full_path).map_err(|e| {
format!("failed to create cgroup {}: {e}", full_path.display())
})?;
for (name, value) in *properties {
let prop_path = full_path.join(name);
fs::write(&prop_path, value).map_err(|e| {
format!("failed to write cgroup property {}: {e}", prop_path.display())
})?;
}
let fd = nix::fcntl::open(
&full_path,
nix::fcntl::OFlag::O_RDONLY,
nix::sys::stat::Mode::empty(),
)
.map_err(|e| format!("failed to open cgroup {}: {e}", full_path.display()))?;
fds.insert(*proc_type, fd);
}
Ok(Self { fds })
}
}
impl CgroupManager for Cgroup2Manager {
fn get_fd(&self, proc_type: ProcessType) -> Option<RawFd> {
use std::os::unix::io::AsRawFd;
self.fds.get(&proc_type).map(|fd| fd.as_raw_fd())
}
}
pub struct NoopCgroupManager;
impl CgroupManager for NoopCgroupManager {
fn get_fd(&self, _proc_type: ProcessType) -> Option<RawFd> {
None
}
}

16
envd-rs/src/config.rs Normal file
View File

@ -0,0 +1,16 @@
use std::time::Duration;
pub const DEFAULT_PORT: u16 = 49983;
pub const IDLE_TIMEOUT: Duration = Duration::from_secs(640);
pub const CORS_MAX_AGE: Duration = Duration::from_secs(7200);
pub const PORT_SCANNER_INTERVAL: Duration = Duration::from_millis(1000);
pub const DEFAULT_USER: &str = "root";
pub const WRENN_RUN_DIR: &str = "/run/wrenn";
pub const KILOBYTE: u64 = 1024;
pub const MEGABYTE: u64 = 1024 * KILOBYTE;
pub const MMDS_ADDRESS: &str = "169.254.169.254";
pub const MMDS_POLL_INTERVAL: Duration = Duration::from_millis(50);
pub const MMDS_TOKEN_EXPIRATION_SECS: u64 = 60;
pub const MMDS_ACCESS_TOKEN_CLIENT_TIMEOUT: Duration = Duration::from_secs(10);

View File

@ -0,0 +1,79 @@
use std::collections::HashSet;
use std::sync::Mutex;
/// Tracks active TCP connections for snapshot/restore lifecycle.
///
/// Before snapshot: close idle connections, record active ones.
/// After restore: close all pre-snapshot connections (zombie TCP sockets).
///
/// In Rust/axum, we don't have Go's ConnState callback. Instead we track
/// connections via a tower middleware that registers connection IDs.
/// For the initial implementation, we track by a simple connection counter
/// and rely on axum's graceful shutdown mechanics.
pub struct ConnTracker {
inner: Mutex<ConnTrackerInner>,
}
struct ConnTrackerInner {
active: HashSet<u64>,
pre_snapshot: Option<HashSet<u64>>,
next_id: u64,
keepalives_enabled: bool,
}
impl ConnTracker {
pub fn new() -> Self {
Self {
inner: Mutex::new(ConnTrackerInner {
active: HashSet::new(),
pre_snapshot: None,
next_id: 0,
keepalives_enabled: true,
}),
}
}
pub fn register_connection(&self) -> u64 {
let mut inner = self.inner.lock().unwrap();
let id = inner.next_id;
inner.next_id += 1;
inner.active.insert(id);
id
}
pub fn remove_connection(&self, id: u64) {
let mut inner = self.inner.lock().unwrap();
inner.active.remove(&id);
if let Some(ref mut pre) = inner.pre_snapshot {
pre.remove(&id);
}
}
pub fn prepare_for_snapshot(&self) {
let mut inner = self.inner.lock().unwrap();
inner.keepalives_enabled = false;
inner.pre_snapshot = Some(inner.active.clone());
tracing::info!(
active_connections = inner.active.len(),
"snapshot: recorded pre-snapshot connections, keep-alives disabled"
);
}
pub fn restore_after_snapshot(&self) {
let mut inner = self.inner.lock().unwrap();
if let Some(pre) = inner.pre_snapshot.take() {
let zombie_count = pre.len();
for id in &pre {
inner.active.remove(id);
}
if zombie_count > 0 {
tracing::info!(zombie_count, "restore: closed zombie connections");
}
}
inner.keepalives_enabled = true;
}
pub fn keepalives_enabled(&self) -> bool {
self.inner.lock().unwrap().keepalives_enabled
}
}

View File

@ -0,0 +1,22 @@
use hmac::{Hmac, Mac};
use sha2::Sha256;
type HmacSha256 = Hmac<Sha256>;
pub fn compute(key: &[u8], data: &[u8]) -> String {
let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length");
mac.update(data);
let result = mac.finalize();
hex::encode(result.into_bytes())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_hmac_sha256() {
let result = compute(b"key", b"message");
assert_eq!(result.len(), 64); // SHA-256 hex = 64 chars
}
}

View File

@ -0,0 +1,3 @@
pub mod sha256;
pub mod sha512;
pub mod hmac_sha256;

View File

@ -0,0 +1,33 @@
use base64::Engine;
use base64::engine::general_purpose::STANDARD_NO_PAD;
use sha2::{Digest, Sha256};
pub fn hash(data: &[u8]) -> String {
let h = Sha256::digest(data);
let encoded = STANDARD_NO_PAD.encode(h);
format!("$sha256${encoded}")
}
pub fn hash_without_prefix(data: &[u8]) -> String {
let h = Sha256::digest(data);
STANDARD_NO_PAD.encode(h)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_hash_format() {
let result = hash(b"test");
assert!(result.starts_with("$sha256$"));
assert!(!result.contains('='));
}
#[test]
fn test_hash_without_prefix() {
let result = hash_without_prefix(b"test");
assert!(!result.starts_with("$sha256$"));
assert!(!result.contains('='));
}
}

View File

@ -0,0 +1,24 @@
use sha2::{Digest, Sha512};
pub fn hash_access_token(token: &str) -> String {
let h = Sha512::digest(token.as_bytes());
hex::encode(h)
}
pub fn hash_access_token_bytes(token: &[u8]) -> String {
let h = Sha512::digest(token);
hex::encode(h)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_hash_access_token() {
let h1 = hash_access_token("test");
let h2 = hash_access_token_bytes(b"test");
assert_eq!(h1, h2);
assert_eq!(h1.len(), 128); // SHA-512 hex = 128 chars
}
}

View File

@ -0,0 +1,42 @@
use dashmap::DashMap;
use std::sync::Arc;
#[derive(Clone)]
pub struct Defaults {
pub env_vars: Arc<DashMap<String, String>>,
pub user: String,
pub workdir: Option<String>,
}
impl Defaults {
pub fn new(user: &str) -> Self {
Self {
env_vars: Arc::new(DashMap::new()),
user: user.to_string(),
workdir: None,
}
}
}
pub fn resolve_default_workdir(workdir: &str, default_workdir: Option<&str>) -> String {
if !workdir.is_empty() {
return workdir.to_string();
}
if let Some(dw) = default_workdir {
return dw.to_string();
}
String::new()
}
pub fn resolve_default_username<'a>(
username: Option<&'a str>,
default_username: &'a str,
) -> Result<&'a str, &'static str> {
if let Some(u) = username {
return Ok(u);
}
if !default_username.is_empty() {
return Ok(default_username);
}
Err("username not provided")
}

View File

@ -0,0 +1,73 @@
use std::ffi::CString;
use std::time::{SystemTime, UNIX_EPOCH};
use serde::Serialize;
#[derive(Serialize)]
pub struct Metrics {
pub ts: i64,
pub cpu_count: u32,
pub cpu_used_pct: f32,
pub mem_total_mib: u64,
pub mem_used_mib: u64,
pub mem_total: u64,
pub mem_used: u64,
pub disk_used: u64,
pub disk_total: u64,
}
pub fn get_metrics() -> Result<Metrics, String> {
use sysinfo::System;
let mut sys = System::new();
sys.refresh_memory();
sys.refresh_cpu_all();
std::thread::sleep(std::time::Duration::from_millis(100));
sys.refresh_cpu_all();
let cpu_count = sys.cpus().len() as u32;
let cpu_used_pct = sys.global_cpu_usage();
let cpu_used_pct_rounded = if cpu_used_pct > 0.0 {
(cpu_used_pct * 100.0).round() / 100.0
} else {
0.0
};
let mem_total = sys.total_memory();
let mem_used = sys.used_memory();
let (disk_total, disk_used) = disk_stats("/")?;
let ts = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs() as i64;
Ok(Metrics {
ts,
cpu_count,
cpu_used_pct: cpu_used_pct_rounded,
mem_total_mib: mem_total / 1024 / 1024,
mem_used_mib: mem_used / 1024 / 1024,
mem_total,
mem_used,
disk_used,
disk_total,
})
}
fn disk_stats(path: &str) -> Result<(u64, u64), String> {
let c_path = CString::new(path).unwrap();
let mut stat: libc::statfs = unsafe { std::mem::zeroed() };
let ret = unsafe { libc::statfs(c_path.as_ptr(), &mut stat) };
if ret != 0 {
return Err(format!("statfs failed: {}", std::io::Error::last_os_error()));
}
let block = stat.f_bsize as u64;
let total = stat.f_blocks * block;
let available = stat.f_bavail * block;
Ok((total, total - available))
}

113
envd-rs/src/host/mmds.rs Normal file
View File

@ -0,0 +1,113 @@
use std::sync::Arc;
use std::time::Duration;
use dashmap::DashMap;
use serde::Deserialize;
use tokio_util::sync::CancellationToken;
use crate::config::{MMDS_ADDRESS, MMDS_POLL_INTERVAL, MMDS_TOKEN_EXPIRATION_SECS, WRENN_RUN_DIR};
#[derive(Debug, Clone, Deserialize)]
pub struct MMDSOpts {
#[serde(rename = "instanceID")]
pub sandbox_id: String,
#[serde(rename = "envID")]
pub template_id: String,
#[serde(rename = "address")]
pub logs_collector_address: String,
#[serde(rename = "accessTokenHash", default)]
pub access_token_hash: String,
}
async fn get_mmds_token(client: &reqwest::Client) -> Result<String, String> {
let resp = client
.put(format!("http://{MMDS_ADDRESS}/latest/api/token"))
.header(
"X-metadata-token-ttl-seconds",
MMDS_TOKEN_EXPIRATION_SECS.to_string(),
)
.send()
.await
.map_err(|e| format!("mmds token request failed: {e}"))?;
let token = resp.text().await.map_err(|e| format!("mmds token read: {e}"))?;
if token.is_empty() {
return Err("mmds token is an empty string".into());
}
Ok(token)
}
async fn get_mmds_opts(client: &reqwest::Client, token: &str) -> Result<MMDSOpts, String> {
let resp = client
.get(format!("http://{MMDS_ADDRESS}"))
.header("X-metadata-token", token)
.header("Accept", "application/json")
.send()
.await
.map_err(|e| format!("mmds opts request failed: {e}"))?;
resp.json::<MMDSOpts>()
.await
.map_err(|e| format!("mmds opts parse: {e}"))
}
pub async fn get_access_token_hash() -> Result<String, String> {
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(10))
.no_proxy()
.build()
.map_err(|e| format!("http client: {e}"))?;
let token = get_mmds_token(&client).await?;
let opts = get_mmds_opts(&client, &token).await?;
Ok(opts.access_token_hash)
}
/// Polls MMDS every 50ms until metadata is available.
/// Stores sandbox_id and template_id in env_vars and writes to /run/wrenn/ files.
pub async fn poll_for_opts(
env_vars: Arc<DashMap<String, String>>,
cancel: CancellationToken,
) -> Option<MMDSOpts> {
let client = reqwest::Client::builder()
.no_proxy()
.build()
.ok()?;
let mut interval = tokio::time::interval(MMDS_POLL_INTERVAL);
loop {
tokio::select! {
_ = cancel.cancelled() => {
tracing::warn!("context cancelled while waiting for mmds opts");
return None;
}
_ = interval.tick() => {
let token = match get_mmds_token(&client).await {
Ok(t) => t,
Err(e) => {
tracing::debug!(error = %e, "mmds token poll");
continue;
}
};
let opts = match get_mmds_opts(&client, &token).await {
Ok(o) => o,
Err(e) => {
tracing::debug!(error = %e, "mmds opts poll");
continue;
}
};
env_vars.insert("WRENN_SANDBOX_ID".into(), opts.sandbox_id.clone());
env_vars.insert("WRENN_TEMPLATE_ID".into(), opts.template_id.clone());
let run_dir = std::path::Path::new(WRENN_RUN_DIR);
let _ = std::fs::write(run_dir.join(".WRENN_SANDBOX_ID"), &opts.sandbox_id);
let _ = std::fs::write(run_dir.join(".WRENN_TEMPLATE_ID"), &opts.template_id);
return Some(opts);
}
}
}
}

2
envd-rs/src/host/mod.rs Normal file
View File

@ -0,0 +1,2 @@
pub mod metrics;
pub mod mmds;

View File

@ -0,0 +1,147 @@
use axum::http::Request;
const ENCODING_GZIP: &str = "gzip";
const ENCODING_IDENTITY: &str = "identity";
const ENCODING_WILDCARD: &str = "*";
const SUPPORTED_ENCODINGS: &[&str] = &[ENCODING_GZIP];
struct EncodingWithQuality {
encoding: String,
quality: f64,
}
fn parse_encoding_with_quality(value: &str) -> EncodingWithQuality {
let value = value.trim();
let mut quality = 1.0;
if let Some(idx) = value.find(';') {
let params = &value[idx + 1..];
let enc = value[..idx].trim();
for param in params.split(';') {
let param = param.trim();
if let Some(stripped) = param.strip_prefix("q=").or_else(|| param.strip_prefix("Q=")) {
if let Ok(q) = stripped.parse::<f64>() {
quality = q;
}
}
}
return EncodingWithQuality {
encoding: enc.to_ascii_lowercase(),
quality,
};
}
EncodingWithQuality {
encoding: value.to_ascii_lowercase(),
quality,
}
}
fn parse_accept_encoding_header(header: &str) -> (Vec<EncodingWithQuality>, bool) {
if header.is_empty() {
return (Vec::new(), false);
}
let encodings: Vec<EncodingWithQuality> =
header.split(',').map(|v| parse_encoding_with_quality(v)).collect();
let mut identity_rejected = false;
let mut identity_explicitly_accepted = false;
let mut wildcard_rejected = false;
for eq in &encodings {
match eq.encoding.as_str() {
ENCODING_IDENTITY => {
if eq.quality == 0.0 {
identity_rejected = true;
} else {
identity_explicitly_accepted = true;
}
}
ENCODING_WILDCARD => {
if eq.quality == 0.0 {
wildcard_rejected = true;
}
}
_ => {}
}
}
if wildcard_rejected && !identity_explicitly_accepted {
identity_rejected = true;
}
(encodings, identity_rejected)
}
pub fn is_identity_acceptable<B>(r: &Request<B>) -> bool {
let header = r
.headers()
.get("accept-encoding")
.and_then(|v| v.to_str().ok())
.unwrap_or("");
let (_, rejected) = parse_accept_encoding_header(header);
!rejected
}
pub fn parse_accept_encoding<B>(r: &Request<B>) -> Result<&'static str, String> {
let header = r
.headers()
.get("accept-encoding")
.and_then(|v| v.to_str().ok())
.unwrap_or("");
if header.is_empty() {
return Ok(ENCODING_IDENTITY);
}
let (mut encodings, identity_rejected) = parse_accept_encoding_header(header);
encodings.sort_by(|a, b| b.quality.partial_cmp(&a.quality).unwrap_or(std::cmp::Ordering::Equal));
for eq in &encodings {
if eq.quality == 0.0 {
continue;
}
if eq.encoding == ENCODING_IDENTITY {
return Ok(ENCODING_IDENTITY);
}
if eq.encoding == ENCODING_WILDCARD {
if identity_rejected && !SUPPORTED_ENCODINGS.is_empty() {
return Ok(SUPPORTED_ENCODINGS[0]);
}
return Ok(ENCODING_IDENTITY);
}
if eq.encoding == ENCODING_GZIP {
return Ok(ENCODING_GZIP);
}
}
if !identity_rejected {
return Ok(ENCODING_IDENTITY);
}
Err(format!("no acceptable encoding found, supported: {SUPPORTED_ENCODINGS:?}"))
}
pub fn parse_content_encoding<B>(r: &Request<B>) -> Result<&'static str, String> {
let header = r
.headers()
.get("content-encoding")
.and_then(|v| v.to_str().ok())
.unwrap_or("");
if header.is_empty() {
return Ok(ENCODING_IDENTITY);
}
let encoding = header.trim().to_ascii_lowercase();
if encoding == ENCODING_IDENTITY {
return Ok(ENCODING_IDENTITY);
}
if SUPPORTED_ENCODINGS.contains(&encoding.as_str()) {
return Ok(ENCODING_GZIP);
}
Err(format!("unsupported Content-Encoding: {header}, supported: {SUPPORTED_ENCODINGS:?}"))
}

25
envd-rs/src/http/envs.rs Normal file
View File

@ -0,0 +1,25 @@
use std::collections::HashMap;
use std::sync::Arc;
use axum::Json;
use axum::extract::State;
use axum::http::header;
use axum::response::IntoResponse;
use crate::state::AppState;
pub async fn get_envs(State(state): State<Arc<AppState>>) -> impl IntoResponse {
tracing::debug!("getting env vars");
let envs: HashMap<String, String> = state
.defaults
.env_vars
.iter()
.map(|entry| (entry.key().clone(), entry.value().clone()))
.collect();
(
[(header::CACHE_CONTROL, "no-store")],
Json(envs),
)
}

20
envd-rs/src/http/error.rs Normal file
View File

@ -0,0 +1,20 @@
use axum::Json;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use serde::Serialize;
#[derive(Serialize)]
struct ErrorBody {
code: u16,
message: String,
}
pub fn json_error(status: StatusCode, message: &str) -> impl IntoResponse {
(
status,
Json(ErrorBody {
code: status.as_u16(),
message: message.to_string(),
}),
)
}

443
envd-rs/src/http/files.rs Normal file
View File

@ -0,0 +1,443 @@
use std::io::Write as _;
use std::path::Path;
use std::sync::Arc;
use axum::body::Body;
use axum::extract::{FromRequest, Query, Request, State};
use axum::http::{StatusCode, header};
use axum::response::{IntoResponse, Response};
use serde::{Deserialize, Serialize};
use crate::auth::signing;
use crate::execcontext;
use crate::http::encoding;
use crate::permissions::path::{ensure_dirs, expand_and_resolve};
use crate::permissions::user::lookup_user;
use crate::state::AppState;
const ACCESS_TOKEN_HEADER: &str = "x-access-token";
#[derive(Deserialize)]
pub struct FileParams {
pub path: Option<String>,
pub username: Option<String>,
pub signature: Option<String>,
pub signature_expiration: Option<i64>,
}
#[derive(Serialize)]
struct EntryInfo {
path: String,
name: String,
r#type: &'static str,
}
fn json_error(status: StatusCode, msg: &str) -> Response {
let body = serde_json::json!({ "code": status.as_u16(), "message": msg });
(status, axum::Json(body)).into_response()
}
fn extract_header_token(req: &Request) -> Option<&str> {
req.headers()
.get(ACCESS_TOKEN_HEADER)
.and_then(|v| v.to_str().ok())
}
fn validate_file_signing(
state: &AppState,
header_token: Option<&str>,
params: &FileParams,
path: &str,
operation: &str,
username: &str,
) -> Result<(), String> {
signing::validate_signing(
&state.access_token,
header_token,
params.signature.as_deref(),
params.signature_expiration,
username,
path,
operation,
)
}
/// GET /files — download a file
pub async fn get_files(
State(state): State<Arc<AppState>>,
Query(params): Query<FileParams>,
req: Request,
) -> Response {
let path_str = params.path.as_deref().unwrap_or("");
let header_token = extract_header_token(&req);
let username = match execcontext::resolve_default_username(
params.username.as_deref(),
&state.defaults.user,
) {
Ok(u) => u.to_string(),
Err(e) => return json_error(StatusCode::BAD_REQUEST, e),
};
if let Err(e) = validate_file_signing(
&state,
header_token,
&params,
path_str,
signing::READ_OPERATION,
&username,
) {
return json_error(StatusCode::UNAUTHORIZED, &e);
}
let user = match lookup_user(&username) {
Ok(u) => u,
Err(e) => return json_error(StatusCode::UNAUTHORIZED, &e),
};
let home_dir = format!("/home/{}", user.name);
let resolved = match expand_and_resolve(path_str, &home_dir, state.defaults.workdir.as_deref())
{
Ok(p) => p,
Err(e) => return json_error(StatusCode::BAD_REQUEST, &e),
};
let meta = match std::fs::metadata(&resolved) {
Ok(m) => m,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
return json_error(
StatusCode::NOT_FOUND,
&format!("path '{}' does not exist", resolved),
);
}
Err(e) => {
return json_error(
StatusCode::INTERNAL_SERVER_ERROR,
&format!("error checking path: {e}"),
);
}
};
if meta.is_dir() {
return json_error(
StatusCode::BAD_REQUEST,
&format!("path '{}' is a directory", resolved),
);
}
if !meta.file_type().is_file() {
return json_error(
StatusCode::BAD_REQUEST,
&format!("path '{}' is not a regular file", resolved),
);
}
let accept_enc = match encoding::parse_accept_encoding(&req) {
Ok(e) => e,
Err(e) => return json_error(StatusCode::NOT_ACCEPTABLE, &e),
};
let has_range_or_conditional = req.headers().get("range").is_some()
|| req.headers().get("if-modified-since").is_some()
|| req.headers().get("if-none-match").is_some()
|| req.headers().get("if-range").is_some();
let use_encoding = if has_range_or_conditional {
if !encoding::is_identity_acceptable(&req) {
return json_error(
StatusCode::NOT_ACCEPTABLE,
"identity encoding not acceptable for Range or conditional request",
);
}
"identity"
} else {
accept_enc
};
let file_data = match std::fs::read(&resolved) {
Ok(d) => d,
Err(e) => {
return json_error(
StatusCode::INTERNAL_SERVER_ERROR,
&format!("error reading file: {e}"),
);
}
};
let filename = Path::new(&resolved)
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_default();
let content_disposition = format!("inline; filename=\"{}\"", filename);
let content_type = mime_guess::from_path(&resolved)
.first_raw()
.unwrap_or("application/octet-stream");
if use_encoding == "gzip" {
let mut encoder =
flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
if let Err(e) = encoder.write_all(&file_data) {
return json_error(
StatusCode::INTERNAL_SERVER_ERROR,
&format!("gzip encoding error: {e}"),
);
}
let compressed = match encoder.finish() {
Ok(d) => d,
Err(e) => {
return json_error(
StatusCode::INTERNAL_SERVER_ERROR,
&format!("gzip finish error: {e}"),
);
}
};
return Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, content_type)
.header(header::CONTENT_ENCODING, "gzip")
.header(header::CONTENT_DISPOSITION, content_disposition)
.header(header::VARY, "Accept-Encoding")
.body(Body::from(compressed))
.unwrap();
}
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, content_type)
.header(header::CONTENT_DISPOSITION, content_disposition)
.header(header::VARY, "Accept-Encoding")
.header(header::CONTENT_LENGTH, file_data.len())
.body(Body::from(file_data))
.unwrap()
}
/// POST /files — upload file(s) via multipart
pub async fn post_files(
State(state): State<Arc<AppState>>,
Query(params): Query<FileParams>,
req: Request,
) -> Response {
let path_str = params.path.as_deref().unwrap_or("");
let header_token = extract_header_token(&req);
let username = match execcontext::resolve_default_username(
params.username.as_deref(),
&state.defaults.user,
) {
Ok(u) => u.to_string(),
Err(e) => return json_error(StatusCode::BAD_REQUEST, e),
};
if let Err(e) = validate_file_signing(
&state,
header_token,
&params,
path_str,
signing::WRITE_OPERATION,
&username,
) {
return json_error(StatusCode::UNAUTHORIZED, &e);
}
let user = match lookup_user(&username) {
Ok(u) => u,
Err(e) => return json_error(StatusCode::UNAUTHORIZED, &e),
};
let home_dir = format!("/home/{}", user.name);
let uid = user.uid;
let gid = user.gid;
let content_enc = match encoding::parse_content_encoding(&req) {
Ok(e) => e,
Err(e) => return json_error(StatusCode::BAD_REQUEST, &e),
};
let mut multipart = match axum::extract::Multipart::from_request(req, &()).await {
Ok(m) => m,
Err(e) => {
return json_error(
StatusCode::INTERNAL_SERVER_ERROR,
&format!("error parsing multipart: {e}"),
);
}
};
let mut uploaded: Vec<EntryInfo> = Vec::new();
while let Ok(Some(field)) = multipart.next_field().await {
let field_name = field.name().unwrap_or("").to_string();
if field_name != "file" {
continue;
}
let file_path = if !path_str.is_empty() {
match expand_and_resolve(path_str, &home_dir, state.defaults.workdir.as_deref()) {
Ok(p) => p,
Err(e) => return json_error(StatusCode::BAD_REQUEST, &e),
}
} else {
let fname = field
.file_name()
.unwrap_or("upload")
.to_string();
match expand_and_resolve(&fname, &home_dir, state.defaults.workdir.as_deref()) {
Ok(p) => p,
Err(e) => return json_error(StatusCode::BAD_REQUEST, &e),
}
};
if uploaded.iter().any(|e| e.path == file_path) {
return json_error(
StatusCode::BAD_REQUEST,
&format!("cannot upload multiple files to same path '{}'", file_path),
);
}
let raw_bytes = match field.bytes().await {
Ok(b) => b,
Err(e) => {
return json_error(
StatusCode::INTERNAL_SERVER_ERROR,
&format!("error reading field: {e}"),
);
}
};
let data = if content_enc == "gzip" {
use std::io::Read;
let mut decoder = flate2::read::GzDecoder::new(&raw_bytes[..]);
let mut buf = Vec::new();
match decoder.read_to_end(&mut buf) {
Ok(_) => buf,
Err(e) => {
return json_error(
StatusCode::BAD_REQUEST,
&format!("gzip decompression failed: {e}"),
);
}
}
} else {
raw_bytes.to_vec()
};
if let Err(e) = process_file(&file_path, &data, uid, gid) {
let (status, msg) = e;
return json_error(status, &msg);
}
let name = Path::new(&file_path)
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_default();
uploaded.push(EntryInfo {
path: file_path,
name,
r#type: "file",
});
}
axum::Json(uploaded).into_response()
}
fn process_file(
path: &str,
data: &[u8],
uid: nix::unistd::Uid,
gid: nix::unistd::Gid,
) -> Result<(), (StatusCode, String)> {
let dir = Path::new(path)
.parent()
.map(|p| p.to_string_lossy().to_string())
.unwrap_or_default();
if !dir.is_empty() {
ensure_dirs(&dir, uid, gid).map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("error ensuring directories: {e}"),
)
})?;
}
let can_pre_chown = match std::fs::metadata(path) {
Ok(meta) => {
if meta.is_dir() {
return Err((
StatusCode::BAD_REQUEST,
format!("path is a directory: {path}"),
));
}
true
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => false,
Err(e) => {
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
format!("error getting file info: {e}"),
))
}
};
let mut chowned = false;
if can_pre_chown {
match std::os::unix::fs::chown(path, Some(uid.as_raw()), Some(gid.as_raw())) {
Ok(()) => chowned = true,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
Err(e) => {
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
format!("error changing ownership: {e}"),
))
}
}
}
let mut file = std::fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.mode(0o666)
.open(path)
.map_err(|e| {
if e.raw_os_error() == Some(libc::ENOSPC) {
return (
StatusCode::INSUFFICIENT_STORAGE,
"not enough disk space available".to_string(),
);
}
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("error opening file: {e}"),
)
})?;
if !chowned {
std::os::unix::fs::chown(path, Some(uid.as_raw()), Some(gid.as_raw())).map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("error changing ownership: {e}"),
)
})?;
}
file.write_all(data).map_err(|e| {
if e.raw_os_error() == Some(libc::ENOSPC) {
return (
StatusCode::INSUFFICIENT_STORAGE,
"not enough disk space available".to_string(),
);
}
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("error writing file: {e}"),
)
})?;
Ok(())
}
use std::os::unix::fs::OpenOptionsExt;

View File

@ -0,0 +1,39 @@
use std::sync::Arc;
use std::sync::atomic::Ordering;
use axum::Json;
use axum::extract::State;
use axum::http::header;
use axum::response::IntoResponse;
use serde_json::json;
use crate::state::AppState;
pub async fn get_health(State(state): State<Arc<AppState>>) -> impl IntoResponse {
if state
.needs_restore
.compare_exchange(true, false, Ordering::AcqRel, Ordering::Relaxed)
.is_ok()
{
post_restore_recovery(&state);
}
tracing::trace!("health check");
(
[(header::CACHE_CONTROL, "no-store")],
Json(json!({ "version": state.version })),
)
}
fn post_restore_recovery(state: &AppState) {
tracing::info!("restore: post-restore recovery (no GC needed in Rust)");
state.conn_tracker.restore_after_snapshot();
tracing::info!("restore: zombie connections closed");
if let Some(ref ps) = state.port_subsystem {
ps.restart();
tracing::info!("restore: port subsystem restarted");
}
}

274
envd-rs/src/http/init.rs Normal file
View File

@ -0,0 +1,274 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use axum::Json;
use axum::extract::State;
use axum::http::{StatusCode, header};
use axum::response::IntoResponse;
use serde::Deserialize;
use crate::crypto;
use crate::host::mmds;
use crate::state::AppState;
#[derive(Deserialize, Default)]
#[serde(rename_all = "camelCase")]
pub struct InitRequest {
pub access_token: Option<String>,
pub default_user: Option<String>,
pub default_workdir: Option<String>,
pub env_vars: Option<HashMap<String, String>>,
pub hyperloop_ip: Option<String>,
pub timestamp: Option<String>,
pub volume_mounts: Option<Vec<VolumeMount>>,
}
#[derive(Deserialize)]
pub struct VolumeMount {
pub nfs_target: String,
pub path: String,
}
/// POST /init — called by host agent after boot and after every resume.
pub async fn post_init(
State(state): State<Arc<AppState>>,
body: Option<Json<InitRequest>>,
) -> impl IntoResponse {
let init_req = body.map(|b| b.0).unwrap_or_default();
// Validate access token if provided
if let Some(ref token_str) = init_req.access_token {
if let Err(e) = validate_init_access_token(&state, token_str).await {
tracing::error!(error = %e, "init: access token validation failed");
return (StatusCode::UNAUTHORIZED, e).into_response();
}
}
// Idempotent timestamp check
if let Some(ref ts_str) = init_req.timestamp {
if let Ok(ts) = chrono_parse_to_nanos(ts_str) {
if !state.last_set_time.set_to_greater(ts) {
// Stale request, skip data updates
return trigger_restore_and_respond(&state).await;
}
}
}
// Apply env vars
if let Some(ref vars) = init_req.env_vars {
tracing::debug!(count = vars.len(), "setting env vars");
for (k, v) in vars {
state.defaults.env_vars.insert(k.clone(), v.clone());
}
}
// Set access token
if let Some(ref token_str) = init_req.access_token {
if !token_str.is_empty() {
tracing::debug!("setting access token");
let _ = state.access_token.set(token_str.as_bytes());
} else if state.access_token.is_set() {
tracing::debug!("clearing access token");
state.access_token.destroy();
}
}
// Set default user
if let Some(ref user) = init_req.default_user {
if !user.is_empty() {
tracing::debug!(user = %user, "setting default user");
let mut defaults = state.defaults.clone();
defaults.user = user.clone();
// Note: In Rust we'd need interior mutability for this.
// For now, env_vars (DashMap) handles concurrent access.
// User/workdir mutation deferred to full state refactor.
}
}
// Hyperloop /etc/hosts setup
if let Some(ref ip) = init_req.hyperloop_ip {
let ip = ip.clone();
let env_vars = Arc::clone(&state.defaults.env_vars);
tokio::spawn(async move {
setup_hyperloop(&ip, &env_vars).await;
});
}
// NFS mounts
if let Some(ref mounts) = init_req.volume_mounts {
for mount in mounts {
let target = mount.nfs_target.clone();
let path = mount.path.clone();
tokio::spawn(async move {
setup_nfs(&target, &path).await;
});
}
}
// Re-poll MMDS in background
if state.is_fc {
let env_vars = Arc::clone(&state.defaults.env_vars);
let cancel = tokio_util::sync::CancellationToken::new();
let cancel_clone = cancel.clone();
tokio::spawn(async move {
tokio::time::timeout(std::time::Duration::from_secs(60), async {
mmds::poll_for_opts(env_vars, cancel_clone).await;
})
.await
.ok();
});
}
trigger_restore_and_respond(&state).await
}
async fn trigger_restore_and_respond(state: &AppState) -> axum::response::Response {
// Safety net: if health check's postRestoreRecovery hasn't run yet
if state
.needs_restore
.compare_exchange(true, false, Ordering::AcqRel, Ordering::Relaxed)
.is_ok()
{
post_restore_recovery(state);
}
state.conn_tracker.restore_after_snapshot();
if let Some(ref ps) = state.port_subsystem {
ps.restart();
}
(
StatusCode::NO_CONTENT,
[(header::CACHE_CONTROL, "no-store")],
)
.into_response()
}
fn post_restore_recovery(state: &AppState) {
tracing::info!("restore: post-restore recovery (no GC needed in Rust)");
state.conn_tracker.restore_after_snapshot();
if let Some(ref ps) = state.port_subsystem {
ps.restart();
tracing::info!("restore: port subsystem restarted");
}
}
async fn validate_init_access_token(state: &AppState, request_token: &str) -> Result<(), String> {
// Fast path: matches existing token
if state.access_token.is_set() && !request_token.is_empty() && state.access_token.equals(request_token) {
return Ok(());
}
// Check MMDS hash
if state.is_fc {
if let Ok(mmds_hash) = mmds::get_access_token_hash().await {
if !mmds_hash.is_empty() {
if request_token.is_empty() {
let empty_hash = crypto::sha512::hash_access_token("");
if mmds_hash == empty_hash {
return Ok(());
}
} else {
let token_hash = crypto::sha512::hash_access_token(request_token);
if mmds_hash == token_hash {
return Ok(());
}
}
return Err("access token validation failed".into());
}
}
}
// First-time setup: no existing token and no MMDS
if !state.access_token.is_set() {
return Ok(());
}
if request_token.is_empty() {
return Err("access token reset not authorized".into());
}
Err("access token validation failed".into())
}
async fn setup_hyperloop(address: &str, env_vars: &dashmap::DashMap<String, String>) {
// Write to /etc/hosts: events.wrenn.local → address
let entry = format!("{address} events.wrenn.local\n");
match std::fs::read_to_string("/etc/hosts") {
Ok(contents) => {
let filtered: String = contents
.lines()
.filter(|line| !line.contains("events.wrenn.local"))
.collect::<Vec<_>>()
.join("\n");
let new_contents = format!("{filtered}\n{entry}");
if let Err(e) = std::fs::write("/etc/hosts", new_contents) {
tracing::error!(error = %e, "failed to modify hosts file");
return;
}
}
Err(e) => {
tracing::error!(error = %e, "failed to read hosts file");
return;
}
}
env_vars.insert(
"WRENN_EVENTS_ADDRESS".into(),
format!("http://{address}"),
);
}
async fn setup_nfs(nfs_target: &str, path: &str) {
let mkdir = tokio::process::Command::new("mkdir")
.args(["-p", path])
.output()
.await;
if let Err(e) = mkdir {
tracing::error!(error = %e, path, "nfs: mkdir failed");
return;
}
let mount = tokio::process::Command::new("mount")
.args([
"-v",
"-t",
"nfs",
"-o",
"mountproto=tcp,mountport=2049,proto=tcp,port=2049,nfsvers=3,noacl",
nfs_target,
path,
])
.output()
.await;
match mount {
Ok(output) => {
let stdout = String::from_utf8_lossy(&output.stdout);
let stderr = String::from_utf8_lossy(&output.stderr);
if output.status.success() {
tracing::info!(nfs_target, path, stdout = %stdout, "nfs: mount success");
} else {
tracing::error!(nfs_target, path, stderr = %stderr, "nfs: mount failed");
}
}
Err(e) => {
tracing::error!(error = %e, nfs_target, path, "nfs: mount command failed");
}
}
}
fn chrono_parse_to_nanos(ts: &str) -> Result<i64, ()> {
// Parse RFC3339 timestamp to nanoseconds since epoch
// Simple approach: parse as seconds + fractional
let secs = ts.parse::<f64>().ok();
if let Some(s) = secs {
return Ok((s * 1_000_000_000.0) as i64);
}
// Try RFC3339 format
// For now, fall back to allowing the update
Err(())
}

102
envd-rs/src/http/metrics.rs Normal file
View File

@ -0,0 +1,102 @@
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use axum::Json;
use axum::extract::State;
use axum::http::{StatusCode, header};
use axum::response::IntoResponse;
use serde::Serialize;
use crate::state::AppState;
#[derive(Serialize)]
pub struct Metrics {
ts: i64,
cpu_count: u32,
cpu_used_pct: f32,
mem_total_mib: u64,
mem_used_mib: u64,
mem_total: u64,
mem_used: u64,
disk_used: u64,
disk_total: u64,
}
pub async fn get_metrics(State(_state): State<Arc<AppState>>) -> impl IntoResponse {
tracing::trace!("get metrics");
match collect_metrics() {
Ok(m) => (
StatusCode::OK,
[(header::CACHE_CONTROL, "no-store")],
Json(m),
)
.into_response(),
Err(e) => {
tracing::error!(error = %e, "failed to get metrics");
StatusCode::INTERNAL_SERVER_ERROR.into_response()
}
}
}
fn collect_metrics() -> Result<Metrics, String> {
use sysinfo::System;
let mut sys = System::new();
sys.refresh_memory();
sys.refresh_cpu_all();
// sysinfo needs a small delay for accurate CPU — first call returns 0.
// In a real daemon this would be cached; for now, report instantaneous.
std::thread::sleep(std::time::Duration::from_millis(100));
sys.refresh_cpu_all();
let cpu_count = sys.cpus().len() as u32;
let cpu_used_pct = sys.global_cpu_usage();
let cpu_used_pct_rounded = if cpu_used_pct > 0.0 {
(cpu_used_pct * 100.0).round() / 100.0
} else {
0.0
};
let mem_total = sys.total_memory();
let mem_used = sys.used_memory();
let mem_total_mib = mem_total / 1024 / 1024;
let mem_used_mib = mem_used / 1024 / 1024;
let (disk_total, disk_used) = disk_stats("/").map_err(|e| e.to_string())?;
let ts = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs() as i64;
Ok(Metrics {
ts,
cpu_count,
cpu_used_pct: cpu_used_pct_rounded,
mem_total_mib,
mem_used_mib,
mem_total,
mem_used,
disk_used,
disk_total,
})
}
fn disk_stats(path: &str) -> Result<(u64, u64), nix::Error> {
use std::ffi::CString;
let c_path = CString::new(path).unwrap();
let mut stat: libc::statfs = unsafe { std::mem::zeroed() };
let ret = unsafe { libc::statfs(c_path.as_ptr(), &mut stat) };
if ret != 0 {
return Err(nix::Error::last());
}
let block = stat.f_bsize as u64;
let total = stat.f_blocks * block;
let available = stat.f_bavail * block;
Ok((total, total - available))
}

56
envd-rs/src/http/mod.rs Normal file
View File

@ -0,0 +1,56 @@
pub mod encoding;
pub mod envs;
pub mod error;
pub mod files;
pub mod health;
pub mod init;
pub mod metrics;
pub mod snapshot;
use std::sync::Arc;
use std::time::Duration;
use axum::Router;
use axum::routing::{get, post};
use http::header::{CACHE_CONTROL, HeaderName};
use http::Method;
use tower_http::cors::{AllowHeaders, AllowMethods, AllowOrigin, CorsLayer};
use crate::config::CORS_MAX_AGE;
use crate::state::AppState;
pub fn router(state: Arc<AppState>) -> Router {
let cors = CorsLayer::new()
.allow_origin(AllowOrigin::any())
.allow_methods(AllowMethods::list([
Method::HEAD,
Method::GET,
Method::POST,
Method::PUT,
Method::PATCH,
Method::DELETE,
]))
.allow_headers(AllowHeaders::any())
.expose_headers([
HeaderName::from_static("location"),
CACHE_CONTROL,
HeaderName::from_static("x-content-type-options"),
HeaderName::from_static("connect-content-encoding"),
HeaderName::from_static("connect-protocol-version"),
HeaderName::from_static("grpc-encoding"),
HeaderName::from_static("grpc-message"),
HeaderName::from_static("grpc-status"),
HeaderName::from_static("grpc-status-details-bin"),
])
.max_age(Duration::from_secs(CORS_MAX_AGE.as_secs()));
Router::new()
.route("/health", get(health::get_health))
.route("/metrics", get(metrics::get_metrics))
.route("/envs", get(envs::get_envs))
.route("/init", post(init::post_init))
.route("/snapshot/prepare", post(snapshot::post_snapshot_prepare))
.route("/files", get(files::get_files).post(files::post_files))
.layer(cors)
.with_state(state)
}

View File

@ -0,0 +1,32 @@
use std::sync::Arc;
use std::sync::atomic::Ordering;
use axum::extract::State;
use axum::http::{StatusCode, header};
use axum::response::IntoResponse;
use crate::state::AppState;
/// POST /snapshot/prepare — quiesce subsystems before Firecracker snapshot.
///
/// In Rust there is no GC dance. We just:
/// 1. Stop port subsystem
/// 2. Close idle connections via conntracker
/// 3. Set needs_restore flag
pub async fn post_snapshot_prepare(State(state): State<Arc<AppState>>) -> impl IntoResponse {
if let Some(ref ps) = state.port_subsystem {
ps.stop();
tracing::info!("snapshot/prepare: port subsystem stopped");
}
state.conn_tracker.prepare_for_snapshot();
tracing::info!("snapshot/prepare: connections prepared");
state.needs_restore.store(true, Ordering::Release);
tracing::info!("snapshot/prepare: ready for freeze");
(
StatusCode::NO_CONTENT,
[(header::CACHE_CONTROL, "no-store")],
)
}

17
envd-rs/src/logging.rs Normal file
View File

@ -0,0 +1,17 @@
use tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt};
pub fn init(json: bool) {
let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
if json {
tracing_subscriber::registry()
.with(filter)
.with(fmt::layer().json().flatten_event(true))
.init();
} else {
tracing_subscriber::registry()
.with(filter)
.with(fmt::layer())
.init();
}
}

224
envd-rs/src/main.rs Normal file
View File

@ -0,0 +1,224 @@
#![allow(dead_code)]
mod auth;
mod cgroups;
mod config;
mod conntracker;
mod crypto;
mod execcontext;
mod host;
mod http;
mod logging;
mod permissions;
mod port;
mod rpc;
mod state;
mod util;
use std::fs;
use std::net::SocketAddr;
use std::path::Path;
use std::sync::Arc;
use clap::Parser;
use tokio::net::TcpListener;
use tokio_util::sync::CancellationToken;
use config::{DEFAULT_PORT, DEFAULT_USER, WRENN_RUN_DIR};
use execcontext::Defaults;
use port::subsystem::PortSubsystem;
use state::AppState;
const VERSION: &str = env!("CARGO_PKG_VERSION");
const COMMIT: &str = {
match option_env!("ENVD_COMMIT") {
Some(c) => c,
None => "unknown",
}
};
#[derive(Parser)]
#[command(name = "envd", about = "Wrenn guest agent daemon")]
struct Cli {
#[arg(long, default_value_t = DEFAULT_PORT)]
port: u16,
#[arg(long = "isnotfc", default_value_t = false)]
is_not_fc: bool,
#[arg(long)]
version: bool,
#[arg(long)]
commit: bool,
#[arg(long = "cmd", default_value = "")]
start_cmd: String,
#[arg(long = "cgroup-root", default_value = "/sys/fs/cgroup")]
cgroup_root: String,
}
#[tokio::main]
async fn main() {
let cli = Cli::parse();
if cli.version {
println!("{VERSION}");
return;
}
if cli.commit {
println!("{COMMIT}");
return;
}
let use_json = !cli.is_not_fc;
logging::init(use_json);
if let Err(e) = fs::create_dir_all(WRENN_RUN_DIR) {
tracing::error!(error = %e, "failed to create wrenn run directory");
}
let defaults = Defaults::new(DEFAULT_USER);
let is_fc_str = if cli.is_not_fc { "false" } else { "true" };
defaults
.env_vars
.insert("WRENN_SANDBOX".into(), is_fc_str.into());
let wrenn_sandbox_path = Path::new(WRENN_RUN_DIR).join(".WRENN_SANDBOX");
if let Err(e) = fs::write(&wrenn_sandbox_path, is_fc_str.as_bytes()) {
tracing::error!(error = %e, "failed to write sandbox file");
}
let cancel = CancellationToken::new();
// MMDS polling (only in FC mode)
if !cli.is_not_fc {
let env_vars = Arc::clone(&defaults.env_vars);
let cancel_clone = cancel.clone();
tokio::spawn(async move {
host::mmds::poll_for_opts(env_vars, cancel_clone).await;
});
}
// Cgroup manager
let cgroup_manager: Arc<dyn cgroups::CgroupManager> =
match cgroups::Cgroup2Manager::new(
&cli.cgroup_root,
&[
(
cgroups::ProcessType::Pty,
"wrenn/pty",
&[] as &[(&str, &str)],
),
(
cgroups::ProcessType::User,
"wrenn/user",
&[] as &[(&str, &str)],
),
(
cgroups::ProcessType::Socat,
"wrenn/socat",
&[] as &[(&str, &str)],
),
],
) {
Ok(m) => {
tracing::info!("cgroup2 manager initialized");
Arc::new(m)
}
Err(e) => {
tracing::warn!(error = %e, "cgroup2 init failed, using noop");
Arc::new(cgroups::NoopCgroupManager)
}
};
// Port subsystem
let port_subsystem = Arc::new(PortSubsystem::new(Arc::clone(&cgroup_manager)));
port_subsystem.start();
tracing::info!("port subsystem started");
let state = AppState::new(
defaults,
VERSION.to_string(),
COMMIT.to_string(),
!cli.is_not_fc,
Some(Arc::clone(&port_subsystem)),
);
// RPC services (Connect protocol — serves Connect + gRPC + gRPC-Web on same port)
let connect_router = rpc::rpc_router(Arc::clone(&state));
let app = http::router(Arc::clone(&state))
.fallback_service(connect_router.into_axum_service());
// --cmd: spawn initial process if specified
if !cli.start_cmd.is_empty() {
let cmd = cli.start_cmd.clone();
let state_clone = Arc::clone(&state);
tokio::spawn(async move {
spawn_initial_command(&cmd, &state_clone);
});
}
let addr = SocketAddr::from(([0, 0, 0, 0], cli.port));
tracing::info!(port = cli.port, version = VERSION, commit = COMMIT, "envd starting");
let listener = TcpListener::bind(addr).await.expect("failed to bind");
let graceful = axum::serve(listener, app).with_graceful_shutdown(async move {
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
.expect("failed to register SIGTERM")
.recv()
.await;
tracing::info!("SIGTERM received, shutting down");
});
if let Err(e) = graceful.await {
tracing::error!(error = %e, "server error");
}
port_subsystem.stop();
cancel.cancel();
}
fn spawn_initial_command(cmd: &str, state: &AppState) {
use crate::permissions::user::lookup_user;
use crate::rpc::process_handler;
use std::collections::HashMap;
let user = match lookup_user(&state.defaults.user) {
Ok(u) => u,
Err(e) => {
tracing::error!(error = %e, "cmd: failed to lookup user");
return;
}
};
let home = format!("/home/{}", user.name);
let cwd = state
.defaults
.workdir
.as_deref()
.unwrap_or(&home);
match process_handler::spawn_process(
cmd,
&[],
&HashMap::new(),
cwd,
None,
false,
Some("init-cmd".to_string()),
&user,
&state.defaults.env_vars,
) {
Ok(handle) => {
tracing::info!(pid = handle.pid, cmd, "initial command spawned");
}
Err(e) => {
tracing::error!(error = %e, cmd, "failed to spawn initial command");
}
}
}

View File

@ -0,0 +1,2 @@
pub mod user;
pub mod path;

View File

@ -0,0 +1,72 @@
use std::fs;
use std::os::unix::fs::chown;
use std::path::{Path, PathBuf};
use nix::unistd::{Gid, Uid};
fn expand_tilde(path: &str, home_dir: &str) -> Result<String, String> {
if path.is_empty() || !path.starts_with('~') {
return Ok(path.to_string());
}
if path.len() > 1 && path.as_bytes()[1] != b'/' && path.as_bytes()[1] != b'\\' {
return Err("cannot expand user-specific home dir".into());
}
Ok(format!("{}{}", home_dir, &path[1..]))
}
pub fn expand_and_resolve(
path: &str,
home_dir: &str,
default_path: Option<&str>,
) -> Result<String, String> {
let path = if path.is_empty() {
default_path.unwrap_or("").to_string()
} else {
path.to_string()
};
let path = expand_tilde(&path, home_dir)?;
if Path::new(&path).is_absolute() {
return Ok(path);
}
let joined = PathBuf::from(home_dir).join(&path);
joined
.canonicalize()
.or_else(|_| Ok(joined))
.map(|p| p.to_string_lossy().to_string())
}
pub fn ensure_dirs(path: &str, uid: Uid, gid: Gid) -> Result<(), String> {
let path = Path::new(path);
let mut current = PathBuf::new();
for component in path.components() {
current.push(component);
let current_str = current.to_string_lossy();
if current_str == "/" {
continue;
}
match fs::metadata(&current) {
Ok(meta) => {
if !meta.is_dir() {
return Err(format!("path is a file: {current_str}"));
}
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
fs::create_dir(&current)
.map_err(|e| format!("failed to create directory {current_str}: {e}"))?;
chown(&current, Some(uid.as_raw()), Some(gid.as_raw()))
.map_err(|e| format!("failed to chown directory {current_str}: {e}"))?;
}
Err(e) => {
return Err(format!("failed to stat directory {current_str}: {e}"));
}
}
}
Ok(())
}

View File

@ -0,0 +1,32 @@
use nix::unistd::{Gid, Group, Uid, User};
pub fn lookup_user(username: &str) -> Result<User, String> {
User::from_name(username)
.map_err(|e| format!("error looking up user '{username}': {e}"))?
.ok_or_else(|| format!("user '{username}' not found"))
}
pub fn get_uid_gid(user: &User) -> (Uid, Gid) {
(user.uid, user.gid)
}
pub fn get_user_groups(user: &User) -> Vec<Gid> {
let c_name = std::ffi::CString::new(user.name.as_str()).unwrap();
nix::unistd::getgrouplist(&c_name, user.gid).unwrap_or_default()
}
pub fn lookup_username_by_uid(uid: Uid) -> String {
User::from_uid(uid)
.ok()
.flatten()
.map(|u| u.name)
.unwrap_or_else(|| uid.to_string())
}
pub fn lookup_groupname_by_gid(gid: Gid) -> String {
Group::from_gid(gid)
.ok()
.flatten()
.map(|g| g.name)
.unwrap_or_else(|| gid.to_string())
}

112
envd-rs/src/port/conn.rs Normal file
View File

@ -0,0 +1,112 @@
use std::io::{self, BufRead};
#[derive(Debug, Clone)]
pub struct ConnStat {
pub local_ip: String,
pub local_port: u32,
pub status: String,
pub family: u32,
pub inode: u64,
}
fn tcp_state_name(hex: &str) -> &'static str {
match hex {
"01" => "ESTABLISHED",
"02" => "SYN_SENT",
"03" => "SYN_RECV",
"04" => "FIN_WAIT1",
"05" => "FIN_WAIT2",
"06" => "TIME_WAIT",
"07" => "CLOSE",
"08" => "CLOSE_WAIT",
"09" => "LAST_ACK",
"0A" => "LISTEN",
"0B" => "CLOSING",
_ => "UNKNOWN",
}
}
pub fn read_tcp_connections() -> Vec<ConnStat> {
let mut conns = Vec::new();
if let Ok(c) = parse_proc_net_tcp("/proc/net/tcp", libc::AF_INET as u32) {
conns.extend(c);
}
if let Ok(c) = parse_proc_net_tcp("/proc/net/tcp6", libc::AF_INET6 as u32) {
conns.extend(c);
}
conns
}
fn parse_proc_net_tcp(path: &str, family: u32) -> io::Result<Vec<ConnStat>> {
let file = std::fs::File::open(path)?;
let reader = io::BufReader::new(file);
let mut conns = Vec::new();
let mut first = true;
for line in reader.lines() {
let line = line?;
if first {
first = false;
continue;
}
let line = line.trim().to_string();
if line.is_empty() {
continue;
}
let fields: Vec<&str> = line.split_whitespace().collect();
if fields.len() < 10 {
continue;
}
let (ip, port) = match parse_hex_addr(fields[1], family) {
Some(v) => v,
None => continue,
};
let state = tcp_state_name(fields[3]);
let inode: u64 = match fields[9].parse() {
Ok(v) => v,
Err(_) => continue,
};
conns.push(ConnStat {
local_ip: ip,
local_port: port,
status: state.to_string(),
family,
inode,
});
}
Ok(conns)
}
fn parse_hex_addr(s: &str, family: u32) -> Option<(String, u32)> {
let (ip_hex, port_hex) = s.split_once(':')?;
let port = u32::from_str_radix(port_hex, 16).ok()?;
let ip_bytes = hex::decode(ip_hex).ok()?;
let ip_str = if family == libc::AF_INET as u32 {
if ip_bytes.len() != 4 {
return None;
}
format!("{}.{}.{}.{}", ip_bytes[3], ip_bytes[2], ip_bytes[1], ip_bytes[0])
} else {
if ip_bytes.len() != 16 {
return None;
}
let mut octets = [0u8; 16];
for i in 0..4 {
octets[i * 4] = ip_bytes[i * 4 + 3];
octets[i * 4 + 1] = ip_bytes[i * 4 + 2];
octets[i * 4 + 2] = ip_bytes[i * 4 + 1];
octets[i * 4 + 3] = ip_bytes[i * 4];
}
let addr = std::net::Ipv6Addr::from(octets);
addr.to_string()
};
Some((ip_str, port))
}

View File

@ -0,0 +1,181 @@
use std::collections::HashMap;
use std::os::unix::process::CommandExt;
use std::process::Command;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use crate::cgroups::{CgroupManager, ProcessType};
use super::conn::ConnStat;
const DEFAULT_GATEWAY_IP: &str = "169.254.0.21";
#[derive(PartialEq)]
enum PortState {
Forward,
Delete,
}
struct PortToForward {
pid: Option<u32>,
inode: u64,
family: u32,
state: PortState,
port: u32,
}
fn family_to_ip_version(family: u32) -> u32 {
if family == libc::AF_INET as u32 {
4
} else if family == libc::AF_INET6 as u32 {
6
} else {
0
}
}
pub struct Forwarder {
cgroup_manager: Arc<dyn CgroupManager>,
ports: HashMap<String, PortToForward>,
source_ip: String,
}
impl Forwarder {
pub fn new(cgroup_manager: Arc<dyn CgroupManager>) -> Self {
Self {
cgroup_manager,
ports: HashMap::new(),
source_ip: DEFAULT_GATEWAY_IP.to_string(),
}
}
pub async fn start_forwarding(
&mut self,
mut rx: mpsc::Receiver<Vec<ConnStat>>,
cancel: CancellationToken,
) {
loop {
tokio::select! {
_ = cancel.cancelled() => {
self.stop_all();
return;
}
msg = rx.recv() => {
match msg {
Some(conns) => self.process_scan(conns),
None => {
self.stop_all();
return;
}
}
}
}
}
}
fn process_scan(&mut self, conns: Vec<ConnStat>) {
for ptf in self.ports.values_mut() {
ptf.state = PortState::Delete;
}
for conn in &conns {
let key = format!("{}-{}", conn.inode, conn.local_port);
if let Some(ptf) = self.ports.get_mut(&key) {
ptf.state = PortState::Forward;
} else {
tracing::debug!(
ip = %conn.local_ip,
port = conn.local_port,
family = family_to_ip_version(conn.family),
"detected new port on localhost"
);
let mut ptf = PortToForward {
pid: None,
inode: conn.inode,
family: family_to_ip_version(conn.family),
state: PortState::Forward,
port: conn.local_port,
};
self.start_port_forwarding(&mut ptf);
self.ports.insert(key, ptf);
}
}
let to_stop: Vec<String> = self
.ports
.iter()
.filter(|(_, v)| v.state == PortState::Delete)
.map(|(k, _)| k.clone())
.collect();
for key in to_stop {
if let Some(ptf) = self.ports.get(&key) {
stop_port_forwarding(ptf);
}
self.ports.remove(&key);
}
}
fn start_port_forwarding(&self, ptf: &mut PortToForward) {
let listen_arg = format!(
"TCP4-LISTEN:{},bind={},reuseaddr,fork",
ptf.port, self.source_ip
);
let connect_arg = format!("TCP{}:localhost:{}", ptf.family, ptf.port);
let mut cmd = Command::new("socat");
cmd.args(["-d", "-d", "-d", &listen_arg, &connect_arg]);
unsafe {
let cgroup_fd = self.cgroup_manager.get_fd(ProcessType::Socat);
cmd.pre_exec(move || {
libc::setpgid(0, 0);
if let Some(fd) = cgroup_fd {
let pid_str = format!("{}", libc::getpid());
let tasks_path = format!("/proc/self/fd/{}/cgroup.procs", fd);
let _ = std::fs::write(&tasks_path, pid_str.as_bytes());
}
Ok(())
});
}
tracing::debug!(
port = ptf.port,
inode = ptf.inode,
family = ptf.family,
source_ip = %self.source_ip,
"starting port forwarding"
);
match cmd.spawn() {
Ok(child) => {
ptf.pid = Some(child.id());
std::thread::spawn(move || {
let mut child = child;
let _ = child.wait();
});
}
Err(e) => {
tracing::error!(error = %e, port = ptf.port, "failed to start socat");
}
}
}
fn stop_all(&mut self) {
for ptf in self.ports.values() {
stop_port_forwarding(ptf);
}
self.ports.clear();
}
}
fn stop_port_forwarding(ptf: &PortToForward) {
if let Some(pid) = ptf.pid {
tracing::debug!(port = ptf.port, pid, "stopping port forwarding");
unsafe {
libc::kill(-(pid as i32), libc::SIGKILL);
}
}
}

4
envd-rs/src/port/mod.rs Normal file
View File

@ -0,0 +1,4 @@
pub mod conn;
pub mod forwarder;
pub mod scanner;
pub mod subsystem;

View File

@ -0,0 +1,79 @@
use std::sync::{Arc, RwLock};
use std::time::Duration;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use super::conn::{ConnStat, read_tcp_connections};
pub struct ScannerFilter {
pub ips: Vec<String>,
pub state: String,
}
impl ScannerFilter {
pub fn matches(&self, conn: &ConnStat) -> bool {
if self.state.is_empty() && self.ips.is_empty() {
return false;
}
self.ips.contains(&conn.local_ip) && self.state == conn.status
}
}
pub struct ScannerSubscriber {
pub tx: mpsc::Sender<Vec<ConnStat>>,
pub filter: Option<ScannerFilter>,
}
pub struct Scanner {
period: Duration,
subs: RwLock<Vec<(String, Arc<ScannerSubscriber>)>>,
}
impl Scanner {
pub fn new(period: Duration) -> Self {
Self {
period,
subs: RwLock::new(Vec::new()),
}
}
pub fn add_subscriber(
&self,
id: &str,
filter: Option<ScannerFilter>,
) -> mpsc::Receiver<Vec<ConnStat>> {
let (tx, rx) = mpsc::channel(4);
let sub = Arc::new(ScannerSubscriber { tx, filter });
let mut subs = self.subs.write().unwrap();
subs.push((id.to_string(), sub));
rx
}
pub fn remove_subscriber(&self, id: &str) {
let mut subs = self.subs.write().unwrap();
subs.retain(|(sid, _)| sid != id);
}
pub async fn scan_and_broadcast(&self, cancel: CancellationToken) {
loop {
let conns = read_tcp_connections();
{
let subs = self.subs.read().unwrap();
for (_, sub) in subs.iter() {
let payload = match &sub.filter {
Some(f) => conns.iter().filter(|c| f.matches(c)).cloned().collect(),
None => conns.clone(),
};
let _ = sub.tx.try_send(payload);
}
}
tokio::select! {
_ = cancel.cancelled() => return,
_ = tokio::time::sleep(self.period) => {}
}
}
}
}

View File

@ -0,0 +1,78 @@
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use crate::cgroups::CgroupManager;
use crate::config::PORT_SCANNER_INTERVAL;
use super::forwarder::Forwarder;
use super::scanner::{Scanner, ScannerFilter};
pub struct PortSubsystem {
cgroup_manager: Arc<dyn CgroupManager>,
cancel: std::sync::Mutex<Option<CancellationToken>>,
}
impl PortSubsystem {
pub fn new(cgroup_manager: Arc<dyn CgroupManager>) -> Self {
Self {
cgroup_manager,
cancel: std::sync::Mutex::new(None),
}
}
pub fn start(&self) {
let mut guard = self.cancel.lock().unwrap();
if guard.is_some() {
return;
}
let cancel = CancellationToken::new();
*guard = Some(cancel.clone());
drop(guard);
let cgroup_manager = Arc::clone(&self.cgroup_manager);
let cancel_scanner = cancel.clone();
let cancel_forwarder = cancel.clone();
tokio::spawn(async move {
let scanner = Arc::new(Scanner::new(PORT_SCANNER_INTERVAL));
let rx = scanner.add_subscriber(
"port-forwarder",
Some(ScannerFilter {
ips: vec![
"127.0.0.1".to_string(),
"localhost".to_string(),
"::1".to_string(),
],
state: "LISTEN".to_string(),
}),
);
let scanner_clone = Arc::clone(&scanner);
let scanner_handle = tokio::spawn(async move {
scanner_clone.scan_and_broadcast(cancel_scanner).await;
});
let forwarder_handle = tokio::spawn(async move {
let mut forwarder = Forwarder::new(cgroup_manager);
forwarder.start_forwarding(rx, cancel_forwarder).await;
});
let _ = tokio::join!(scanner_handle, forwarder_handle);
});
}
pub fn stop(&self) {
let mut guard = self.cancel.lock().unwrap();
if let Some(cancel) = guard.take() {
cancel.cancel();
}
}
pub fn restart(&self) {
self.stop();
self.start();
}
}

142
envd-rs/src/rpc/entry.rs Normal file
View File

@ -0,0 +1,142 @@
use std::os::unix::fs::MetadataExt;
use std::path::Path;
use connectrpc::{ConnectError, ErrorCode};
use crate::permissions::user::{lookup_groupname_by_gid, lookup_username_by_uid};
use crate::rpc::pb::filesystem::{EntryInfo, FileType};
use nix::unistd::{Gid, Uid};
const NFS_SUPER_MAGIC: i64 = 0x6969;
const CIFS_MAGIC: i64 = 0xFF534D42;
const SMB_SUPER_MAGIC: i64 = 0x517B;
const SMB2_MAGIC_NUMBER: i64 = 0xFE534D42;
const FUSE_SUPER_MAGIC: i64 = 0x65735546;
pub fn is_network_mount(path: &str) -> Result<bool, String> {
let c_path = std::ffi::CString::new(path).map_err(|e| e.to_string())?;
let mut stat: libc::statfs = unsafe { std::mem::zeroed() };
let ret = unsafe { libc::statfs(c_path.as_ptr(), &mut stat) };
if ret != 0 {
return Err(format!(
"statfs {path}: {}",
std::io::Error::last_os_error()
));
}
let fs_type = stat.f_type as i64;
Ok(matches!(
fs_type,
NFS_SUPER_MAGIC | CIFS_MAGIC | SMB_SUPER_MAGIC | SMB2_MAGIC_NUMBER | FUSE_SUPER_MAGIC
))
}
pub fn build_entry_info(path: &str) -> Result<EntryInfo, ConnectError> {
let p = Path::new(path);
let lstat = std::fs::symlink_metadata(p).map_err(|e| {
if e.kind() == std::io::ErrorKind::NotFound {
ConnectError::new(ErrorCode::NotFound, format!("file not found: {e}"))
} else {
ConnectError::new(ErrorCode::Internal, format!("error getting file info: {e}"))
}
})?;
let is_symlink = lstat.file_type().is_symlink();
let (file_type, mode, symlink_target) = if is_symlink {
let target = std::fs::canonicalize(p)
.map(|t| t.to_string_lossy().to_string())
.unwrap_or_else(|_| path.to_string());
let target_type = match std::fs::metadata(p) {
Ok(meta) => meta_to_file_type(&meta),
Err(_) => FileType::FILE_TYPE_UNSPECIFIED,
};
let target_mode = std::fs::metadata(p)
.map(|m| m.mode() & 0o7777)
.unwrap_or(0);
(target_type, target_mode, Some(target))
} else {
let ft = meta_to_file_type(&lstat);
let mode = lstat.mode() & 0o7777;
(ft, mode, None)
};
let uid = lstat.uid();
let gid = lstat.gid();
let owner = lookup_username_by_uid(Uid::from_raw(uid));
let group = lookup_groupname_by_gid(Gid::from_raw(gid));
let modified_time = {
let mtime_sec = lstat.mtime();
let mtime_nsec = lstat.mtime_nsec() as i32;
if mtime_sec == 0 && mtime_nsec == 0 {
None
} else {
Some(buffa_types::google::protobuf::Timestamp {
seconds: mtime_sec,
nanos: mtime_nsec,
..Default::default()
})
}
};
let name = p
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_default();
let permissions = format_permissions(lstat.mode());
Ok(EntryInfo {
name,
r#type: buffa::EnumValue::Known(file_type),
path: path.to_string(),
size: lstat.len() as i64,
mode,
permissions,
owner,
group,
modified_time: modified_time.into(),
symlink_target: symlink_target,
..Default::default()
})
}
fn meta_to_file_type(meta: &std::fs::Metadata) -> FileType {
if meta.is_file() {
FileType::FILE_TYPE_FILE
} else if meta.is_dir() {
FileType::FILE_TYPE_DIRECTORY
} else if meta.file_type().is_symlink() {
FileType::FILE_TYPE_SYMLINK
} else {
FileType::FILE_TYPE_UNSPECIFIED
}
}
fn format_permissions(mode: u32) -> String {
let file_type = match mode & libc::S_IFMT {
libc::S_IFDIR => 'd',
libc::S_IFLNK => 'L',
libc::S_IFREG => '-',
libc::S_IFBLK => 'b',
libc::S_IFCHR => 'c',
libc::S_IFIFO => 'p',
libc::S_IFSOCK => 'S',
_ => '?',
};
let perms = mode & 0o777;
let mut s = String::with_capacity(10);
s.push(file_type);
for shift in [6, 3, 0] {
let bits = (perms >> shift) & 7;
s.push(if bits & 4 != 0 { 'r' } else { '-' });
s.push(if bits & 2 != 0 { 'w' } else { '-' });
s.push(if bits & 1 != 0 { 'x' } else { '-' });
}
s
}

View File

@ -0,0 +1,402 @@
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use connectrpc::{ConnectError, Context, ErrorCode};
use dashmap::DashMap;
use futures::Stream;
use crate::permissions::path::{ensure_dirs, expand_and_resolve};
use crate::permissions::user::lookup_user;
use crate::rpc::entry::build_entry_info;
use crate::rpc::pb::filesystem::*;
use crate::state::AppState;
pub struct FilesystemServiceImpl {
state: Arc<AppState>,
watchers: DashMap<String, WatcherHandle>,
}
struct WatcherHandle {
events: Arc<Mutex<Vec<FilesystemEvent>>>,
_watcher: notify::RecommendedWatcher,
}
impl FilesystemServiceImpl {
pub fn new(state: Arc<AppState>) -> Self {
Self {
state,
watchers: DashMap::new(),
}
}
fn resolve_path(&self, path: &str, ctx: &Context) -> Result<String, ConnectError> {
let username = extract_username(ctx).unwrap_or_else(|| self.state.defaults.user.clone());
let user = lookup_user(&username).map_err(|e| {
ConnectError::new(ErrorCode::Unauthenticated, format!("invalid user: {e}"))
})?;
let home_dir = format!("/home/{}", user.name);
let default_workdir = self.state.defaults.workdir.as_deref();
expand_and_resolve(path, &home_dir, default_workdir)
.map_err(|e| ConnectError::new(ErrorCode::InvalidArgument, e))
}
}
fn extract_username(ctx: &Context) -> Option<String> {
ctx.extensions.get::<AuthUser>().map(|u| u.0.clone())
}
#[derive(Clone)]
pub struct AuthUser(pub String);
impl Filesystem for FilesystemServiceImpl {
async fn stat(
&self,
ctx: Context,
request: buffa::view::OwnedView<StatRequestView<'static>>,
) -> Result<(StatResponse, Context), ConnectError> {
let path = self.resolve_path(request.path, &ctx)?;
let entry = build_entry_info(&path)?;
Ok((
StatResponse {
entry: entry.into(),
..Default::default()
},
ctx,
))
}
async fn make_dir(
&self,
ctx: Context,
request: buffa::view::OwnedView<MakeDirRequestView<'static>>,
) -> Result<(MakeDirResponse, Context), ConnectError> {
let path = self.resolve_path(request.path, &ctx)?;
match std::fs::metadata(&path) {
Ok(meta) => {
if meta.is_dir() {
return Err(ConnectError::new(
ErrorCode::AlreadyExists,
format!("directory already exists: {path}"),
));
}
return Err(ConnectError::new(
ErrorCode::InvalidArgument,
format!("path exists but is not a directory: {path}"),
));
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
Err(e) => {
return Err(ConnectError::new(
ErrorCode::Internal,
format!("error getting file info: {e}"),
));
}
}
let username = extract_username(&ctx).unwrap_or_else(|| self.state.defaults.user.clone());
let user =
lookup_user(&username).map_err(|e| ConnectError::new(ErrorCode::Internal, e))?;
ensure_dirs(&path, user.uid, user.gid)
.map_err(|e| ConnectError::new(ErrorCode::Internal, e))?;
let entry = build_entry_info(&path)?;
Ok((
MakeDirResponse {
entry: entry.into(),
..Default::default()
},
ctx,
))
}
async fn r#move(
&self,
ctx: Context,
request: buffa::view::OwnedView<MoveRequestView<'static>>,
) -> Result<(MoveResponse, Context), ConnectError> {
let source = self.resolve_path(request.source, &ctx)?;
let destination = self.resolve_path(request.destination, &ctx)?;
let username = extract_username(&ctx).unwrap_or_else(|| self.state.defaults.user.clone());
let user =
lookup_user(&username).map_err(|e| ConnectError::new(ErrorCode::Internal, e))?;
if let Some(parent) = Path::new(&destination).parent() {
ensure_dirs(&parent.to_string_lossy(), user.uid, user.gid)
.map_err(|e| ConnectError::new(ErrorCode::Internal, e))?;
}
std::fs::rename(&source, &destination).map_err(|e| {
if e.kind() == std::io::ErrorKind::NotFound {
ConnectError::new(ErrorCode::NotFound, format!("source not found: {e}"))
} else {
ConnectError::new(ErrorCode::Internal, format!("error renaming: {e}"))
}
})?;
let entry = build_entry_info(&destination)?;
Ok((
MoveResponse {
entry: entry.into(),
..Default::default()
},
ctx,
))
}
async fn list_dir(
&self,
ctx: Context,
request: buffa::view::OwnedView<ListDirRequestView<'static>>,
) -> Result<(ListDirResponse, Context), ConnectError> {
let mut depth = request.depth as usize;
if depth == 0 {
depth = 1;
}
let path = self.resolve_path(request.path, &ctx)?;
let resolved = std::fs::canonicalize(&path).map_err(|e| {
if e.kind() == std::io::ErrorKind::NotFound {
ConnectError::new(ErrorCode::NotFound, format!("path not found: {e}"))
} else {
ConnectError::new(ErrorCode::Internal, format!("error resolving path: {e}"))
}
})?;
let resolved_str = resolved.to_string_lossy().to_string();
let meta = std::fs::metadata(&resolved).map_err(|e| {
ConnectError::new(ErrorCode::Internal, format!("error getting file info: {e}"))
})?;
if !meta.is_dir() {
return Err(ConnectError::new(
ErrorCode::InvalidArgument,
format!("path is not a directory: {path}"),
));
}
let entries = walk_dir(&path, &resolved_str, depth)?;
Ok((
ListDirResponse {
entries,
..Default::default()
},
ctx,
))
}
async fn remove(
&self,
ctx: Context,
request: buffa::view::OwnedView<RemoveRequestView<'static>>,
) -> Result<(RemoveResponse, Context), ConnectError> {
let path = self.resolve_path(request.path, &ctx)?;
if let Err(e1) = std::fs::remove_dir_all(&path) {
if let Err(e2) = std::fs::remove_file(&path) {
return Err(ConnectError::new(
ErrorCode::Internal,
format!("error removing: {e1}; also tried as file: {e2}"),
));
}
}
Ok((RemoveResponse { ..Default::default() }, ctx))
}
async fn watch_dir(
&self,
_ctx: Context,
_request: buffa::view::OwnedView<WatchDirRequestView<'static>>,
) -> Result<
(
Pin<Box<dyn Stream<Item = Result<WatchDirResponse, ConnectError>> + Send>>,
Context,
),
ConnectError,
> {
Err(ConnectError::new(
ErrorCode::Unimplemented,
"watch_dir streaming not yet implemented",
))
}
async fn create_watcher(
&self,
ctx: Context,
request: buffa::view::OwnedView<CreateWatcherRequestView<'static>>,
) -> Result<(CreateWatcherResponse, Context), ConnectError> {
use notify::{RecursiveMode, Watcher};
let path = self.resolve_path(request.path, &ctx)?;
let recursive = request.recursive;
if let Ok(true) = crate::rpc::entry::is_network_mount(&path) {
return Err(ConnectError::new(
ErrorCode::FailedPrecondition,
"watching network mounts is not supported",
));
}
let watcher_id = simple_id();
let events: Arc<Mutex<Vec<FilesystemEvent>>> = Arc::new(Mutex::new(Vec::new()));
let events_cb = Arc::clone(&events);
let mut watcher = notify::recommended_watcher(
move |res: Result<notify::Event, notify::Error>| {
if let Ok(event) = res {
let event_type = match event.kind {
notify::EventKind::Create(_) => EventType::EVENT_TYPE_CREATE,
notify::EventKind::Modify(notify::event::ModifyKind::Data(_)) => {
EventType::EVENT_TYPE_WRITE
}
notify::EventKind::Modify(notify::event::ModifyKind::Metadata(_)) => {
EventType::EVENT_TYPE_CHMOD
}
notify::EventKind::Remove(_) => EventType::EVENT_TYPE_REMOVE,
notify::EventKind::Modify(notify::event::ModifyKind::Name(_)) => {
EventType::EVENT_TYPE_RENAME
}
_ => return,
};
for p in &event.paths {
if let Ok(mut guard) = events_cb.lock() {
guard.push(FilesystemEvent {
name: p.to_string_lossy().to_string(),
r#type: buffa::EnumValue::Known(event_type),
..Default::default()
});
}
}
}
},
)
.map_err(|e| {
ConnectError::new(ErrorCode::Internal, format!("failed to create watcher: {e}"))
})?;
let mode = if recursive {
RecursiveMode::Recursive
} else {
RecursiveMode::NonRecursive
};
watcher.watch(Path::new(&path), mode).map_err(|e| {
ConnectError::new(ErrorCode::Internal, format!("failed to watch path: {e}"))
})?;
self.watchers.insert(
watcher_id.clone(),
WatcherHandle {
events,
_watcher: watcher,
},
);
Ok((
CreateWatcherResponse {
watcher_id,
..Default::default()
},
ctx,
))
}
async fn get_watcher_events(
&self,
ctx: Context,
request: buffa::view::OwnedView<GetWatcherEventsRequestView<'static>>,
) -> Result<(GetWatcherEventsResponse, Context), ConnectError> {
let watcher_id: &str = request.watcher_id;
let handle = self.watchers.get(watcher_id).ok_or_else(|| {
ConnectError::new(
ErrorCode::NotFound,
format!("watcher not found: {watcher_id}"),
)
})?;
let events = {
let mut guard = handle.events.lock().unwrap();
std::mem::take(&mut *guard)
};
Ok((
GetWatcherEventsResponse {
events,
..Default::default()
},
ctx,
))
}
async fn remove_watcher(
&self,
ctx: Context,
request: buffa::view::OwnedView<RemoveWatcherRequestView<'static>>,
) -> Result<(RemoveWatcherResponse, Context), ConnectError> {
let watcher_id: &str = request.watcher_id;
self.watchers.remove(watcher_id);
Ok((RemoveWatcherResponse { ..Default::default() }, ctx))
}
}
fn walk_dir(
requested_path: &str,
resolved_path: &str,
depth: usize,
) -> Result<Vec<EntryInfo>, ConnectError> {
let mut entries = Vec::new();
let base = Path::new(resolved_path);
for result in walkdir::WalkDir::new(resolved_path)
.min_depth(1)
.max_depth(depth)
.follow_links(false)
{
let dir_entry = match result {
Ok(e) => e,
Err(e) => {
if e.io_error()
.is_some_and(|io| io.kind() == std::io::ErrorKind::NotFound)
{
continue;
}
return Err(ConnectError::new(
ErrorCode::Internal,
format!("error reading directory: {e}"),
));
}
};
let entry_path = dir_entry.path();
let mut entry = match build_entry_info(&entry_path.to_string_lossy()) {
Ok(e) => e,
Err(e) if e.code == ErrorCode::NotFound => continue,
Err(e) => return Err(e),
};
if let Ok(rel) = entry_path.strip_prefix(base) {
let remapped = PathBuf::from(requested_path).join(rel);
entry.path = remapped.to_string_lossy().to_string();
}
entries.push(entry);
}
Ok(entries)
}
fn simple_id() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
format!("w-{nanos:x}")
}

26
envd-rs/src/rpc/mod.rs Normal file
View File

@ -0,0 +1,26 @@
pub mod pb;
pub mod entry;
pub mod process_handler;
pub mod process_service;
pub mod filesystem_service;
use std::sync::Arc;
use crate::rpc::process_service::ProcessServiceImpl;
use crate::rpc::filesystem_service::FilesystemServiceImpl;
use crate::state::AppState;
use pb::process::ProcessExt;
use pb::filesystem::FilesystemExt;
/// Build the connect-rust Router with both RPC services registered.
pub fn rpc_router(state: Arc<AppState>) -> connectrpc::Router {
let process_svc = Arc::new(ProcessServiceImpl::new(Arc::clone(&state)));
let filesystem_svc = Arc::new(FilesystemServiceImpl::new(Arc::clone(&state)));
let router = connectrpc::Router::new();
let router = process_svc.register(router);
let router = filesystem_svc.register(router);
router
}

10
envd-rs/src/rpc/pb.rs Normal file
View File

@ -0,0 +1,10 @@
#![allow(dead_code, non_camel_case_types, unused_imports, clippy::derivable_impls)]
use ::buffa;
use ::buffa_types;
use ::connectrpc;
use ::futures;
use ::http_body;
use ::serde;
include!(concat!(env!("OUT_DIR"), "/_connectrpc.rs"));

View 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
}
}

View 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()
}
}

42
envd-rs/src/state.rs Normal file
View File

@ -0,0 +1,42 @@
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use crate::auth::token::SecureToken;
use crate::conntracker::ConnTracker;
use crate::execcontext::Defaults;
use crate::port::subsystem::PortSubsystem;
use crate::util::AtomicMax;
pub struct AppState {
pub defaults: Defaults,
pub version: String,
pub commit: String,
pub is_fc: bool,
pub needs_restore: AtomicBool,
pub last_set_time: AtomicMax,
pub access_token: SecureToken,
pub conn_tracker: ConnTracker,
pub port_subsystem: Option<Arc<PortSubsystem>>,
}
impl AppState {
pub fn new(
defaults: Defaults,
version: String,
commit: String,
is_fc: bool,
port_subsystem: Option<Arc<PortSubsystem>>,
) -> Arc<Self> {
Arc::new(Self {
defaults,
version,
commit,
is_fc,
needs_restore: AtomicBool::new(false),
last_set_time: AtomicMax::new(),
access_token: SecureToken::new(),
conn_tracker: ConnTracker::new(),
port_subsystem,
})
}
}

33
envd-rs/src/util.rs Normal file
View File

@ -0,0 +1,33 @@
use std::sync::atomic::{AtomicI64, Ordering};
pub struct AtomicMax {
val: AtomicI64,
}
impl AtomicMax {
pub fn new() -> Self {
Self {
val: AtomicI64::new(i64::MIN),
}
}
/// Sets the stored value to `new` if `new` is strictly greater than
/// the current value. Returns `true` if the value was updated.
pub fn set_to_greater(&self, new: i64) -> bool {
loop {
let current = self.val.load(Ordering::Acquire);
if new <= current {
return false;
}
match self.val.compare_exchange_weak(
current,
new,
Ordering::Release,
Ordering::Relaxed,
) {
Ok(_) => return true,
Err(_) => continue,
}
}
}
}