Check the input buffer size against the available CPUs in sched_get/setaffinity

There are scenarios where the available CPUs are less than all the CPUs
on the machine. Therefore, sched_get/setaffinity should be allowed when
the input buffer size is no less than the available CPUs but less than
all the CPUs.
This commit is contained in:
He Sun 2020-07-13 21:49:39 +08:00 committed by zongmin.gzm
parent 306c0333ad
commit cfda47b316
2 changed files with 9 additions and 4 deletions

@ -52,6 +52,11 @@ impl CpuSet {
self.bits.count_ones() == 0 self.bits.count_ones() == 0
} }
/// Returns the number of CPUs in set.
pub fn cpu_count(&self) -> usize {
self.bits.count_ones()
}
// Returns if the CpuSet is a subset of available cpu set // Returns if the CpuSet is a subset of available cpu set
pub fn is_subset_of(&self, other: &CpuSet) -> bool { pub fn is_subset_of(&self, other: &CpuSet) -> bool {
(self.bits.clone() & other.bits.clone()) == self.bits (self.bits.clone() & other.bits.clone()) == self.bits
@ -137,7 +142,7 @@ lazy_static! {
/// cloud platform, the container or vm is usually given access to a subset of the CPU cores on /// cloud platform, the container or vm is usually given access to a subset of the CPU cores on
/// the host machine. /// the host machine.
/// ///
/// Property: `AVAIL_CPU.empty() == false`. /// Property: `AVAIL_CPUSET.empty() == false`.
pub static ref AVAIL_CPUSET: CpuSet = { pub static ref AVAIL_CPUSET: CpuSet = {
extern "C" { extern "C" {
fn occlum_ocall_sched_getaffinity( fn occlum_ocall_sched_getaffinity(

@ -1,4 +1,4 @@
use super::cpu_set::CpuSet; use super::cpu_set::{CpuSet, AVAIL_CPUSET};
use crate::prelude::*; use crate::prelude::*;
use crate::util::mem_util::from_user::*; use crate::util::mem_util::from_user::*;
@ -10,7 +10,7 @@ pub fn do_sched_yield() -> Result<isize> {
pub fn do_sched_getaffinity(pid: pid_t, buf_size: size_t, buf_ptr: *mut u8) -> Result<isize> { pub fn do_sched_getaffinity(pid: pid_t, buf_size: size_t, buf_ptr: *mut u8) -> Result<isize> {
// Construct safe Rust types // Construct safe Rust types
let buf_size = { let buf_size = {
if buf_size < CpuSet::len() { if buf_size * 8 < AVAIL_CPUSET.cpu_count() {
return_errno!(EINVAL, "buf size is not big enough"); return_errno!(EINVAL, "buf size is not big enough");
} }
@ -39,7 +39,7 @@ pub fn do_sched_getaffinity(pid: pid_t, buf_size: size_t, buf_ptr: *mut u8) -> R
pub fn do_sched_setaffinity(pid: pid_t, buf_size: size_t, buf_ptr: *const u8) -> Result<isize> { pub fn do_sched_setaffinity(pid: pid_t, buf_size: size_t, buf_ptr: *const u8) -> Result<isize> {
// Convert unsafe C types into safe Rust types // Convert unsafe C types into safe Rust types
let buf_size = { let buf_size = {
if buf_size < CpuSet::len() { if buf_size * 8 < AVAIL_CPUSET.cpu_count() {
return_errno!(EINVAL, "buf size is not big enough"); return_errno!(EINVAL, "buf size is not big enough");
} }
CpuSet::len() CpuSet::len()