Fix checkptr error with > 1 process in job object (#1284)

The `AllPids` method on JOBOBJECT_BASIC_PROCESS_ID_LIST would allocate a
very large array to store any pids in the job object, cast the memory to
this array, and then slice down to what elements it actually took up in
the array based off the value of what was in the NumberOfProcessIdsInList
field. The checkptr compile option doesn't like when slices don't have
an explicit length and capacity so this change updates the slicing to use
a three-index slice to set the capacity to the same as the length.

Before for fmt.Println(len(arr), cap(arr)) -> 2, 134217727
After: -> 2, 2

This change additionally adds a test in internal/jobobject and modifies the
TestExecWithJob test in internal/exec to verify that checkptr doesn't get angry
when we hit a codepath that performs the cast

Signed-off-by: Daniel Canter <[email protected]>
(cherry picked from commit d082725)
Signed-off-by: Daniel Canter <[email protected]>

internal/exec/exec_test.go

@@ -130,8 +130,8 @@ func TestExecStdinPowershell(t *testing.T) {
 	t.Logf("exit code was: %d", e.ExitCode())
 }
 
-func TestExecWithJob(t *testing.T) {
-	// Test that we can assign a process to a job object at creation time.
+func TestExecsWithJob(t *testing.T) {
+	// Test that we can assign processes to a job object at creation time.
 	job, err := jobobject.Create(context.Background(), &jobobject.Options{Name: "test"})
 	if err != nil {
 		log.Fatal(err)
@@ -140,9 +140,9 @@ func TestExecWithJob(t *testing.T) {
 
 	e, err := New(
 		`C:\Windows\System32\ping.exe`,
-		"ping 127.0.0.1",
+		"ping -t 127.0.0.1",
 		WithJobObject(job),
-		WithStdio(true, false, false),
+		WithStdio(false, false, false),
 		WithEnv(os.Environ()),
 	)
 	if err != nil {
@@ -154,25 +154,62 @@ func TestExecWithJob(t *testing.T) {
 		t.Fatalf("failed to start process: %v", err)
 	}
 
-	pids, err := job.Pids()
+	// Launch a second process and check pids.
+	e2, err := New(
+		`C:\Windows\System32\ping.exe`,
+		"ping -t 127.0.0.1",
+		WithJobObject(job),
+		WithStdio(false, false, false),
+		WithEnv(os.Environ()),
+	)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	if len(pids) == 0 {
-		t.Fatal("no pids found in job object")
+	err = e2.Start()
+	if err != nil {
+		t.Fatalf("failed to start process: %v", err)
+	}
+
+	pidMap := map[int]struct{}{
+		e.Pid():  {},
+		e2.Pid(): {},
 	}
 
-	// Should only be one process in the job, being the process we launched and added to it.
-	if pids[0] != uint32(e.Pid()) {
+	pids, err := job.Pids()
+	if err != nil {
 		t.Fatal(err)
 	}
 
-	err = e.Wait()
+	if len(pids) != 2 {
+		t.Fatalf("should be two pids in job object, got: %d", len(pids))
+	}
+
+	for _, pid := range pids {
+		if _, ok := pidMap[int(pid)]; !ok {
+			t.Fatalf("failed to find pid %d in job object", pid)
+		}
+	}
+
+	err = e.Kill()
 	if err != nil {
-		t.Fatalf("error waiting for process: %v", err)
+		t.Fatalf("error killing process: %v", err)
+	}
+
+	err = e2.Kill()
+	if err != nil {
+		t.Fatalf("error killing process: %v", err)
+	}
+
+	_ = e.Wait()
+	_ = e2.Wait()
+
+	if !e.Exited() {
+		t.Fatalf("Process %d should have exited after kill", e.Pid())
+	}
+	if !e2.Exited() {
+		t.Fatalf("Process %d should have exited after kill", e2.Pid())
 	}
-	t.Logf("exit code was: %d", e.ExitCode())
 }
 
 func TestPseudoConsolePowershell(t *testing.T) {

internal/jobobject/jobobject.go

@@ -364,9 +364,8 @@ func (job *JobObject) Pids() ([]uint32, error) {
 	}
 
 	bufInfo := (*winapi.JOBOBJECT_BASIC_PROCESS_ID_LIST)(unsafe.Pointer(&buf[0]))
-	bufPids := bufInfo.AllPids()
 	pids := make([]uint32, bufInfo.NumberOfProcessIdsInList)
-	for i, bufPid := range bufPids {
+	for i, bufPid := range bufInfo.AllPids() {
 		pids[i] = uint32(bufPid)
 	}
 	return pids, nil

internal/jobobject/jobobject_test.go

@@ -261,3 +261,35 @@ func TestNoMoreProcessesMessageTerminate(t *testing.T) {
 		t.Fatal("didn't receive no more processes message within timeout")
 	}
 }
+
+func TestVerifyPidCount(t *testing.T) {
+	// This test verifies that job.Pids() returns the right info and works with > 1 process.
+	options := &Options{
+		Name:          "test",
+		Notifications: true,
+	}
+	job, err := Create(context.Background(), options)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer job.Close()
+
+	numProcs := 2
+	_, err = createProcsAndAssign(numProcs, job)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	pids, err := job.Pids()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(pids) != numProcs {
+		t.Fatalf("expected %d processes in the job, got: %d", numProcs, len(pids))
+	}
+
+	if err := job.Terminate(1); err != nil {
+		t.Fatal(err)
+	}
+}

internal/winapi/jobobject.go

@@ -93,7 +93,7 @@ type JOBOBJECT_BASIC_PROCESS_ID_LIST struct {
 
 // AllPids returns all the process Ids in the job object.
 func (p *JOBOBJECT_BASIC_PROCESS_ID_LIST) AllPids() []uintptr {
-	return (*[(1 << 27) - 1]uintptr)(unsafe.Pointer(&p.ProcessIdList[0]))[:p.NumberOfProcessIdsInList]
+	return (*[(1 << 27) - 1]uintptr)(unsafe.Pointer(&p.ProcessIdList[0]))[:p.NumberOfProcessIdsInList:p.NumberOfProcessIdsInList]
 }
 
 // https://docs.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-jobobject_basic_accounting_information