From 380ba564f4d8ee5d457d8be256436827e6693ebc Mon Sep 17 00:00:00 2001
From: Praveen K Paladugu <prapal@linux.microsoft.com>
Date: Thu, 2 May 2024 20:02:25 +0000
Subject: [PATCH] vmm: populate console_info during vm actions

Use pre_create_console_devices method to create and populate console
device FDs into console_info in Vmm Object.

Signed-off-by: Praveen K Paladugu <prapal@linux.microsoft.com>
---
 vmm/src/console_devices.rs |  1 -
 vmm/src/lib.rs             | 28 ++++++++++++++++++++++++++++
 vmm/src/vm.rs              |  4 ++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/vmm/src/console_devices.rs b/vmm/src/console_devices.rs
index 5462d1fb5..93f581171 100644
--- a/vmm/src/console_devices.rs
+++ b/vmm/src/console_devices.rs
@@ -166,7 +166,6 @@ pub fn create_pty() -> io::Result<(File, File, PathBuf)> {
     Ok((main, unsafe { File::from_raw_fd(sub_fd) }, path))
 }
 
-#[allow(dead_code)]
 pub(crate) fn pre_create_console_devices(vmm: &mut Vmm) -> ConsoleDeviceResult<ConsoleInfo> {
     let vm_config = vmm.vm_config.as_mut().unwrap().clone();
     let mut vmconfig = vm_config.lock().unwrap();
diff --git a/vmm/src/lib.rs b/vmm/src/lib.rs
index f442f1db6..6dfb088af 100644
--- a/vmm/src/lib.rs
+++ b/vmm/src/lib.rs
@@ -28,6 +28,7 @@ use anyhow::anyhow;
 #[cfg(feature = "dbus_api")]
 use api::dbus::{DBusApiOptions, DBusApiShutdownChannels};
 use api::http::HttpApiHandle;
+use console_devices::{pre_create_console_devices, ConsoleInfo};
 use libc::{tcsetattr, termios, EFD_NONBLOCK, SIGINT, SIGTERM, TCSANOW};
 use memory_manager::MemoryManagerSnapshotData;
 use pci::PciBdf;
@@ -550,6 +551,7 @@ pub struct Vmm {
     threads: Vec<thread::JoinHandle<()>>,
     original_termios_opt: Arc<Mutex<Option<termios>>>,
     console_resize_pipe: Option<File>,
+    console_info: Option<ConsoleInfo>,
 }
 
 impl Vmm {
@@ -685,6 +687,7 @@ impl Vmm {
             threads: vec![],
             original_termios_opt: Arc::new(Mutex::new(None)),
             console_resize_pipe: None,
+            console_info: None,
         })
     }
 
@@ -1211,6 +1214,8 @@ impl RequestHandler for Vmm {
         // The VM will be created when being asked to boot it.
         if self.vm_config.is_none() {
             self.vm_config = Some(config);
+            self.console_info =
+                Some(pre_create_console_devices(self).map_err(VmError::CreateConsoleDevices)?);
             Ok(())
         } else {
             Err(VmError::VmAlreadyCreated)
@@ -1228,6 +1233,12 @@ impl RequestHandler for Vmm {
                 return Err(VmError::VmMissingConfig);
             };
 
+            // console_info is set to None in vm_shutdown. re-populate here if empty
+            if self.console_info.is_none() {
+                self.console_info =
+                    Some(pre_create_console_devices(self).map_err(VmError::CreateConsoleDevices)?);
+            }
+
             // Create a new VM if we don't have one yet.
             if self.vm.is_none() {
                 let exit_evt = self.exit_evt.try_clone().map_err(VmError::EventFdClone)?;
@@ -1298,6 +1309,8 @@ impl RequestHandler for Vmm {
 
     fn vm_snapshot(&mut self, destination_url: &str) -> result::Result<(), VmError> {
         if let Some(ref mut vm) = self.vm {
+            // Drain console_info so that FDs are not reused
+            let _ = self.console_info.take();
             vm.snapshot()
                 .map_err(VmError::Snapshot)
                 .and_then(|snapshot| {
@@ -1352,6 +1365,12 @@ impl RequestHandler for Vmm {
 
         self.vm_config = Some(Arc::clone(&vm_config));
 
+        // console_info is set to None in vm_snapshot. re-populate here if empty
+        if self.console_info.is_none() {
+            self.console_info =
+                Some(pre_create_console_devices(self).map_err(VmError::CreateConsoleDevices)?);
+        }
+
         let exit_evt = self.exit_evt.try_clone().map_err(VmError::EventFdClone)?;
         let reset_evt = self.reset_evt.try_clone().map_err(VmError::EventFdClone)?;
         #[cfg(feature = "guest_debug")]
@@ -1403,6 +1422,8 @@ impl RequestHandler for Vmm {
 
     fn vm_shutdown(&mut self) -> result::Result<(), VmError> {
         let r = if let Some(ref mut vm) = self.vm.take() {
+            // Drain console_info so that the FDs are not reused
+            let _ = self.console_info.take();
             vm.shutdown()
         } else {
             Err(VmError::VmNotRunning)
@@ -1441,6 +1462,10 @@ impl RequestHandler for Vmm {
                 return Err(VmError::VmNotCreated);
             };
 
+        // vm.shutdown() closes all the console devices, so set console_info to None
+        // so that the closed FD #s are not reused.
+        let _ = self.console_info.take();
+
         let exit_evt = self.exit_evt.try_clone().map_err(VmError::EventFdClone)?;
         let reset_evt = self.reset_evt.try_clone().map_err(VmError::EventFdClone)?;
         #[cfg(feature = "guest_debug")]
@@ -1460,6 +1485,9 @@ impl RequestHandler for Vmm {
             warn!("Spurious second reset event received. Ignoring.");
         }
 
+        self.console_info =
+            Some(pre_create_console_devices(self).map_err(VmError::CreateConsoleDevices)?);
+
         // Then we create the new VM
         let mut vm = Vm::new(
             config,
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index 6c183f807..338a5274c 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -16,6 +16,7 @@ use crate::config::{
     UserDeviceConfig, ValidationError, VdpaConfig, VmConfig, VsockConfig,
 };
 use crate::config::{NumaConfig, PayloadConfig};
+use crate::console_devices::ConsoleDeviceError;
 #[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
 use crate::coredump::{
     CpuElf64Writable, DumpState, Elf64Writable, GuestDebuggable, GuestDebuggableError, NoteDescType,
@@ -322,6 +323,9 @@ pub enum Error {
 
     #[error("Error resuming the VM: {0}")]
     ResumeVm(#[source] hypervisor::HypervisorVmError),
+
+    #[error("Error creating console devices")]
+    CreateConsoleDevices(ConsoleDeviceError),
 }
 pub type Result<T> = result::Result<T, Error>;