Checks if CPU count is valid for KVM and Hypervisor (obhq#742)

src/hv/src/darwin.rs

@@ -12,6 +12,15 @@ impl Vm {
             v => Err(NewError::CreateVmFailed(v)),
         }
     }
+
+    pub fn capability(&self, id: u64) -> Result<u64, c_int> {
+        let mut value = 0;
+
+        match unsafe { hv_capability(id, &mut value) } {
+            0 => Ok(value),
+            v => Err(v),
+        }
+    }
 }
 
 impl Drop for Vm {
@@ -27,4 +36,5 @@ impl Drop for Vm {
 extern "C" {
     fn hv_vm_create(config: *mut ()) -> c_int;
     fn hv_vm_destroy() -> c_int;
+    fn hv_capability(capability: u64, value: *mut u64) -> c_int;
 }

src/hv/src/lib.rs

@@ -55,26 +55,34 @@ impl Hypervisor {
 
         // Initialize platform hypervisor.
         #[cfg(any(target_os = "linux", target_os = "android"))]
-        return Self::new_linux(active, ram, addr, len.get());
+        return Self::new_linux(active, cpu, ram, addr, len.get());
 
         #[cfg(target_os = "windows")]
         return Self::new_windows(active, cpu, ram, addr, len.get());
 
         #[cfg(target_os = "macos")]
-        return Self::new_mac(active, ram, addr, len.get());
+        return Self::new_mac(active, cpu, ram, addr, len.get());
     }
 
     #[cfg(any(target_os = "linux", target_os = "android"))]
     unsafe fn new_linux(
         active: Active,
+        cpu: NonZeroUsize,
         ram: *mut u8,
         addr: usize,
         len: usize,
     ) -> Result<Self, NewError> {
         use std::os::fd::AsFd;
 
+        // Open KVM device.
         let kvm = self::linux::open_kvm()?;
-        let vm = self::linux::create_vm(kvm.as_fd())?;
+
+        if cpu.get() > self::linux::max_vcpus(kvm.as_fd()).map_err(NewError::GetMaxCpuFailed)? {
+            return Err(NewError::InvalidCpuCount);
+        }
+
+        // Create a new VM.
+        let vm = self::linux::create_vm(kvm.as_fd()).map_err(NewError::CreateVmFailed)?;
 
         Ok(Self { vm, kvm, active })
     }
@@ -97,11 +105,18 @@ impl Hypervisor {
     #[cfg(target_os = "macos")]
     unsafe fn new_mac(
         active: Active,
+        cpu: NonZeroUsize,
         ram: *mut u8,
         addr: usize,
         len: usize,
     ) -> Result<Self, NewError> {
+        // Create a VM.
         let vm = self::darwin::Vm::new()?;
+        let cpu: u64 = cpu.get().try_into().unwrap();
+
+        if cpu > vm.capability(0).map_err(NewError::GetMaxCpuFailed)? {
+            return Err(NewError::InvalidCpuCount);
+        }
 
         Ok(Self { vm, active })
     }
@@ -156,7 +171,6 @@ pub enum NewError {
     #[error("couldn't determine page size of the host")]
     GetHostPageSizeFailed(#[source] std::io::Error),
 
-    #[cfg(target_os = "windows")]
     #[error("the number of CPU is not valid")]
     InvalidCpuCount,
 
@@ -175,6 +189,10 @@ pub enum NewError {
     #[error("unexpected KVM version")]
     KvmVersionMismatched,
 
+    #[cfg(any(target_os = "linux", target_os = "android"))]
+    #[error("couldn't get maximum number of CPU for a VM")]
+    GetMaxCpuFailed(#[source] std::io::Error),
+
     #[cfg(any(target_os = "linux", target_os = "android"))]
     #[error("couldn't create a VM")]
     CreateVmFailed(#[source] std::io::Error),
@@ -194,6 +212,10 @@ pub enum NewError {
     #[cfg(target_os = "macos")]
     #[error("couldn't create a VM ({0:#x})")]
     CreateVmFailed(std::ffi::c_int),
+
+    #[cfg(target_os = "macos")]
+    #[error("couldn't get maximum number of CPU for a VM")]
+    GetMaxCpuFailed(std::ffi::c_int),
 }
 
 static ACTIVE: AtomicBool = AtomicBool::new(false);

src/hv/src/linux/kvm.cpp

@@ -1,8 +1,10 @@
+// The reason we need C++ here is because the KVM ioctls is not available in the libc binding.
 #include <linux/kvm.h>
 
 #include <sys/ioctl.h>
 
 #include <errno.h>
+#include <stddef.h>
 
 extern "C" int kvm_check_version(int kvm, bool *compat)
 {
@@ -16,6 +18,18 @@ extern "C" int kvm_check_version(int kvm, bool *compat)
     return 0;
 }
 
+extern "C" int kvm_max_vcpus(int kvm, size_t *max)
+{
+    auto num = ioctl(kvm, KVM_CHECK_EXTENSION, KVM_CAP_MAX_VCPUS);
+
+    if (num < 0) {
+        return errno;
+    }
+
+    *max = static_cast<size_t>(num);
+    return 0;
+}
+
 extern "C" int kvm_create_vm(int kvm, int *fd)
 {
     auto vm = ioctl(kvm, KVM_CREATE_VM, 0);

src/hv/src/linux/mod.rs

@@ -28,16 +28,26 @@ pub fn open_kvm() -> Result<OwnedFd, NewError> {
     Ok(fd)
 }
 
-pub fn create_vm(kvm: BorrowedFd) -> Result<OwnedFd, NewError> {
+pub fn max_vcpus(kvm: BorrowedFd) -> Result<usize, Error> {
+    let mut max = 0;
+
+    match unsafe { kvm_max_vcpus(kvm.as_raw_fd(), &mut max) } {
+        0 => Ok(max),
+        v => Err(Error::from_raw_os_error(v)),
+    }
+}
+
+pub fn create_vm(kvm: BorrowedFd) -> Result<OwnedFd, Error> {
     let mut vm = -1;
 
     match unsafe { kvm_create_vm(kvm.as_raw_fd(), &mut vm) } {
         0 => Ok(unsafe { OwnedFd::from_raw_fd(vm) }),
-        v => Err(NewError::CreateVmFailed(Error::from_raw_os_error(v))),
+        v => Err(Error::from_raw_os_error(v)),
     }
 }
 
 extern "C" {
     fn kvm_check_version(kvm: c_int, compat: *mut bool) -> c_int;
+    fn kvm_max_vcpus(kvm: c_int, max: *mut usize) -> c_int;
     fn kvm_create_vm(kvm: c_int, fd: *mut c_int) -> c_int;
 }