From 76d5bb52439924b0349da238856e07b620c60877 Mon Sep 17 00:00:00 2001 From: Putta Khunchalee Date: Sat, 16 Mar 2024 22:59:19 +0700 Subject: [PATCH] Checks if CPU count is valid for KVM and Hypervisor (#742) --- src/hv/src/darwin.rs | 10 ++++++++++ src/hv/src/lib.rs | 30 ++++++++++++++++++++++++++---- src/hv/src/linux/kvm.cpp | 14 ++++++++++++++ src/hv/src/linux/mod.rs | 14 ++++++++++++-- 4 files changed, 62 insertions(+), 6 deletions(-) diff --git a/src/hv/src/darwin.rs b/src/hv/src/darwin.rs index bd5d84922..990fa55ff 100644 --- a/src/hv/src/darwin.rs +++ b/src/hv/src/darwin.rs @@ -12,6 +12,15 @@ impl Vm { v => Err(NewError::CreateVmFailed(v)), } } + + pub fn capability(&self, id: u64) -> Result { + 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; } diff --git a/src/hv/src/lib.rs b/src/hv/src/lib.rs index 89e6a2144..7613f8e1f 100644 --- a/src/hv/src/lib.rs +++ b/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 { 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 { + // 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); diff --git a/src/hv/src/linux/kvm.cpp b/src/hv/src/linux/kvm.cpp index 72cfc6c50..bfd6976cc 100644 --- a/src/hv/src/linux/kvm.cpp +++ b/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 #include #include +#include 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(num); + return 0; +} + extern "C" int kvm_create_vm(int kvm, int *fd) { auto vm = ioctl(kvm, KVM_CREATE_VM, 0); diff --git a/src/hv/src/linux/mod.rs b/src/hv/src/linux/mod.rs index aa0862e0f..394858fe7 100644 --- a/src/hv/src/linux/mod.rs +++ b/src/hv/src/linux/mod.rs @@ -28,16 +28,26 @@ pub fn open_kvm() -> Result { Ok(fd) } -pub fn create_vm(kvm: BorrowedFd) -> Result { +pub fn max_vcpus(kvm: BorrowedFd) -> Result { + 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 { 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; }