From cfda47b31682b7610c15df63046f651035da72e9 Mon Sep 17 00:00:00 2001
From: He Sun <bochang.sh@antfin.com>
Date: Mon, 13 Jul 2020 21:49:39 +0800
Subject: [PATCH] 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.
---
 src/libos/src/sched/cpu_set.rs  | 7 ++++++-
 src/libos/src/sched/syscalls.rs | 6 +++---
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/libos/src/sched/cpu_set.rs b/src/libos/src/sched/cpu_set.rs
index 4a5d0628..9a9ffcc4 100644
--- a/src/libos/src/sched/cpu_set.rs
+++ b/src/libos/src/sched/cpu_set.rs
@@ -52,6 +52,11 @@ impl CpuSet {
         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
     pub fn is_subset_of(&self, other: &CpuSet) -> bool {
         (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
     /// the host machine.
     ///
-    /// Property: `AVAIL_CPU.empty() == false`.
+    /// Property: `AVAIL_CPUSET.empty() == false`.
     pub static ref AVAIL_CPUSET: CpuSet = {
         extern "C" {
             fn occlum_ocall_sched_getaffinity(
diff --git a/src/libos/src/sched/syscalls.rs b/src/libos/src/sched/syscalls.rs
index 3d5778e6..bb96340e 100644
--- a/src/libos/src/sched/syscalls.rs
+++ b/src/libos/src/sched/syscalls.rs
@@ -1,4 +1,4 @@
-use super::cpu_set::CpuSet;
+use super::cpu_set::{CpuSet, AVAIL_CPUSET};
 use crate::prelude::*;
 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> {
     // Construct safe Rust types
     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");
         }
 
@@ -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> {
     // Convert unsafe C types into safe Rust types
     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");
         }
         CpuSet::len()