Skip to main content

a3s_box_runtime/vm/
mod.rs

1//! VM Manager - Lifecycle management for MicroVM instances.
2
3mod layout;
4mod network;
5mod ready;
6pub mod reap;
7mod spec;
8
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11
12/// Callback type for image pull progress: `(current, total, digest, size_bytes)`.
13type PullProgressFn = Arc<dyn Fn(usize, usize, &str, i64) + Send + Sync>;
14
15use a3s_box_core::config::BoxConfig;
16#[cfg(unix)]
17use a3s_box_core::config::TeeConfig;
18use a3s_box_core::error::{BoxError, Result};
19use a3s_box_core::event::{BoxEvent, EventEmitter};
20use serde::{Deserialize, Serialize};
21use tokio::sync::RwLock;
22use tracing::Instrument;
23
24#[cfg(unix)]
25use libc;
26
27#[cfg(unix)]
28use crate::grpc::ExecClient;
29#[cfg(unix)]
30use crate::tee::TeeExtension;
31use crate::vmm::{VmController, VmHandler, VmmProvider, DEFAULT_SHUTDOWN_TIMEOUT_MS};
32
33/// Box state machine.
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
35pub enum BoxState {
36    /// Config captured, no VM started
37    Created,
38
39    /// VM booted, container initialized, gRPC healthy
40    Ready,
41
42    /// A session is actively processing a prompt
43    Busy,
44
45    /// A session is compressing its context
46    Compacting,
47
48    /// VM terminated, resources freed
49    Stopped,
50}
51
52/// Layout of directories for a box instance.
53pub(crate) struct BoxLayout {
54    /// Path to the root filesystem
55    pub(crate) rootfs_path: PathBuf,
56    /// Path to the exec Unix socket
57    pub(crate) exec_socket_path: PathBuf,
58    /// Path to the PTY Unix socket
59    pub(crate) pty_socket_path: PathBuf,
60    /// Path to the attestation Unix socket
61    pub(crate) attest_socket_path: PathBuf,
62    /// Path to the CRI port-forward Unix socket
63    pub(crate) port_forward_socket_path: PathBuf,
64    /// Path to the workspace directory
65    pub(crate) workspace_path: PathBuf,
66    /// Path to console output file (optional)
67    pub(crate) console_output: Option<PathBuf>,
68    /// OCI image config (entrypoint, env, working dir, volumes)
69    pub(crate) oci_config: Option<crate::oci::OciImageConfig>,
70    /// TEE instance configuration (if TEE is enabled)
71    pub(crate) tee_instance_config: Option<crate::vmm::TeeInstanceConfig>,
72}
73
74/// VM manager - orchestrates VM lifecycle.
75pub struct VmManager {
76    /// Box configuration
77    pub(crate) config: BoxConfig,
78
79    /// Unique box identifier
80    pub(crate) box_id: String,
81
82    /// Current state
83    pub(crate) state: Arc<RwLock<BoxState>>,
84
85    /// Event emitter
86    pub(crate) event_emitter: EventEmitter,
87
88    /// VMM provider (spawns VMs via pluggable backend)
89    pub(crate) provider: Option<Box<dyn VmmProvider>>,
90
91    /// VM handler (runtime operations on running VM)
92    pub(crate) handler: Arc<RwLock<Option<Box<dyn VmHandler>>>>,
93
94    /// Exec client for executing commands in the guest
95    #[cfg(unix)]
96    pub(crate) exec_client: Option<ExecClient>,
97
98    /// Network backend manager for bridge networking (None if TSI mode).
99    /// Platform-specific: passt on Linux, gvproxy on macOS.
100    pub(crate) net_manager: Option<Box<dyn crate::network::NetworkBackend>>,
101
102    /// A3S home directory (~/.a3s)
103    pub(crate) home_dir: PathBuf,
104
105    /// Anonymous volume names created during boot (from OCI VOLUME directives)
106    pub(crate) anonymous_volumes: Vec<String>,
107
108    /// Anonymous volumes newly created by the current boot attempt.
109    ///
110    /// Reused anonymous volumes must survive failed restarts because they may
111    /// contain data from an existing stopped box.
112    pub(crate) created_anonymous_volumes: Vec<String>,
113
114    /// OCI image config resolved during the last successful boot.
115    pub(crate) image_config: Option<crate::oci::OciImageConfig>,
116
117    /// TEE extension (attestation, sealing, secret injection)
118    #[cfg(unix)]
119    pub(crate) tee: Option<Box<dyn TeeExtension>>,
120
121    /// Rootfs provider (overlay or copy)
122    pub(crate) rootfs_provider: Box<dyn crate::rootfs::RootfsProvider>,
123
124    /// Path to the exec Unix socket (set after boot)
125    pub(crate) exec_socket_path: Option<PathBuf>,
126
127    /// Path to the PTY Unix socket (set after boot)
128    pub(crate) pty_socket_path: Option<PathBuf>,
129
130    /// Path to the CRI port-forward Unix socket (set after boot)
131    pub(crate) port_forward_socket_path: Option<PathBuf>,
132
133    /// Prometheus metrics (optional, for instrumented deployments).
134    pub(crate) prom: Option<crate::prom::RuntimeMetrics>,
135
136    /// Exit code captured from the shim process after it exits.
137    pub(crate) shim_exit_code: Option<i32>,
138
139    /// Optional progress callback for image pulls: `(current, total, digest, size_bytes)`.
140    pub(crate) pull_progress_fn: Option<PullProgressFn>,
141
142    /// Logging driver config, threaded into the InstanceSpec so the shim runs
143    /// the log processor for the box's lifetime (set by the CLI via
144    /// [`VmManager::set_log_config`]).
145    pub(crate) log_config: a3s_box_core::log::LogConfig,
146}
147
148impl VmManager {
149    /// Create a new VM manager.
150    pub fn new(config: BoxConfig, event_emitter: EventEmitter) -> Self {
151        let box_id = uuid::Uuid::new_v4().to_string();
152        let home_dir = a3s_box_core::dirs_home();
153
154        Self {
155            config,
156            box_id,
157            state: Arc::new(RwLock::new(BoxState::Created)),
158            event_emitter,
159            provider: None,
160            handler: Arc::new(RwLock::new(None)),
161            #[cfg(unix)]
162            exec_client: None,
163            net_manager: None,
164            home_dir,
165            anonymous_volumes: Vec::new(),
166            created_anonymous_volumes: Vec::new(),
167            image_config: None,
168            #[cfg(unix)]
169            tee: None,
170            rootfs_provider: crate::rootfs::default_provider(),
171            exec_socket_path: None,
172            pty_socket_path: None,
173            port_forward_socket_path: None,
174            prom: None,
175            shim_exit_code: None,
176            pull_progress_fn: None,
177            log_config: a3s_box_core::log::LogConfig::default(),
178        }
179    }
180
181    /// Create a new VM manager with a specific box ID.
182    pub fn with_box_id(config: BoxConfig, event_emitter: EventEmitter, box_id: String) -> Self {
183        let home_dir = a3s_box_core::dirs_home();
184
185        Self {
186            config,
187            box_id,
188            state: Arc::new(RwLock::new(BoxState::Created)),
189            event_emitter,
190            provider: None,
191            handler: Arc::new(RwLock::new(None)),
192            #[cfg(unix)]
193            exec_client: None,
194            net_manager: None,
195            home_dir,
196            anonymous_volumes: Vec::new(),
197            created_anonymous_volumes: Vec::new(),
198            image_config: None,
199            #[cfg(unix)]
200            tee: None,
201            rootfs_provider: crate::rootfs::default_provider(),
202            exec_socket_path: None,
203            pty_socket_path: None,
204            port_forward_socket_path: None,
205            prom: None,
206            shim_exit_code: None,
207            pull_progress_fn: None,
208            log_config: a3s_box_core::log::LogConfig::default(),
209        }
210    }
211
212    /// Remove host-side boot artifacts after a failed boot attempt.
213    async fn cleanup_boot_failure(&mut self) {
214        if let Some(mut handler) = self.handler.write().await.take() {
215            if let Err(error) = handler.stop(default_stop_signal(), DEFAULT_SHUTDOWN_TIMEOUT_MS) {
216                tracing::warn!(
217                    box_id = %self.box_id,
218                    error = %error,
219                    "Failed to stop VM handler after boot failure"
220                );
221            }
222            self.shim_exit_code = handler.exit_code();
223        }
224
225        if let Some(mut net_manager) = self.net_manager.take() {
226            net_manager.stop();
227        }
228
229        self.cleanup_created_anonymous_volumes();
230        self.cleanup_box_dir();
231    }
232
233    fn cleanup_created_anonymous_volumes(&mut self) {
234        if self.created_anonymous_volumes.is_empty() {
235            return;
236        }
237
238        let created = std::mem::take(&mut self.created_anonymous_volumes);
239        let created_set: std::collections::HashSet<_> = created.iter().cloned().collect();
240        let store = crate::volume::VolumeStore::new(
241            self.home_dir.join("volumes.json"),
242            self.home_dir.join("volumes"),
243        );
244
245        for volume_name in &created {
246            if let Err(error) = store.remove(volume_name, true) {
247                tracing::debug!(
248                    box_id = %self.box_id,
249                    volume = volume_name,
250                    error = %error,
251                    "Failed to remove anonymous volume after boot failure"
252                );
253            }
254        }
255
256        self.anonymous_volumes
257            .retain(|name| !created_set.contains(name));
258    }
259
260    /// Remove the box directory on the host.
261    fn cleanup_box_dir(&self) {
262        let box_dir = self.home_dir.join("boxes").join(&self.box_id);
263
264        // Reap the box's passt daemon (Linux bridge mode) BEFORE removing its
265        // socket dir. A boot that fails after passt spawned but before
266        // `self.net_manager` was assigned leaves `net_manager.stop()` a no-op, so
267        // passt would otherwise survive holding the published port — the
268        // "Address already in use" on the next start. terminate_passt reads
269        // `socket_dir/passt.pid` and is a no-op when there is no passt.
270        #[cfg(target_os = "linux")]
271        crate::network::terminate_passt(&self.socket_dir());
272
273        if let Err(error) = self.rootfs_provider.cleanup(&box_dir, false) {
274            tracing::warn!(
275                box_id = %self.box_id,
276                path = %box_dir.display(),
277                error = %error,
278                "Failed to cleanup rootfs provider after boot failure"
279            );
280        }
281
282        match std::fs::remove_dir_all(&box_dir) {
283            Ok(()) => {}
284            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
285            Err(error) => {
286                tracing::warn!(
287                    box_id = %self.box_id,
288                    path = %box_dir.display(),
289                    error = %error,
290                    "Failed to cleanup box directory after boot failure"
291                );
292            }
293        }
294
295        let socket_dir = self.socket_dir();
296        if socket_dir != box_dir.join("sockets") {
297            match std::fs::remove_dir_all(&socket_dir) {
298                Ok(()) => {}
299                Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
300                Err(error) => {
301                    tracing::debug!(
302                        box_id = %self.box_id,
303                        path = %socket_dir.display(),
304                        error = %error,
305                        "Failed to cleanup socket directory after boot failure"
306                    );
307                }
308            }
309        }
310    }
311
312    /// Create a new VM manager with a custom VMM provider.
313    pub fn with_provider(
314        config: BoxConfig,
315        event_emitter: EventEmitter,
316        provider: Box<dyn VmmProvider>,
317    ) -> Self {
318        let box_id = uuid::Uuid::new_v4().to_string();
319        let home_dir = a3s_box_core::dirs_home();
320        Self {
321            config,
322            box_id,
323            state: Arc::new(RwLock::new(BoxState::Created)),
324            event_emitter,
325            provider: Some(provider),
326            handler: Arc::new(RwLock::new(None)),
327            #[cfg(unix)]
328            exec_client: None,
329            net_manager: None,
330            home_dir,
331            anonymous_volumes: Vec::new(),
332            created_anonymous_volumes: Vec::new(),
333            image_config: None,
334            #[cfg(unix)]
335            tee: None,
336            rootfs_provider: crate::rootfs::default_provider(),
337            exec_socket_path: None,
338            pty_socket_path: None,
339            port_forward_socket_path: None,
340            prom: None,
341            shim_exit_code: None,
342            pull_progress_fn: None,
343            log_config: a3s_box_core::log::LogConfig::default(),
344        }
345    }
346
347    /// Get the box ID.
348    pub fn box_id(&self) -> &str {
349        &self.box_id
350    }
351
352    /// Get current state.
353    pub async fn state(&self) -> BoxState {
354        *self.state.read().await
355    }
356
357    /// Get the exec client, if connected.
358    #[cfg(unix)]
359    pub fn exec_client(&self) -> Option<&ExecClient> {
360        self.exec_client.as_ref()
361    }
362
363    #[cfg(unix)]
364    async fn connect_exec_client_for_request(socket_path: &Path) -> Result<ExecClient> {
365        const ATTEMPT_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(500);
366
367        let client = ExecClient::connect(socket_path).await?;
368        match tokio::time::timeout(ATTEMPT_TIMEOUT, client.heartbeat()).await {
369            Ok(Ok(true)) => Ok(client),
370            Ok(Ok(false)) => Err(BoxError::ExecError(format!(
371                "Exec client not connected: heartbeat failed at {}",
372                socket_path.display()
373            ))),
374            Ok(Err(error)) => Err(error),
375            Err(_) => Err(BoxError::ExecError(format!(
376                "Exec client not connected: heartbeat timed out at {}",
377                socket_path.display()
378            ))),
379        }
380    }
381
382    /// Wait until the guest exec server can complete a heartbeat.
383    ///
384    /// Cold foreground boots may proceed after the short diagnostic readiness
385    /// cap so logs remain visible. A warm pool has a stronger contract: an idle
386    /// VM must actually be executable before it is published to callers.
387    #[cfg(unix)]
388    pub async fn wait_for_exec_available(&mut self, timeout: std::time::Duration) -> Result<()> {
389        let socket_path = self
390            .exec_socket_path
391            .clone()
392            .ok_or_else(|| BoxError::ExecError("Exec socket path is unavailable".to_string()))?;
393        let deadline = tokio::time::Instant::now() + timeout;
394        loop {
395            match Self::connect_exec_client_for_request(&socket_path).await {
396                Ok(client) => {
397                    self.exec_client = Some(client);
398                    return Ok(());
399                }
400                Err(error) if tokio::time::Instant::now() < deadline => {
401                    tracing::debug!(%error, "Waiting for pooled VM exec readiness");
402                    tokio::time::sleep(std::time::Duration::from_millis(200)).await;
403                }
404                Err(error) => return Err(error),
405            }
406        }
407    }
408
409    #[cfg(not(unix))]
410    pub async fn wait_for_exec_available(&mut self, _timeout: std::time::Duration) -> Result<()> {
411        Ok(())
412    }
413
414    /// Attach this manager to an already-running shim process.
415    ///
416    /// This is useful for crash recovery or control-plane restart flows where
417    /// the workload VM is still alive and only the host-side manager state
418    /// needs to be reconstructed.
419    #[cfg(unix)]
420    pub async fn attach_running_process(
421        &mut self,
422        pid: u32,
423        exec_socket_path: PathBuf,
424        pty_socket_path: Option<PathBuf>,
425    ) -> Result<()> {
426        let port_forward_socket_path = exec_socket_path.with_file_name("portfwd.sock");
427        let handler = crate::vmm::ShimHandler::from_pid(pid, self.box_id.clone());
428        if !handler.is_running() {
429            return Err(BoxError::StateError(format!(
430                "Cannot attach to non-running VM process {pid}"
431            )));
432        }
433
434        self.exec_client = match ExecClient::connect(&exec_socket_path).await {
435            Ok(client) => Some(client),
436            Err(error) => {
437                tracing::debug!(
438                    box_id = %self.box_id,
439                    socket_path = %exec_socket_path.display(),
440                    error = %error,
441                    "Failed to reconnect exec client while attaching to running VM"
442                );
443                None
444            }
445        };
446        self.exec_socket_path = Some(exec_socket_path);
447        self.pty_socket_path = pty_socket_path;
448        self.port_forward_socket_path = Some(port_forward_socket_path);
449        *self.handler.write().await = Some(Box::new(handler));
450        *self.state.write().await = BoxState::Ready;
451        Ok(())
452    }
453
454    /// Get the exec socket path, if the VM has been booted.
455    pub fn exec_socket_path(&self) -> Option<&Path> {
456        self.exec_socket_path.as_deref()
457    }
458
459    /// Get the PTY socket path, if the VM has been booted.
460    pub fn pty_socket_path(&self) -> Option<&Path> {
461        self.pty_socket_path.as_deref()
462    }
463
464    /// Get the CRI port-forward socket path, if the VM has been booted.
465    pub fn port_forward_socket_path(&self) -> Option<&Path> {
466        self.port_forward_socket_path.as_deref()
467    }
468
469    /// Inject a custom VMM provider (e.g., a VmController with a known shim path).
470    ///
471    /// If set before `boot()`, the injected provider is used instead of the
472    /// default `VmController::find_shim()` fallback.
473    pub fn set_provider(&mut self, provider: Box<dyn VmmProvider>) {
474        self.provider = Some(provider);
475    }
476
477    /// Override the rootfs provider (overlay or copy).
478    ///
479    /// By default, `default_provider()` auto-detects the best available provider.
480    /// Call this before `boot()` to force a specific provider.
481    pub fn set_rootfs_provider(&mut self, provider: Box<dyn crate::rootfs::RootfsProvider>) {
482        self.rootfs_provider = provider;
483    }
484
485    /// Get the name of the active rootfs provider.
486    pub fn rootfs_provider_name(&self) -> &str {
487        self.rootfs_provider.name()
488    }
489
490    /// Set a progress callback for image pulls: `(current, total, digest, size_bytes)`.
491    /// Called once per layer when `run` pulls an image that is not yet cached.
492    pub fn set_pull_progress_fn(&mut self, f: PullProgressFn) {
493        self.pull_progress_fn = Some(f);
494    }
495
496    /// Attach Prometheus metrics to this VM manager.
497    pub fn set_metrics(&mut self, metrics: crate::prom::RuntimeMetrics) {
498        self.prom = Some(metrics);
499    }
500
501    /// Set the logging driver config. Threaded into the InstanceSpec so the shim
502    /// runs the log processor for the box's lifetime.
503    pub fn set_log_config(&mut self, log_config: a3s_box_core::log::LogConfig) {
504        self.log_config = log_config;
505    }
506
507    /// Get the attached Prometheus metrics (if any).
508    pub fn metrics_prom(&self) -> Option<&crate::prom::RuntimeMetrics> {
509        self.prom.as_ref()
510    }
511
512    /// Get the names of anonymous volumes created during boot.
513    ///
514    /// These are auto-created from OCI VOLUME directives and should be tracked
515    /// for cleanup when the box is removed.
516    pub fn anonymous_volumes(&self) -> &[String] {
517        &self.anonymous_volumes
518    }
519
520    /// Get the OCI image config resolved during boot.
521    pub fn image_config(&self) -> Option<&crate::oci::OciImageConfig> {
522        self.image_config.as_ref()
523    }
524
525    /// Get the exit code of the container, if it has exited.
526    ///
527    /// Returns `Some(code)` after `destroy()` has been called and the shim
528    /// process exited naturally (not killed). Returns `None` if the VM has not
529    /// yet stopped or the exit code could not be determined.
530    pub fn exit_code(&self) -> Option<i32> {
531        self.shim_exit_code
532    }
533
534    fn persisted_exit_code(&self) -> Option<i32> {
535        std::fs::read_to_string(
536            self.home_dir
537                .join("boxes")
538                .join(&self.box_id)
539                .join("upper")
540                .join(".a3s_exit_code"),
541        )
542        .ok()
543        .and_then(|contents| contents.trim().parse::<i32>().ok())
544    }
545
546    /// Poll the owned VM process for natural exit without sending a signal.
547    ///
548    /// This is used by foreground CLI flows where the container command may
549    /// finish on its own and the CLI should clean up instead of waiting for
550    /// a Ctrl-C.
551    pub async fn try_wait_exit(&mut self) -> Result<Option<i32>> {
552        if let Some(code) = self.persisted_exit_code() {
553            self.shim_exit_code = Some(code);
554            return Ok(Some(code));
555        }
556
557        let mut handler = self.handler.write().await;
558        let Some(handler) = handler.as_mut() else {
559            return Ok(self.shim_exit_code);
560        };
561
562        if let Some(code) = handler.try_wait_exit()? {
563            self.shim_exit_code = Some(code);
564            return Ok(Some(code));
565        }
566
567        Ok(None)
568    }
569
570    /// Return true once the foreground container is known to have finished, even
571    /// if the shim exit status has not been reaped yet.
572    pub async fn has_exited(&self) -> bool {
573        if self.shim_exit_code.is_some() || self.persisted_exit_code().is_some() {
574            return true;
575        }
576
577        self.handler
578            .read()
579            .await
580            .as_ref()
581            .map(|handler| handler.has_exited())
582            .unwrap_or(false)
583    }
584
585    /// Run a command as the container MAIN in an IDLE-booted (deferred-main) VM.
586    ///
587    /// Sends the `spawn-main` control frame carrying `spec_json` (the command),
588    /// waits for the main to exit (which halts the VM), and returns its real exit
589    /// code + the box's json-file console logs split by stream. This is the full-
590    /// box-semantics counterpart to [`Self::exec_command`] (whose output is piped
591    /// over the exec stream, not the json-file logs).
592    #[cfg(unix)]
593    pub async fn run_deferred_main(
594        &mut self,
595        spec_json: &[u8],
596        timeout: std::time::Duration,
597    ) -> Result<a3s_box_core::exec::ExecOutput> {
598        let log_dir = self.home_dir.join("boxes").join(&self.box_id).join("logs");
599        let console_out_path = log_dir.join("console.log");
600        let console_err_path = a3s_box_core::log::stderr_console_path(&console_out_path);
601        let console_out_start = std::fs::metadata(&console_out_path)
602            .map(|metadata| metadata.len())
603            .unwrap_or(0);
604        let console_err_start = std::fs::metadata(&console_err_path)
605            .map(|metadata| metadata.len())
606            .unwrap_or(0);
607
608        let acked = {
609            let owned_client;
610            let client = if let Some(client) = self.exec_client.as_ref() {
611                client
612            } else {
613                let socket_path = self
614                    .exec_socket_path
615                    .as_deref()
616                    .ok_or_else(|| BoxError::ExecError("Exec client not connected".to_string()))?;
617                owned_client = Self::connect_exec_client_for_request(socket_path).await?;
618                &owned_client
619            };
620            client.spawn_main(Some(spec_json)).await?
621        };
622        let exit_wait_timeout = if acked {
623            timeout
624        } else {
625            // Very short deferred mains can exit and halt the VM before the
626            // guest's ACK frame makes it back to the host. Treat a missing ACK as
627            // provisional: if the VM exits promptly, the spawn succeeded and the
628            // real exit code/logs are authoritative; otherwise fail quickly
629            // instead of waiting the full command timeout for an IDLE VM.
630            tracing::debug!(
631                box_id = %self.box_id,
632                "spawn-main was not acknowledged; waiting briefly for main exit"
633            );
634            timeout.min(std::time::Duration::from_secs(2))
635        };
636
637        // Wait for the main to exit — guest-init persists the code and halts the VM.
638        let start = std::time::Instant::now();
639        let exit_code = loop {
640            if let Some(code) = self.try_wait_exit().await? {
641                break code;
642            }
643            if start.elapsed() >= exit_wait_timeout {
644                let message = if acked {
645                    "deferred main did not exit within the timeout"
646                } else {
647                    "spawn-main was not acknowledged by the guest"
648                };
649                return Err(BoxError::ExecError(message.to_string()));
650            }
651            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
652        };
653
654        // Let the shim's log processor finish draining console.log into the json
655        // file (it flushes as the VM halts). A single short "stable length"
656        // sample is not enough here: deferred-main can persist its exit code
657        // before the final stdout/stderr bytes have reached the host tailer,
658        // especially with pre-warmed pools. Require a small quiet window before
659        // reading logs, bounded so no-output commands still return promptly.
660        let json_path = log_dir.join("container.json");
661        let drain_start = std::time::Instant::now();
662        let max_wait = std::time::Duration::from_secs(2);
663        let min_wait = std::time::Duration::from_millis(500);
664        let quiet_window = std::time::Duration::from_millis(200);
665        let mut last_len: Option<u64> = None;
666        let mut last_change = drain_start;
667        loop {
668            let len = std::fs::metadata(&json_path).map(|m| m.len()).unwrap_or(0);
669            if last_len != Some(len) {
670                last_len = Some(len);
671                last_change = std::time::Instant::now();
672            }
673            let elapsed = drain_start.elapsed();
674            if elapsed >= max_wait || (elapsed >= min_wait && last_change.elapsed() >= quiet_window)
675            {
676                break;
677            }
678            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
679        }
680        let (mut stdout, mut stderr) = self.read_container_logs();
681        if stdout.is_empty() {
682            stdout = Self::read_file_from_offset(&console_out_path, console_out_start);
683        }
684        if stderr.is_empty() {
685            stderr = Self::read_file_from_offset(&console_err_path, console_err_start);
686        }
687        Ok(a3s_box_core::exec::ExecOutput {
688            stdout,
689            stderr,
690            exit_code,
691        })
692    }
693
694    fn read_file_from_offset(path: &Path, offset: u64) -> Vec<u8> {
695        use std::io::{Read, Seek, SeekFrom};
696
697        let mut file = match std::fs::File::open(path) {
698            Ok(file) => file,
699            Err(_) => return vec![],
700        };
701        if file.seek(SeekFrom::Start(offset)).is_err() {
702            return vec![];
703        }
704
705        let mut bytes = Vec::new();
706        if file.read_to_end(&mut bytes).is_err() {
707            return vec![];
708        }
709        bytes
710    }
711
712    /// Read the box's json-file console logs, split into stdout/stderr by stream.
713    fn read_container_logs(&self) -> (Vec<u8>, Vec<u8>) {
714        let path = self
715            .home_dir
716            .join("boxes")
717            .join(&self.box_id)
718            .join("logs")
719            .join("container.json");
720        let (mut out, mut err) = (Vec::new(), Vec::new());
721        if let Ok(content) = std::fs::read_to_string(&path) {
722            for line in content.lines() {
723                if let Ok(entry) = serde_json::from_str::<a3s_box_core::log::LogEntry>(line) {
724                    if entry.stream == "stderr" {
725                        err.extend_from_slice(entry.log.as_bytes());
726                    } else {
727                        out.extend_from_slice(entry.log.as_bytes());
728                    }
729                }
730            }
731        }
732        (out, err)
733    }
734
735    /// Execute a command in the guest VM.
736    ///
737    /// Requires the VM to be in Ready, Busy, or Compacting state.
738    #[cfg(unix)]
739    #[tracing::instrument(skip(self, request), fields(box_id = %self.box_id))]
740    pub async fn exec_request(
741        &self,
742        request: &a3s_box_core::exec::ExecRequest,
743    ) -> Result<a3s_box_core::exec::ExecOutput> {
744        if request.cmd.is_empty() {
745            return Err(BoxError::ExecError(
746                "Exec request requires a non-empty command".to_string(),
747            ));
748        }
749
750        let state = self.state.read().await;
751        match *state {
752            BoxState::Ready | BoxState::Busy | BoxState::Compacting => {}
753            BoxState::Created => {
754                return Err(BoxError::ExecError("VM not yet booted".to_string()));
755            }
756            BoxState::Stopped => {
757                return Err(BoxError::ExecError("VM is stopped".to_string()));
758            }
759        }
760        drop(state);
761
762        let owned_client;
763        let client = if let Some(client) = self.exec_client.as_ref() {
764            client
765        } else {
766            let socket_path = self
767                .exec_socket_path
768                .as_deref()
769                .ok_or_else(|| BoxError::ExecError("Exec client not connected".to_string()))?;
770            owned_client = Self::connect_exec_client_for_request(socket_path).await?;
771            &owned_client
772        };
773
774        let exec_start = std::time::Instant::now();
775        let result = client.exec_command(request).await;
776
777        // Record Prometheus metrics
778        if let Some(ref prom) = self.prom {
779            prom.exec_total.inc();
780            prom.exec_duration
781                .observe(exec_start.elapsed().as_secs_f64());
782            if result.is_err() || result.as_ref().is_ok_and(|o| o.exit_code != 0) {
783                prom.exec_errors_total.inc();
784            }
785        }
786
787        result
788    }
789
790    /// Execute a command in the guest VM.
791    ///
792    /// Requires the VM to be in Ready, Busy, or Compacting state.
793    #[cfg(unix)]
794    #[tracing::instrument(skip(self, cmd), fields(box_id = %self.box_id))]
795    pub async fn exec_command(
796        &self,
797        cmd: Vec<String>,
798        timeout_ns: u64,
799    ) -> Result<a3s_box_core::exec::ExecOutput> {
800        let request = a3s_box_core::exec::ExecRequest {
801            cmd,
802            timeout_ns,
803            env: vec![],
804            working_dir: None,
805            rootfs: None,
806            stdin: None,
807            stdin_streaming: false,
808            user: None,
809            streaming: false,
810        };
811
812        self.exec_request(&request).await
813    }
814
815    /// Boot the VM.
816    pub async fn boot(&mut self) -> Result<()> {
817        let boot_span = tracing::info_span!("vm_boot", box_id = %self.box_id);
818        // Check and transition state: Created → booting
819        {
820            let state = self.state.read().await;
821            if *state != BoxState::Created {
822                return Err(BoxError::StateError("VM already booted".to_string()));
823            }
824        }
825
826        let boot_start = std::time::Instant::now();
827
828        tracing::info!(parent: &boot_span, box_id = %self.box_id, "Booting VM");
829
830        // 1. Prepare filesystem layout
831        let layout = match self
832            .prepare_layout()
833            .instrument(tracing::info_span!(parent: &boot_span, "prepare_layout"))
834            .await
835        {
836            Ok(layout) => layout,
837            Err(error) => {
838                self.cleanup_boot_failure().await;
839                return Err(error);
840            }
841        };
842        self.image_config = layout.oci_config.clone();
843
844        // 1.5. Override /etc/resolv.conf with configured DNS
845        let resolv_content = a3s_box_core::dns::generate_resolv_conf(&self.config.dns);
846        let resolv_path = layout.rootfs_path.join("etc/resolv.conf");
847        if let Err(e) = tokio::fs::write(&resolv_path, &resolv_content).await {
848            self.cleanup_boot_failure().await;
849            return Err(BoxError::IoError(e));
850        }
851        tracing::debug!(parent: &boot_span, dns = %resolv_content.trim(), "Configured guest DNS");
852
853        // 1.6. Apply hostname and static hosts entries before the VM starts.
854        if let Err(e) = self.write_hostname_file(&layout) {
855            self.cleanup_boot_failure().await;
856            return Err(e);
857        }
858        if let Err(e) = self.write_standalone_hosts_file(&layout) {
859            self.cleanup_boot_failure().await;
860            return Err(e);
861        }
862
863        // 2. Build InstanceSpec
864        let mut spec = match self.build_instance_spec(&layout) {
865            Ok(s) => s,
866            Err(e) => {
867                self.cleanup_boot_failure().await;
868                return Err(e);
869            }
870        };
871
872        // 2.5. Configure bridge networking if requested
873        let bridge_network = match &self.config.network {
874            a3s_box_core::NetworkMode::Bridge { network } => Some(network.clone()),
875            _ => None,
876        };
877        if let Some(network_name) = bridge_network {
878            let net_config = match self.setup_bridge_network(&network_name) {
879                Ok(n) => n,
880                Err(e) => {
881                    self.cleanup_boot_failure().await;
882                    return Err(e);
883                }
884            };
885
886            // Write /etc/hosts for DNS service discovery
887            match self.write_hosts_file(&layout, &network_name) {
888                Ok(()) => (),
889                Err(e) => {
890                    self.cleanup_boot_failure().await;
891                    return Err(e);
892                }
893            };
894
895            // Inject network env vars into entrypoint so they are passed via
896            // krun_set_exec's envp (not krun_set_env which overwrites all vars).
897            let ip_cidr = format!("{}/{}", net_config.ip_address, net_config.prefix_len);
898            spec.entrypoint
899                .env
900                .push(("A3S_NET_IP".to_string(), ip_cidr));
901            spec.entrypoint.env.push((
902                "A3S_NET_GATEWAY".to_string(),
903                net_config.gateway.to_string(),
904            ));
905            spec.entrypoint.env.push((
906                "A3S_NET_DNS".to_string(),
907                net_config
908                    .dns_servers
909                    .iter()
910                    .map(|s| s.to_string())
911                    .collect::<Vec<_>>()
912                    .join(","),
913            ));
914
915            spec.network = Some(net_config);
916        }
917
918        #[cfg(target_os = "macos")]
919        if spec.network.is_none()
920            && matches!(self.config.network, a3s_box_core::NetworkMode::Tsi)
921            && !self.config.port_map.is_empty()
922        {
923            let net_config = match self.setup_published_default_network() {
924                Ok(network) => network,
925                Err(error) => {
926                    self.cleanup_boot_failure().await;
927                    return Err(error);
928                }
929            };
930            let ip_cidr = format!("{}/{}", net_config.ip_address, net_config.prefix_len);
931            spec.entrypoint
932                .env
933                .push(("A3S_NET_IP".to_string(), ip_cidr));
934            spec.entrypoint.env.push((
935                "A3S_NET_GATEWAY".to_string(),
936                net_config.gateway.to_string(),
937            ));
938            spec.entrypoint.env.push((
939                "A3S_NET_DNS".to_string(),
940                net_config
941                    .dns_servers
942                    .iter()
943                    .map(ToString::to_string)
944                    .collect::<Vec<_>>()
945                    .join(","),
946            ));
947            spec.network = Some(net_config);
948        }
949
950        // 3. Initialize VMM provider (use injected provider or default to VmController)
951        if self.provider.is_none() {
952            let shim_path = match VmController::find_shim() {
953                Ok(p) => p,
954                Err(e) => {
955                    self.cleanup_boot_failure().await;
956                    return Err(e);
957                }
958            };
959            let controller = match VmController::new(shim_path) {
960                Ok(c) => c,
961                Err(e) => {
962                    self.cleanup_boot_failure().await;
963                    return Err(e);
964                }
965            };
966            self.provider = Some(Box::new(controller));
967        }
968
969        // 4. Start VM via provider
970        let handler = {
971            let provider = self
972                .provider
973                .as_ref()
974                .ok_or_else(|| BoxError::BoxBootError {
975                    message: "VMM provider not initialized".to_string(),
976                    hint: Some("Ensure VmManager has a provider set before boot".to_string()),
977                })?;
978            let vm_start_span = tracing::info_span!(parent: &boot_span, "vm_start");
979            match async { provider.start(&spec).await }
980                .instrument(vm_start_span)
981                .await
982            {
983                Ok(h) => h,
984                Err(e) => {
985                    self.cleanup_boot_failure().await;
986                    return Err(e);
987                }
988            }
989        };
990
991        // Store handler
992        *self.handler.write().await = Some(handler);
993
994        // 5. Wait for guest ready
995        {
996            let wait_span = tracing::info_span!(parent: &boot_span, "wait_for_ready");
997            if let Err(e) = async {
998                self.wait_for_vm_running().await?;
999
1000                // 5b. Become ready. A snapshot-restore boot resumes an already-booted
1001                // guest whose exec server won't re-signal readiness, so the cold-boot
1002                // wait would stall registration on its safety cap — do one best-effort
1003                // probe instead. A normal boot waits for the Heartbeat health check.
1004                #[cfg(unix)]
1005                if is_restore_mode(&self.config) {
1006                    self.probe_exec_ready_once(&layout.exec_socket_path).await;
1007                } else {
1008                    self.wait_for_exec_ready(&layout.exec_socket_path).await?;
1009                }
1010                Ok::<(), BoxError>(())
1011            }
1012            .instrument(wait_span)
1013            .await
1014            {
1015                self.cleanup_boot_failure().await;
1016                return Err(e);
1017            }
1018        }
1019
1020        // Prototype: deferred-main-spawn. The guest booted IDLE (BOX_DEFERRED_MAIN);
1021        // now that the exec server is ready, tell it to spawn the container command
1022        // (already passed via BOX_EXEC_*) as the MAIN process — full box semantics
1023        // (exit code + json-file console logs) without a cold boot.
1024        // Auto-trigger spawn-main only for the env-driven `run` path, where the
1025        // command is known at boot. The pool sets config.deferred_main to boot the
1026        // VM IDLE but drives spawn-main EXPLICITLY per request (the per-request
1027        // command isn't known at pre-warm), so a pool VM must NOT auto-trigger here.
1028        // A restored guest's main is ALREADY running (captured in the snapshot), so
1029        // it must never re-spawn — doing so would start a duplicate main.
1030        #[cfg(unix)]
1031        if !is_restore_mode(&self.config)
1032            && std::env::var("BOX_DEFERRED_MAIN")
1033                .map(|v| v == "1")
1034                .unwrap_or(false)
1035        {
1036            if let Some(client) = self.exec_client.as_ref() {
1037                match client.spawn_main(None).await {
1038                    Ok(true) => tracing::info!("deferred container main spawned"),
1039                    Ok(false) => tracing::warn!("deferred spawn-main not acknowledged"),
1040                    Err(e) => tracing::warn!(error = %e, "deferred spawn-main failed"),
1041                }
1042            }
1043        }
1044
1045        // 5b2. Store socket paths for CRI streaming access
1046        self.exec_socket_path = Some(layout.exec_socket_path.clone());
1047        self.pty_socket_path = Some(layout.pty_socket_path.clone());
1048        self.port_forward_socket_path = Some(layout.port_forward_socket_path.clone());
1049
1050        // 5c. Initialize TEE extension for TEE environments
1051        #[cfg(unix)]
1052        if !matches!(self.config.tee, TeeConfig::None) {
1053            self.tee = Some(Box::new(crate::tee::SnpTeeExtension::new(
1054                self.box_id.clone(),
1055                layout.attest_socket_path.clone(),
1056            )));
1057        }
1058
1059        // 6. Update state to Ready
1060        *self.state.write().await = BoxState::Ready;
1061
1062        // Record Prometheus metrics
1063        if let Some(ref prom) = self.prom {
1064            let boot_duration = boot_start.elapsed().as_secs_f64();
1065            prom.vm_boot_duration.observe(boot_duration);
1066            prom.vm_created_total.inc();
1067            prom.vm_count.with_label_values(&["ready"]).inc();
1068        }
1069
1070        // Emit ready event
1071        self.event_emitter.emit(BoxEvent::empty("box.ready"));
1072
1073        tracing::info!(parent: &boot_span, box_id = %self.box_id, "VM ready");
1074
1075        Ok(())
1076    }
1077
1078    /// Destroy the VM with the default shutdown timeout and SIGTERM.
1079    pub async fn destroy(&mut self) -> Result<()> {
1080        self.destroy_with_options(default_stop_signal(), DEFAULT_SHUTDOWN_TIMEOUT_MS)
1081            .await
1082    }
1083
1084    /// Destroy the VM with a custom shutdown timeout and SIGTERM.
1085    pub async fn destroy_with_timeout(&mut self, timeout_ms: u64) -> Result<()> {
1086        self.destroy_with_options(libc::SIGTERM, timeout_ms).await
1087    }
1088
1089    /// Destroy the VM with a specific stop signal and timeout.
1090    ///
1091    /// Sends `signal` to the shim process and waits up to `timeout_ms` for it
1092    /// to exit gracefully before sending SIGKILL.
1093    #[tracing::instrument(skip(self), fields(box_id = %self.box_id))]
1094    pub async fn destroy_with_options(&mut self, signal: i32, timeout_ms: u64) -> Result<()> {
1095        let mut state = self.state.write().await;
1096
1097        if *state == BoxState::Stopped {
1098            return Ok(());
1099        }
1100
1101        tracing::info!(box_id = %self.box_id, signal, timeout_ms, "Destroying VM");
1102
1103        // Mark as stopped first — ensures state is correct even if handler.stop() fails.
1104        *state = BoxState::Stopped;
1105
1106        // Stop the VM handler and capture its exit code before it's dropped.
1107        // A stop failure must NOT skip the host-resource teardown below (network
1108        // backend, overlay unmount, socket + box dirs) — those are already
1109        // best-effort and would otherwise leak on every wedged stop. Capture the
1110        // error and surface it after teardown instead of returning early.
1111        let mut stop_error = None;
1112        if let Some(mut handler) = self.handler.write().await.take() {
1113            if let Err(e) = handler.stop(signal, timeout_ms) {
1114                tracing::error!(box_id = %self.box_id, error = %e, "Failed to stop VM handler; continuing teardown");
1115                stop_error = Some(e);
1116            }
1117            self.shim_exit_code = handler.exit_code();
1118        }
1119
1120        // Stop network backend if running
1121        if let Some(ref mut net) = self.net_manager {
1122            net.stop();
1123        }
1124        self.net_manager = None;
1125
1126        // Cleanup rootfs provider (unmount overlay if applicable)
1127        let box_dir = self.home_dir.join("boxes").join(&self.box_id);
1128        if let Err(e) = self
1129            .rootfs_provider
1130            .cleanup(&box_dir, self.config.persistent)
1131        {
1132            tracing::warn!(
1133                box_id = %self.box_id,
1134                error = %e,
1135                "Failed to cleanup rootfs provider"
1136            );
1137        }
1138
1139        let socket_dir = self.socket_dir();
1140        if let Err(e) = std::fs::remove_dir_all(&socket_dir) {
1141            tracing::debug!(
1142                box_id = %self.box_id,
1143                path = %socket_dir.display(),
1144                error = %e,
1145                "Failed to cleanup VM socket directory"
1146            );
1147        }
1148
1149        // Remove the box working directory itself (overlay upper/work, logs,
1150        // leftover metadata) for non-persistent boxes. Without this, ephemeral
1151        // CRI pods leak their `boxes/<id>` directory on every destroy; the
1152        // accumulation slows later RunPodSandbox calls until they time out
1153        // (observed: pod #21 after churning 20). Persistent boxes keep their
1154        // dir intentionally.
1155        if !self.config.persistent {
1156            match std::fs::remove_dir_all(&box_dir) {
1157                Ok(()) => {}
1158                Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
1159                Err(e) => {
1160                    tracing::warn!(
1161                        box_id = %self.box_id,
1162                        path = %box_dir.display(),
1163                        error = %e,
1164                        "Failed to remove box directory on destroy"
1165                    );
1166                }
1167            }
1168        }
1169
1170        // Record Prometheus metrics
1171        if let Some(ref prom) = self.prom {
1172            prom.vm_destroyed_total.inc();
1173            prom.vm_count.with_label_values(&["ready"]).dec();
1174        }
1175
1176        // Emit stopped event
1177        self.event_emitter.emit(BoxEvent::empty("box.stopped"));
1178
1179        // Host teardown above is complete; surface a handler-stop failure now so
1180        // the caller still learns the stop was imperfect.
1181        match stop_error {
1182            Some(e) => Err(e),
1183            None => Ok(()),
1184        }
1185    }
1186
1187    /// Transition to busy state.
1188    pub async fn set_busy(&self) -> Result<()> {
1189        let mut state = self.state.write().await;
1190
1191        if *state != BoxState::Ready {
1192            return Err(BoxError::StateError("VM not ready".to_string()));
1193        }
1194
1195        *state = BoxState::Busy;
1196        Ok(())
1197    }
1198
1199    /// Transition back to ready state.
1200    pub async fn set_ready(&self) -> Result<()> {
1201        let mut state = self.state.write().await;
1202
1203        if *state != BoxState::Busy && *state != BoxState::Compacting {
1204            return Err(BoxError::StateError("Invalid state transition".to_string()));
1205        }
1206
1207        *state = BoxState::Ready;
1208        Ok(())
1209    }
1210
1211    /// Transition to compacting state.
1212    pub async fn set_compacting(&self) -> Result<()> {
1213        let mut state = self.state.write().await;
1214
1215        if *state != BoxState::Busy {
1216            return Err(BoxError::StateError("VM not busy".to_string()));
1217        }
1218
1219        *state = BoxState::Compacting;
1220        Ok(())
1221    }
1222
1223    /// Pause the VM by sending SIGSTOP to the shim process.
1224    ///
1225    /// The VM must be in Ready, Busy, or Compacting state.
1226    #[cfg(unix)]
1227    pub async fn pause(&self) -> Result<()> {
1228        let state = self.state.read().await;
1229        match *state {
1230            BoxState::Ready | BoxState::Busy | BoxState::Compacting => {}
1231            BoxState::Created => {
1232                return Err(BoxError::StateError("VM not yet booted".to_string()));
1233            }
1234            BoxState::Stopped => {
1235                return Err(BoxError::StateError("VM is stopped".to_string()));
1236            }
1237        }
1238        drop(state);
1239
1240        if let Some(pid) = self.pid().await {
1241            // Safety: sending SIGSTOP to pause the process
1242            let ret = unsafe { libc::kill(pid as i32, libc::SIGSTOP) };
1243            if ret != 0 {
1244                let err = std::io::Error::last_os_error();
1245                return Err(BoxError::ExecError(format!(
1246                    "Failed to send SIGSTOP to pid {}: {}",
1247                    pid, err
1248                )));
1249            }
1250            tracing::info!(box_id = %self.box_id, pid, "VM paused");
1251            Ok(())
1252        } else {
1253            Err(BoxError::StateError(
1254                "VM has no running process".to_string(),
1255            ))
1256        }
1257    }
1258
1259    /// Resume the VM by sending SIGCONT to the shim process.
1260    ///
1261    /// Can be called on a paused VM to resume execution.
1262    #[cfg(unix)]
1263    pub async fn resume(&self) -> Result<()> {
1264        if let Some(pid) = self.pid().await {
1265            // Safety: sending SIGCONT to resume the process
1266            let ret = unsafe { libc::kill(pid as i32, libc::SIGCONT) };
1267            if ret != 0 {
1268                let err = std::io::Error::last_os_error();
1269                return Err(BoxError::ExecError(format!(
1270                    "Failed to send SIGCONT to pid {}: {}",
1271                    pid, err
1272                )));
1273            }
1274            tracing::info!(box_id = %self.box_id, pid, "VM resumed");
1275            Ok(())
1276        } else {
1277            Err(BoxError::StateError(
1278                "VM has no running process".to_string(),
1279            ))
1280        }
1281    }
1282
1283    /// Pause the VM (Windows stub - not yet implemented).
1284    #[cfg(windows)]
1285    pub async fn pause(&self) -> Result<()> {
1286        Err(BoxError::StateError(
1287            "VM pause is not yet supported on Windows".to_string(),
1288        ))
1289    }
1290
1291    /// Resume the VM (Windows stub - not yet implemented).
1292    #[cfg(windows)]
1293    pub async fn resume(&self) -> Result<()> {
1294        Err(BoxError::StateError(
1295            "VM resume is not yet supported on Windows".to_string(),
1296        ))
1297    }
1298
1299    /// Check if VM is healthy.
1300    pub async fn health_check(&self) -> Result<bool> {
1301        let state = self.state.read().await;
1302
1303        match *state {
1304            BoxState::Ready | BoxState::Busy | BoxState::Compacting => {
1305                // Check if handler reports VM is running
1306                if let Some(ref handler) = *self.handler.read().await {
1307                    Ok(handler.is_running())
1308                } else {
1309                    Ok(false)
1310                }
1311            }
1312            _ => Ok(false),
1313        }
1314    }
1315
1316    /// Get VM metrics.
1317    pub async fn metrics(&self) -> Option<crate::vmm::VmMetrics> {
1318        let vm_metrics = self
1319            .handler
1320            .read()
1321            .await
1322            .as_ref()
1323            .map(|handler| handler.metrics())?;
1324
1325        // Update per-VM Prometheus gauges if metrics are attached
1326        if let Some(ref prom) = self.prom {
1327            prom.vm_cpu_percent
1328                .with_label_values(&[&self.box_id])
1329                .set(vm_metrics.cpu_percent.unwrap_or(0.0) as f64);
1330            prom.vm_memory_bytes
1331                .with_label_values(&[&self.box_id])
1332                .set(vm_metrics.memory_bytes.unwrap_or(0) as f64);
1333        }
1334
1335        Some(vm_metrics)
1336    }
1337
1338    /// Get the PID of the VM shim process.
1339    pub async fn pid(&self) -> Option<u32> {
1340        self.handler
1341            .read()
1342            .await
1343            .as_ref()
1344            .map(|handler| handler.pid())
1345    }
1346
1347    /// Get the TEE extension, if TEE is configured and VM is booted.
1348    #[cfg(unix)]
1349    pub fn tee(&self) -> Option<&dyn TeeExtension> {
1350        self.tee.as_deref()
1351    }
1352
1353    /// Get the TEE extension or return an error.
1354    #[cfg(unix)]
1355    pub fn require_tee(&self) -> Result<&dyn TeeExtension> {
1356        self.tee.as_deref().ok_or_else(|| {
1357            BoxError::AttestationError("TEE is not configured for this box".to_string())
1358        })
1359    }
1360
1361    /// Apply a live resource update to the running VM.
1362    ///
1363    /// Tier 1 changes (vCPU count, memory size) are rejected with a clear error
1364    /// because libkrun does not expose a hot-resize API.
1365    ///
1366    /// Tier 2 changes (cgroup-based limits) are applied by executing shell
1367    /// commands inside the guest that write to cgroup v2 control files.
1368    #[cfg(unix)]
1369    pub async fn update_resources(
1370        &self,
1371        update: &crate::resize::ResourceUpdate,
1372    ) -> Result<crate::resize::ResizeResult> {
1373        // Reject Tier 1 changes upfront
1374        crate::resize::validate_update(update)?;
1375
1376        let mut result = crate::resize::ResizeResult {
1377            applied: Vec::new(),
1378            rejected: Vec::new(),
1379        };
1380
1381        if !update.has_tier2_changes() {
1382            return Ok(result);
1383        }
1384
1385        // Build cgroup commands and execute them inside the guest
1386        let commands = update.build_cgroup_commands();
1387        for cmd_str in &commands {
1388            let shell_cmd = vec!["sh".to_string(), "-c".to_string(), cmd_str.clone()];
1389
1390            match self.exec_command(shell_cmd, 5_000_000_000).await {
1391                Ok(output) if output.exit_code == 0 => {
1392                    result.applied.push(cmd_str.clone());
1393                }
1394                Ok(output) => {
1395                    let stderr = String::from_utf8_lossy(&output.stderr);
1396                    let reason = if stderr.trim().is_empty() {
1397                        format!("exit code {}", output.exit_code)
1398                    } else {
1399                        stderr.trim().to_string()
1400                    };
1401                    tracing::warn!(
1402                        box_id = %self.box_id,
1403                        cmd = %cmd_str,
1404                        exit_code = output.exit_code,
1405                        stderr = %stderr,
1406                        "Cgroup update failed inside guest"
1407                    );
1408                    result.rejected.push((cmd_str.clone(), reason));
1409                }
1410                Err(e) => {
1411                    tracing::warn!(
1412                        box_id = %self.box_id,
1413                        cmd = %cmd_str,
1414                        error = %e,
1415                        "Failed to exec cgroup update in guest"
1416                    );
1417                    result.rejected.push((cmd_str.clone(), e.to_string()));
1418                }
1419            }
1420        }
1421
1422        Ok(result)
1423    }
1424}
1425
1426/// Whether this boot is a snapshot-fork restore (the guest is resumed already-booted
1427/// rather than cold-booted). PER-VM: a pool / fork daemon sets `config.restore_from`
1428/// so one process can restore different VMs; the single-VM `run` path uses the
1429/// `KRUN_RESTORE_FROM` env. Either source means restore mode.
1430#[cfg(unix)]
1431fn is_restore_mode(config: &BoxConfig) -> bool {
1432    config
1433        .restore_from
1434        .as_deref()
1435        .is_some_and(|s| !s.is_empty())
1436        || std::env::var("KRUN_RESTORE_FROM")
1437            .map(|v| !v.is_empty())
1438            .unwrap_or(false)
1439}
1440
1441/// Simple FNV-1a hash for generating short deterministic hashes from strings.
1442pub(crate) fn fnv1a_hash(input: &str) -> u64 {
1443    let mut hash: u64 = 0xcbf29ce484222325;
1444    for byte in input.bytes() {
1445        hash ^= byte as u64;
1446        hash = hash.wrapping_mul(0x100000001b3);
1447    }
1448    hash
1449}
1450
1451#[cfg(unix)]
1452fn default_stop_signal() -> i32 {
1453    libc::SIGTERM
1454}
1455
1456#[cfg(windows)]
1457fn default_stop_signal() -> i32 {
1458    15
1459}
1460
1461#[cfg(test)]
1462mod tests {
1463    use super::*;
1464    use a3s_box_core::event::EventEmitter;
1465    use std::sync::{
1466        atomic::{AtomicBool, Ordering},
1467        Arc,
1468    };
1469
1470    struct RecordingHandler {
1471        stopped: Arc<AtomicBool>,
1472    }
1473
1474    impl VmHandler for RecordingHandler {
1475        fn stop(&mut self, _signal: i32, _timeout_ms: u64) -> Result<()> {
1476            self.stopped.store(true, Ordering::SeqCst);
1477            Ok(())
1478        }
1479
1480        fn metrics(&self) -> crate::vmm::VmMetrics {
1481            crate::vmm::VmMetrics::default()
1482        }
1483
1484        fn is_running(&self) -> bool {
1485            true
1486        }
1487
1488        fn pid(&self) -> u32 {
1489            42
1490        }
1491    }
1492
1493    struct ExitStateHandler {
1494        exited: bool,
1495    }
1496
1497    impl VmHandler for ExitStateHandler {
1498        fn stop(&mut self, _signal: i32, _timeout_ms: u64) -> Result<()> {
1499            Ok(())
1500        }
1501
1502        fn metrics(&self) -> crate::vmm::VmMetrics {
1503            crate::vmm::VmMetrics::default()
1504        }
1505
1506        fn is_running(&self) -> bool {
1507            !self.exited
1508        }
1509
1510        fn has_exited(&self) -> bool {
1511            self.exited
1512        }
1513
1514        fn pid(&self) -> u32 {
1515            42
1516        }
1517    }
1518
1519    /// A handler whose `stop` always fails — models a wedged VM that won't halt.
1520    struct FailingHandler;
1521
1522    impl VmHandler for FailingHandler {
1523        fn stop(&mut self, _signal: i32, _timeout_ms: u64) -> Result<()> {
1524            Err(BoxError::StateError("simulated stop failure".to_string()))
1525        }
1526
1527        fn metrics(&self) -> crate::vmm::VmMetrics {
1528            crate::vmm::VmMetrics::default()
1529        }
1530
1531        fn is_running(&self) -> bool {
1532            true
1533        }
1534
1535        fn pid(&self) -> u32 {
1536            42
1537        }
1538    }
1539
1540    // Regression: a handler-stop failure must still run the host teardown
1541    // (overlay unmount, socket + box dirs). Pre-fix, destroy_with_options
1542    // returned early on the stop error and leaked the box directory on every
1543    // wedged stop.
1544    #[tokio::test]
1545    async fn test_destroy_runs_host_teardown_even_when_handler_stop_fails() {
1546        let tmp = tempfile::tempdir().unwrap();
1547        let box_id = "box-stopfail".to_string();
1548        // persistent=false (the default) → the box dir is removed on destroy.
1549        let mut vm =
1550            VmManager::with_box_id(BoxConfig::default(), EventEmitter::new(16), box_id.clone());
1551        vm.home_dir = tmp.path().to_path_buf();
1552        *vm.handler.write().await = Some(Box::new(FailingHandler));
1553
1554        let box_dir = tmp.path().join("boxes").join(&box_id);
1555        std::fs::create_dir_all(box_dir.join("logs")).unwrap();
1556
1557        let result = vm.destroy_with_options(default_stop_signal(), 100).await;
1558
1559        // The stop failure is still surfaced to the caller...
1560        assert!(
1561            result.is_err(),
1562            "a handler-stop failure must still be reported"
1563        );
1564        // ...but the host teardown ran anyway: handler taken + box dir removed.
1565        assert!(vm.handler.read().await.is_none());
1566        assert!(
1567            !box_dir.exists(),
1568            "non-persistent box dir must be removed even when the stop failed"
1569        );
1570    }
1571
1572    #[tokio::test]
1573    async fn test_cleanup_boot_failure_stops_handler_and_removes_created_volumes() {
1574        let tmp = tempfile::tempdir().unwrap();
1575        let box_id = "box-test".to_string();
1576        let mut vm =
1577            VmManager::with_box_id(BoxConfig::default(), EventEmitter::new(16), box_id.clone());
1578        vm.home_dir = tmp.path().to_path_buf();
1579        vm.anonymous_volumes = vec!["created-volume".to_string(), "reused-volume".to_string()];
1580        vm.created_anonymous_volumes = vec!["created-volume".to_string()];
1581
1582        let stopped = Arc::new(AtomicBool::new(false));
1583        *vm.handler.write().await = Some(Box::new(RecordingHandler {
1584            stopped: stopped.clone(),
1585        }));
1586
1587        let box_dir = tmp.path().join("boxes").join(&box_id);
1588        std::fs::create_dir_all(box_dir.join("logs")).unwrap();
1589
1590        let store = crate::volume::VolumeStore::new(
1591            tmp.path().join("volumes.json"),
1592            tmp.path().join("volumes"),
1593        );
1594        store
1595            .create(a3s_box_core::volume::VolumeConfig::new(
1596                "created-volume",
1597                "",
1598            ))
1599            .unwrap();
1600        store
1601            .create(a3s_box_core::volume::VolumeConfig::new("reused-volume", ""))
1602            .unwrap();
1603
1604        vm.cleanup_boot_failure().await;
1605
1606        assert!(stopped.load(Ordering::SeqCst));
1607        assert!(vm.handler.read().await.is_none());
1608        assert!(vm.created_anonymous_volumes.is_empty());
1609        assert_eq!(vm.anonymous_volumes, vec!["reused-volume".to_string()]);
1610        assert!(store.get("created-volume").unwrap().is_none());
1611        assert!(store.get("reused-volume").unwrap().is_some());
1612        assert!(!box_dir.exists());
1613    }
1614
1615    #[tokio::test]
1616    async fn test_wait_for_vm_running_returns_error_when_handler_exited() {
1617        let vm = VmManager::with_box_id(
1618            BoxConfig::default(),
1619            EventEmitter::new(16),
1620            "box-exited".to_string(),
1621        );
1622        *vm.handler.write().await = Some(Box::new(ExitStateHandler { exited: true }));
1623
1624        let err = vm.wait_for_vm_running().await.unwrap_err();
1625
1626        assert!(err
1627            .to_string()
1628            .contains("VM process exited immediately after start"));
1629    }
1630
1631    #[tokio::test]
1632    async fn test_wait_for_vm_running_succeeds_when_handler_stays_running() {
1633        let config = BoxConfig {
1634            restore_from: Some("snapshot-path".to_string()),
1635            ..BoxConfig::default()
1636        };
1637        let vm = VmManager::with_box_id(config, EventEmitter::new(16), "box-running".to_string());
1638        *vm.handler.write().await = Some(Box::new(ExitStateHandler { exited: false }));
1639
1640        vm.wait_for_vm_running().await.unwrap();
1641    }
1642
1643    #[tokio::test]
1644    async fn test_try_wait_exit_reads_guest_persisted_exit_code() {
1645        let tmp = tempfile::tempdir().unwrap();
1646        let box_id = "box-exit-code".to_string();
1647        let mut vm =
1648            VmManager::with_box_id(BoxConfig::default(), EventEmitter::new(16), box_id.clone());
1649        vm.home_dir = tmp.path().to_path_buf();
1650
1651        let exit_path = tmp
1652            .path()
1653            .join("boxes")
1654            .join(&box_id)
1655            .join("upper")
1656            .join(".a3s_exit_code");
1657        std::fs::create_dir_all(exit_path.parent().unwrap()).unwrap();
1658        std::fs::write(&exit_path, "42\n").unwrap();
1659
1660        assert_eq!(vm.try_wait_exit().await.unwrap(), Some(42));
1661        assert_eq!(vm.exit_code(), Some(42));
1662        assert!(vm.has_exited().await);
1663    }
1664
1665    #[cfg(unix)]
1666    #[tokio::test]
1667    async fn test_wait_for_exec_ready_returns_when_handler_already_exited() {
1668        let mut vm = VmManager::with_box_id(
1669            BoxConfig::default(),
1670            EventEmitter::new(16),
1671            "box-exec-exited".to_string(),
1672        );
1673        *vm.handler.write().await = Some(Box::new(ExitStateHandler { exited: true }));
1674        let tmp = tempfile::tempdir().unwrap();
1675
1676        vm.wait_for_exec_ready(&tmp.path().join("missing-exec.sock"))
1677            .await
1678            .unwrap();
1679
1680        assert!(vm.exec_client.is_none());
1681    }
1682
1683    #[cfg(unix)]
1684    #[tokio::test]
1685    async fn test_wait_for_exec_ready_returns_when_guest_exit_code_persisted() {
1686        let tmp = tempfile::tempdir().unwrap();
1687        let box_id = "box-exec-finished".to_string();
1688        let mut vm =
1689            VmManager::with_box_id(BoxConfig::default(), EventEmitter::new(16), box_id.clone());
1690        vm.home_dir = tmp.path().to_path_buf();
1691
1692        let exit_path = tmp
1693            .path()
1694            .join("boxes")
1695            .join(&box_id)
1696            .join("upper")
1697            .join(".a3s_exit_code");
1698        std::fs::create_dir_all(exit_path.parent().unwrap()).unwrap();
1699        std::fs::write(&exit_path, "17\n").unwrap();
1700
1701        tokio::time::timeout(
1702            std::time::Duration::from_secs(1),
1703            vm.wait_for_exec_ready(&tmp.path().join("missing-exec.sock")),
1704        )
1705        .await
1706        .unwrap()
1707        .unwrap();
1708
1709        assert_eq!(vm.exit_code(), Some(17));
1710        assert!(vm.exec_client.is_none());
1711    }
1712
1713    #[cfg(unix)]
1714    #[tokio::test]
1715    async fn test_probe_exec_ready_once_ignores_missing_socket() {
1716        let mut vm = VmManager::with_box_id(
1717            BoxConfig::default(),
1718            EventEmitter::new(16),
1719            "box-probe".to_string(),
1720        );
1721        let tmp = tempfile::tempdir().unwrap();
1722
1723        vm.probe_exec_ready_once(&tmp.path().join("missing-exec.sock"))
1724            .await;
1725
1726        assert!(vm.exec_client.is_none());
1727    }
1728
1729    #[cfg(unix)]
1730    #[tokio::test]
1731    async fn test_attach_running_process_infers_port_forward_socket_path() {
1732        let mut vm = VmManager::with_box_id(
1733            BoxConfig::default(),
1734            EventEmitter::new(16),
1735            "box-test".to_string(),
1736        );
1737        let tmp = tempfile::tempdir().unwrap();
1738        let exec_socket_path = tmp.path().join("exec.sock");
1739        let pty_socket_path = Some(tmp.path().join("pty.sock"));
1740
1741        vm.attach_running_process(
1742            std::process::id(),
1743            exec_socket_path.clone(),
1744            pty_socket_path.clone(),
1745        )
1746        .await
1747        .unwrap();
1748
1749        assert_eq!(vm.exec_socket_path(), Some(exec_socket_path.as_path()));
1750        assert_eq!(vm.pty_socket_path(), pty_socket_path.as_deref());
1751        assert_eq!(
1752            vm.port_forward_socket_path(),
1753            Some(exec_socket_path.with_file_name("portfwd.sock").as_path())
1754        );
1755        assert_eq!(vm.state().await, BoxState::Ready);
1756    }
1757}