From 34bb3319d489b9c0120d8414096581d05cad58c1 Mon Sep 17 00:00:00 2001
From: Rob Bradford <rbradford@rivosinc.com>
Date: Tue, 11 Apr 2023 15:52:01 +0100
Subject: [PATCH] hypervisor, vmm: Limit max number of vCPUs to hypervisor
 maximum

On KVM this is provided by an ioctl, on MSHV this is constant. Although
there is a HV_MAXIMUM_PROCESSORS constant the MSHV ioctl API is limited
to u8.

Signed-off-by: Rob Bradford <rbradford@rivosinc.com>
---
 hypervisor/src/hypervisor.rs | 3 +++
 hypervisor/src/kvm/mod.rs    | 5 +++++
 hypervisor/src/mshv/mod.rs   | 7 +++++++
 vmm/src/cpu.rs               | 7 +++++++
 4 files changed, 22 insertions(+)

diff --git a/hypervisor/src/hypervisor.rs b/hypervisor/src/hypervisor.rs
index fb93a622e..778ba7b95 100644
--- a/hypervisor/src/hypervisor.rs
+++ b/hypervisor/src/hypervisor.rs
@@ -133,4 +133,7 @@ pub trait Hypervisor: Send + Sync {
     fn get_guest_debug_hw_bps(&self) -> usize {
         unimplemented!()
     }
+
+    /// Get maximum number of vCPUs
+    fn get_max_vcpus(&self) -> u32;
 }
diff --git a/hypervisor/src/kvm/mod.rs b/hypervisor/src/kvm/mod.rs
index f3e33f586..3aa868579 100644
--- a/hypervisor/src/kvm/mod.rs
+++ b/hypervisor/src/kvm/mod.rs
@@ -1084,6 +1084,11 @@ impl hypervisor::Hypervisor for KvmHypervisor {
             self.kvm.get_guest_debug_hw_bps() as usize
         }
     }
+
+    /// Get maximum number of vCPUs
+    fn get_max_vcpus(&self) -> u32 {
+        self.kvm.get_max_vcpus().min(u32::MAX as usize) as u32
+    }
 }
 /// Vcpu struct for KVM
 pub struct KvmVcpu {
diff --git a/hypervisor/src/mshv/mod.rs b/hypervisor/src/mshv/mod.rs
index b048a4a4e..8dc987aac 100644
--- a/hypervisor/src/mshv/mod.rs
+++ b/hypervisor/src/mshv/mod.rs
@@ -279,6 +279,13 @@ impl hypervisor::Hypervisor for MshvHypervisor {
     fn get_supported_cpuid(&self) -> hypervisor::Result<Vec<CpuIdEntry>> {
         Ok(Vec::new())
     }
+
+    /// Get maximum number of vCPUs
+    fn get_max_vcpus(&self) -> u32 {
+        // TODO: Using HV_MAXIMUM_PROCESSORS would be better
+        // but the ioctl API is limited to u8
+        256
+    }
 }
 
 /// Vcpu struct for Microsoft Hypervisor
diff --git a/vmm/src/cpu.rs b/vmm/src/cpu.rs
index 4775aff5d..afb247841 100644
--- a/vmm/src/cpu.rs
+++ b/vmm/src/cpu.rs
@@ -167,6 +167,9 @@ pub enum Error {
     #[cfg(target_arch = "x86_64")]
     #[error("Error setting up AMX: {0}")]
     AmxEnable(#[source] anyhow::Error),
+
+    #[error("Maximum number of vCPUs exceeds host limit")]
+    MaximumVcpusExceeded,
 }
 pub type Result<T> = result::Result<T, Error>;
 
@@ -589,6 +592,10 @@ impl CpuManager {
         #[cfg(feature = "tdx")] tdx_enabled: bool,
         numa_nodes: &NumaNodes,
     ) -> Result<Arc<Mutex<CpuManager>>> {
+        if u32::from(config.max_vcpus) > hypervisor.get_max_vcpus() {
+            return Err(Error::MaximumVcpusExceeded);
+        }
+
         let mut vcpu_states = Vec::with_capacity(usize::from(config.max_vcpus));
         vcpu_states.resize_with(usize::from(config.max_vcpus), VcpuState::default);
         let hypervisor_type = hypervisor.hypervisor_type();