Make sure we will always pick a valid CPU

When tid % allowed.NumCPUs() equals 0, e.g. when there is only one
CPU allowed for this thread, we will always pick CPU 0 even if it's
not in the allowed list. This patch fixes this issue.

Fixes #3133

Signed-off-by: Tiwei Bie <tiwei.btw@antgroup.com>

pkg/sentry/kernel/task_sched.go

@@ -621,9 +621,10 @@ func assignCPU(allowed sched.CPUSet, tid ThreadID) (cpu int32) {
 	n := int(tid) % int(allowed.NumCPUs())
 	// ... then pick the nth CPU in allowed.
 	allowed.ForEachCPU(func(c uint) {
-		if n--; n == 0 {
+		if n == 0 {
 			cpu = int32(c)
 		}
+		n--
 	})
 	return cpu
 }

pkg/sentry/kernel/task_test.go

@@ -29,35 +29,41 @@ func TestTaskCPU(t *testing.T) {
 		{
 			mask: []byte{0xff},
 			tid:  1,
-			cpu:  0,
+			cpu:  1,
 		},
 		{
 			mask: []byte{0xff},
 			tid:  10,
-			cpu:  1,
+			cpu:  2,
 		},
 		{
 			// more than 8 cpus.
 			mask: []byte{0xff, 0xff},
 			tid:  10,
-			cpu:  9,
+			cpu:  10,
 		},
 		{
 			// missing the first cpu.
 			mask: []byte{0xfe},
 			tid:  1,
-			cpu:  1,
+			cpu:  2,
 		},
 		{
 			mask: []byte{0xfe},
 			tid:  10,
-			cpu:  3,
+			cpu:  4,
 		},
 		{
 			// missing the fifth cpu.
 			mask: []byte{0xef},
 			tid:  10,
-			cpu:  2,
+			cpu:  3,
+		},
+		{
+			// only the fifth cpu.
+			mask: []byte{0x10},
+			tid:  10,
+			cpu:  4,
 		},
 	} {
 		assigned := assignCPU(test.mask, test.tid)

test/syscalls/linux/getcpu.cc

@@ -34,6 +34,19 @@ TEST(GetcpuTest, IsValidCpuStress) {
   }
 }
 
+TEST(GetcpuTest, IsValidCpu) {
+  const int num_cpus = NumCPUs();
+  for (int i = 0; i < num_cpus; i++) {
+    cpu_set_t set = {};
+    int cpu;
+    CPU_SET(i, &set);
+    ASSERT_THAT(sched_setaffinity(getpid(), sizeof(set), &set),
+                SyscallSucceeds());
+    ASSERT_THAT(cpu = sched_getcpu(), SyscallSucceeds());
+    ASSERT_EQ(cpu, i);
+  }
+}
+
 }  // namespace
 
 }  // namespace testing