feat: implement cold boot and snapshot-based warm start in node daemon

Author: richiemcilroy

Cargo.lock

@@ -15,6 +15,11 @@ tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 tokio-stream = "0.1"
 libc = "0.2"
+hyper = { version = "1", features = ["client", "http1"] }
+hyper-util = { version = "0.1", features = ["client-legacy", "tokio"] }
+http-body-util = "0.1"
+tower = "0.5"
+hyper-unix-socket = "0.3"
 
 [build-dependencies]
 tonic-build = "0.12"

crates/sandchest-node/Cargo.toml

@@ -1,10 +1,21 @@
 fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // Node proto — server stubs (control plane connects to us)
     tonic_build::configure()
         .build_server(true)
         .build_client(false)
         .compile_protos(
             &["sandchest/node/v1/node.proto"],
             &["../../packages/contract/proto"],
         )?;
+
+    // Agent proto — client stubs (we connect to guest agents)
+    tonic_build::configure()
+        .build_server(false)
+        .build_client(true)
+        .compile_protos(
+            &["sandchest/agent/v1/agent.proto"],
+            &["../../packages/contract/proto"],
+        )?;
+
     Ok(())
 }

crates/sandchest-node/build.rs

@@ -0,0 +1,131 @@
+use std::time::Duration;
+
+use tracing::{info, warn};
+
+pub mod agent_proto {
+    tonic::include_proto!("sandchest.agent.v1");
+}
+
+/// Client for communicating with the guest agent inside a Firecracker microVM.
+///
+/// In production, connects via vsock. In dev mode (TCP), connects to localhost.
+pub struct AgentClient {
+    endpoint: String,
+}
+
+impl AgentClient {
+    /// Create a new agent client.
+    ///
+    /// `endpoint` is a gRPC endpoint URI, e.g. `http://127.0.0.1:8052` for TCP dev mode.
+    /// Vsock connections will be added when vsock support is wired up.
+    pub fn new(endpoint: &str) -> Self {
+        Self {
+            endpoint: endpoint.to_string(),
+        }
+    }
+
+    /// Construct the TCP dev-mode endpoint from a vsock path.
+    ///
+    /// In dev mode, the guest agent listens on TCP instead of vsock. We connect
+    /// to localhost on the agent's dev port (default 8052).
+    pub fn dev_endpoint() -> String {
+        let port = std::env::var("SANDCHEST_AGENT_DEV_PORT")
+            .ok()
+            .and_then(|s| s.parse::<u16>().ok())
+            .unwrap_or(8052);
+        format!("http://127.0.0.1:{}", port)
+    }
+
+    /// Poll the guest agent's Health RPC until it reports ready.
+    ///
+    /// Retries every 100ms up to `timeout`. Used after VM boot to confirm
+    /// the guest agent is running and accepting requests.
+    pub async fn wait_for_health(
+        endpoint: &str,
+        timeout: Duration,
+    ) -> Result<(), AgentClientError> {
+        let start = tokio::time::Instant::now();
+        let interval = Duration::from_millis(100);
+
+        info!(endpoint = %endpoint, timeout_ms = timeout.as_millis(), "waiting for guest agent health");
+
+        while start.elapsed() < timeout {
+            match Self::check_health_once(endpoint).await {
+                Ok(true) => {
+                    let elapsed = start.elapsed().as_millis();
+                    info!(endpoint = %endpoint, elapsed_ms = elapsed, "guest agent is healthy");
+                    return Ok(());
+                }
+                Ok(false) => {
+                    warn!(endpoint = %endpoint, "agent responded but not ready");
+                }
+                Err(_) => {
+                    // Connection refused or timeout — agent not ready yet
+                }
+            }
+            tokio::time::sleep(interval).await;
+        }
+
+        Err(AgentClientError::HealthTimeout(format!(
+            "guest agent at {} did not become healthy within {:?}",
+            endpoint, timeout
+        )))
+    }
+
+    async fn check_health_once(endpoint: &str) -> Result<bool, AgentClientError> {
+        let channel = tonic::transport::Channel::from_shared(endpoint.to_string())
+            .map_err(|e| AgentClientError::Connection(format!("invalid endpoint: {}", e)))?
+            .connect_timeout(Duration::from_secs(2))
+            .timeout(Duration::from_secs(5))
+            .connect()
+            .await
+            .map_err(|e| AgentClientError::Connection(format!("connect failed: {}", e)))?;
+
+        let mut client = agent_proto::guest_agent_client::GuestAgentClient::new(channel);
+        let response = client.health(()).await.map_err(|e| {
+            AgentClientError::Rpc(format!("health RPC failed: {}", e))
+        })?;
+
+        Ok(response.into_inner().ready)
+    }
+
+    /// Connect and return a reusable gRPC client handle.
+    pub async fn connect(
+        &self,
+    ) -> Result<agent_proto::guest_agent_client::GuestAgentClient<tonic::transport::Channel>, AgentClientError>
+    {
+        let channel = tonic::transport::Channel::from_shared(self.endpoint.clone())
+            .map_err(|e| AgentClientError::Connection(format!("invalid endpoint: {}", e)))?
+            .connect_timeout(Duration::from_secs(5))
+            .timeout(Duration::from_secs(300))
+            .connect()
+            .await
+            .map_err(|e| {
+                AgentClientError::Connection(format!(
+                    "failed to connect to agent at {}: {}",
+                    self.endpoint, e
+                ))
+            })?;
+
+        Ok(agent_proto::guest_agent_client::GuestAgentClient::new(channel))
+    }
+}
+
+#[derive(Debug)]
+pub enum AgentClientError {
+    HealthTimeout(String),
+    Connection(String),
+    Rpc(String),
+}
+
+impl std::fmt::Display for AgentClientError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            AgentClientError::HealthTimeout(msg) => write!(f, "health timeout: {}", msg),
+            AgentClientError::Connection(msg) => write!(f, "connection error: {}", msg),
+            AgentClientError::Rpc(msg) => write!(f, "RPC error: {}", msg),
+        }
+    }
+}
+
+impl std::error::Error for AgentClientError {}

crates/sandchest-node/src/agent_client.rs

@@ -0,0 +1,151 @@
+use std::path::Path;
+
+use tracing::{info, warn};
+
+/// Clone a base ext4 image into a per-sandbox ext4 file using reflink copy.
+///
+/// On XFS/btrfs this is an instant CoW clone. On other filesystems it falls
+/// back to a regular copy. The cloned file is passed directly to Firecracker
+/// as the drive's `path_on_host`.
+pub async fn clone_disk(src_ext4: &str, sandbox_id: &str, data_dir: &str) -> Result<String, DiskError> {
+    let sandbox_dir = format!("{}/sandboxes/{}", data_dir, sandbox_id);
+    let dest = format!("{}/rootfs.ext4", sandbox_dir);
+
+    // Create sandbox directory
+    tokio::fs::create_dir_all(&sandbox_dir).await.map_err(|e| {
+        DiskError::Io(format!("failed to create sandbox directory {}: {}", sandbox_dir, e))
+    })?;
+
+    if !Path::new(src_ext4).exists() {
+        return Err(DiskError::SourceNotFound(src_ext4.to_string()));
+    }
+
+    info!(
+        src = %src_ext4,
+        dest = %dest,
+        sandbox_id = %sandbox_id,
+        "cloning disk with reflink"
+    );
+
+    let src = src_ext4.to_string();
+    let dst = dest.clone();
+
+    // Use --reflink=auto on Linux for instant CoW clones on XFS/btrfs.
+    // On macOS/other platforms, fall back to regular cp.
+    let output = if cfg!(target_os = "linux") {
+        tokio::process::Command::new("cp")
+            .arg("--reflink=auto")
+            .arg(&src)
+            .arg(&dst)
+            .output()
+            .await
+            .map_err(|e| DiskError::Io(format!("failed to run cp: {}", e)))?
+    } else {
+        tokio::process::Command::new("cp")
+            .arg(&src)
+            .arg(&dst)
+            .output()
+            .await
+            .map_err(|e| DiskError::Io(format!("failed to run cp: {}", e)))?
+    };
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        return Err(DiskError::Io(format!("cp failed: {}", stderr)));
+    }
+
+    info!(sandbox_id = %sandbox_id, dest = %dest, "disk clone complete");
+    Ok(dest)
+}
+
+/// Remove a sandbox's data directory and its contents.
+pub async fn cleanup_disk(sandbox_id: &str, data_dir: &str) -> Result<(), DiskError> {
+    let sandbox_dir = format!("{}/sandboxes/{}", data_dir, sandbox_id);
+
+    if !Path::new(&sandbox_dir).exists() {
+        warn!(sandbox_id = %sandbox_id, "sandbox directory already absent");
+        return Ok(());
+    }
+
+    tokio::fs::remove_dir_all(&sandbox_dir).await.map_err(|e| {
+        DiskError::Io(format!("failed to remove {}: {}", sandbox_dir, e))
+    })?;
+
+    info!(sandbox_id = %sandbox_id, "disk cleanup complete");
+    Ok(())
+}
+
+#[derive(Debug)]
+pub enum DiskError {
+    SourceNotFound(String),
+    Io(String),
+}
+
+impl std::fmt::Display for DiskError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            DiskError::SourceNotFound(path) => write!(f, "source image not found: {}", path),
+            DiskError::Io(msg) => write!(f, "disk I/O error: {}", msg),
+        }
+    }
+}
+
+impl std::error::Error for DiskError {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn clone_disk_fails_for_missing_source() {
+        let result = clone_disk("/nonexistent/rootfs.ext4", "sb_test", "/tmp/sandchest-test").await;
+        assert!(result.is_err());
+        let err = result.unwrap_err();
+        assert!(matches!(err, DiskError::SourceNotFound(_)));
+    }
+
+    #[tokio::test]
+    async fn clone_disk_creates_sandbox_dir_and_copies() {
+        let tmp = std::env::temp_dir().join("sandchest-disk-test");
+        let _ = std::fs::remove_dir_all(&tmp);
+
+        // Create a small source file
+        let src_dir = tmp.join("images");
+        std::fs::create_dir_all(&src_dir).unwrap();
+        let src_file = src_dir.join("rootfs.ext4");
+        std::fs::write(&src_file, b"fake-ext4-data").unwrap();
+
+        let data_dir = tmp.to_str().unwrap();
+        let result = clone_disk(src_file.to_str().unwrap(), "sb_clone_test", data_dir).await;
+        assert!(result.is_ok());
+
+        let dest = result.unwrap();
+        assert!(Path::new(&dest).exists());
+        let content = std::fs::read(&dest).unwrap();
+        assert_eq!(content, b"fake-ext4-data");
+
+        // Cleanup
+        let _ = std::fs::remove_dir_all(&tmp);
+    }
+
+    #[tokio::test]
+    async fn cleanup_disk_removes_directory() {
+        let tmp = std::env::temp_dir().join("sandchest-cleanup-test");
+        let sandbox_dir = tmp.join("sandboxes").join("sb_cleanup");
+        std::fs::create_dir_all(&sandbox_dir).unwrap();
+        std::fs::write(sandbox_dir.join("rootfs.ext4"), b"data").unwrap();
+
+        let data_dir = tmp.to_str().unwrap();
+        let result = cleanup_disk("sb_cleanup", data_dir).await;
+        assert!(result.is_ok());
+        assert!(!sandbox_dir.exists());
+
+        let _ = std::fs::remove_dir_all(&tmp);
+    }
+
+    #[tokio::test]
+    async fn cleanup_disk_is_idempotent() {
+        let result = cleanup_disk("sb_nonexistent", "/tmp/sandchest-idempotent-test").await;
+        assert!(result.is_ok());
+    }
+}

crates/sandchest-node/src/disk.rs

@@ -16,6 +16,23 @@ pub struct FirecrackerVm {
 }
 
 impl FirecrackerVm {
+    /// Construct a FirecrackerVm from pre-existing parts (used for snapshot warm start).
+    pub fn from_parts(
+        sandbox_id: String,
+        api_socket_path: String,
+        vsock_path: String,
+        data_dir: String,
+        child: Child,
+    ) -> Self {
+        Self {
+            sandbox_id,
+            api_socket_path,
+            vsock_path,
+            data_dir,
+            child,
+        }
+    }
+
     /// Start a new Firecracker VM with the given configuration.
     ///
     /// 1. Creates the sandbox data directory