Skip to content

Commit a0e73e1

Browse files
appleboyappleboy
authored
refactor: isolate platform-specific process termination logic (#259)
* refactor: isolate platform-specific process termination logic - Refactor platform-specific API key helper logic into separate files for Unix and Windows - Remove direct process management and command execution code from the shared helper file - Implement Unix-specific process group termination using signals for timeouts - Implement Windows-specific Job Object termination for timeouts - Improve code clarity and maintainability by isolating OS-specific behaviors Signed-off-by: Bo-Yi Wu <[email protected]> * fix: improve process handle safety and validation - Assign the result of CloseHandle to underscore to avoid unchecked errors - Add PID validation to prevent integer overflow and invalid process IDs - Refactor resource cleanup to use deferred anonymous functions for handle closing Signed-off-by: Bo-Yi Wu <[email protected]> * fix: enhance API key retrieval with robust timeout handling - Improve timeout handling in GetAPIKeyFromHelper by ensuring no API key is returned after a timeout - Add explicit check for uninitialized process handle and report timeout error Signed-off-by: Bo-Yi Wu <[email protected]> * docs: standardize error message terminology for API key helper - Fix typo in error messages: change "timed out" to "timeout" for api_key_helper command errors Signed-off-by: Bo-Yi Wu <[email protected]> * test: expand Windows API key helper and process management tests - Add comprehensive Windows-specific tests for API key helper functionality - Test job object creation and validation of process termination on job close - Verify error handling for invalid and non-existent process IDs - Ensure API key extraction works for various Windows command scenarios - Test timeout behavior and process tree termination using job objects - Check error handling for failed commands and empty output cases - Validate that sensitive stderr output is not leaked in error messages - Confirm correct handling of context cancellation and concurrent invocations Signed-off-by: Bo-Yi Wu <[email protected]> --------- Signed-off-by: Bo-Yi Wu <[email protected]>
1 parent e4a9254 commit a0e73e1

File tree

4 files changed

+622
-85
lines changed

4 files changed

+622
-85
lines changed

‎util/api_key_helper.go‎

Lines changed: 6 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,13 @@
11
package util
22

33
import (
4-
"bytes"
54
"context"
65
"crypto/sha256"
76
"encoding/hex"
87
"encoding/json"
98
"fmt"
109
"os"
11-
"os/exec"
1210
"path/filepath"
13-
"strings"
14-
"syscall"
1511
"time"
1612
)
1713

@@ -127,91 +123,16 @@ func needsRefresh(cache *apiKeyCache, refreshInterval time.Duration) bool {
127123
}
128124

129125
// GetAPIKeyFromHelper executes a shell command to dynamically generate an API key.
130-
// The command is executed in /bin/sh with a timeout controlled by the provided context.
126+
// Platform-specific implementations are in api_key_helper_unix.go and api_key_helper_windows.go.
127+
//
128+
// The command is executed with a timeout controlled by the provided context.
131129
// It returns the trimmed output from stdout, or an error if the command fails.
132130
//
133-
// On timeout, it kills the entire process group (shell and all descendants) using
134-
// a two-phase approach: SIGTERM for graceful termination, then SIGKILL if needed.
131+
// On timeout:
132+
// - Unix/Linux/macOS: kills the entire process group (shell and all descendants)
133+
// - Windows: terminates the Job Object (cmd.exe and all descendants)
135134
//
136135
// Security note: The returned API key is sensitive and should not be logged.
137-
func GetAPIKeyFromHelper(ctx context.Context, helperCmd string) (string, error) {
138-
if helperCmd == "" {
139-
return "", fmt.Errorf("api_key_helper command is empty")
140-
}
141-
142-
// Create context with timeout if not already set
143-
if _, hasDeadline := ctx.Deadline(); !hasDeadline {
144-
var cancel context.CancelFunc
145-
ctx, cancel = context.WithTimeout(ctx, HelperTimeout)
146-
defer cancel()
147-
}
148-
149-
// Execute command in /bin/sh
150-
cmd := exec.CommandContext(ctx, "/bin/sh", "-c", helperCmd)
151-
152-
// Create a new process group so we can kill all descendants on timeout
153-
cmd.SysProcAttr = &syscall.SysProcAttr{
154-
Setpgid: true,
155-
}
156-
157-
var stdout, stderr bytes.Buffer
158-
cmd.Stdout = &stdout
159-
cmd.Stderr = &stderr
160-
161-
// Start the command
162-
if err := cmd.Start(); err != nil {
163-
return "", fmt.Errorf("api_key_helper start failed: %w", err)
164-
}
165-
166-
// Wait for command completion in a goroutine
167-
done := make(chan error, 1)
168-
go func() {
169-
// Always Wait to avoid zombie processes
170-
done <- cmd.Wait()
171-
}()
172-
173-
select {
174-
case err := <-done:
175-
// Command completed normally
176-
if err != nil {
177-
// Don't include stderr in error message as it might contain sensitive info
178-
return "", fmt.Errorf("api_key_helper command failed: %w", err)
179-
}
180-
apiKey := strings.TrimSpace(stdout.String())
181-
if apiKey == "" {
182-
return "", fmt.Errorf("api_key_helper command returned empty output")
183-
}
184-
return apiKey, nil
185-
186-
case <-ctx.Done():
187-
// Timeout or cancellation: terminate the process group gracefully, then forcefully
188-
pgid := cmd.Process.Pid
189-
190-
// First attempt: send SIGTERM to the entire process group for graceful shutdown
191-
_ = syscall.Kill(-pgid, syscall.SIGTERM)
192-
193-
// Wait for graceful termination with a grace period
194-
select {
195-
case err := <-done:
196-
if err != nil {
197-
return "", fmt.Errorf("api_key_helper terminated after timeout: %w", err)
198-
}
199-
apiKey := strings.TrimSpace(stdout.String())
200-
if apiKey == "" {
201-
return "", fmt.Errorf(
202-
"api_key_helper command returned empty output after timeout termination",
203-
)
204-
}
205-
return apiKey, nil
206-
207-
case <-time.After(2 * time.Second):
208-
// Grace period expired: send SIGKILL to force termination
209-
_ = syscall.Kill(-pgid, syscall.SIGKILL)
210-
<-done // Wait for cleanup
211-
return "", fmt.Errorf("api_key_helper command timed out after %v", HelperTimeout)
212-
}
213-
}
214-
}
215136

216137
// GetAPIKeyFromHelperWithCache executes a shell command to dynamically generate an API key,
217138
// with file-based caching support. The API key is cached for the specified refresh interval.

‎util/api_key_helper_unix.go‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
//go:build !windows
2+
3+
package util
4+
5+
import (
6+
"bytes"
7+
"context"
8+
"fmt"
9+
"os/exec"
10+
"strings"
11+
"syscall"
12+
"time"
13+
)
14+
15+
// GetAPIKeyFromHelper executes a shell command to dynamically generate an API key.
16+
// The command is executed in /bin/sh with a timeout controlled by the provided context.
17+
// It returns the trimmed output from stdout, or an error if the command fails.
18+
//
19+
// On timeout, it kills the entire process group (shell and all descendants) using
20+
// a two-phase approach: SIGTERM for graceful termination, then SIGKILL if needed.
21+
//
22+
// Security note: The returned API key is sensitive and should not be logged.
23+
func GetAPIKeyFromHelper(ctx context.Context, helperCmd string) (string, error) {
24+
if helperCmd == "" {
25+
return "", fmt.Errorf("api_key_helper command is empty")
26+
}
27+
28+
// Create context with timeout if not already set
29+
if _, hasDeadline := ctx.Deadline(); !hasDeadline {
30+
var cancel context.CancelFunc
31+
ctx, cancel = context.WithTimeout(ctx, HelperTimeout)
32+
defer cancel()
33+
}
34+
35+
// Execute command in /bin/sh
36+
cmd := exec.CommandContext(ctx, "/bin/sh", "-c", helperCmd)
37+
38+
// Create a new process group so we can kill all descendants on timeout
39+
cmd.SysProcAttr = &syscall.SysProcAttr{
40+
Setpgid: true,
41+
}
42+
43+
var stdout, stderr bytes.Buffer
44+
cmd.Stdout = &stdout
45+
cmd.Stderr = &stderr
46+
47+
// Start the command
48+
if err := cmd.Start(); err != nil {
49+
return "", fmt.Errorf("api_key_helper start failed: %w", err)
50+
}
51+
52+
// Wait for command completion in a goroutine
53+
done := make(chan error, 1)
54+
go func() {
55+
// Always Wait to avoid zombie processes
56+
done <- cmd.Wait()
57+
}()
58+
59+
select {
60+
case err := <-done:
61+
// Command completed normally
62+
if err != nil {
63+
// Don't include stderr in error message as it might contain sensitive info
64+
return "", fmt.Errorf("api_key_helper command failed: %w", err)
65+
}
66+
apiKey := strings.TrimSpace(stdout.String())
67+
if apiKey == "" {
68+
return "", fmt.Errorf("api_key_helper command returned empty output")
69+
}
70+
return apiKey, nil
71+
72+
case <-ctx.Done():
73+
// Timeout or cancellation: terminate the process group gracefully, then forcefully
74+
if cmd.Process == nil {
75+
// Process handle not initialized; wait for cleanup and report timeout
76+
<-done
77+
return "", fmt.Errorf("api_key_helper command timeout after %v", HelperTimeout)
78+
}
79+
pgid := cmd.Process.Pid
80+
81+
// First attempt: send SIGTERM to the entire process group for graceful shutdown
82+
_ = syscall.Kill(-pgid, syscall.SIGTERM)
83+
84+
// Wait for graceful termination with a grace period
85+
select {
86+
case <-done:
87+
// Process exited after timeout was reached; treat as timeout regardless of exit status.
88+
// We intentionally ignore stdout/stderr here to avoid returning a key after a timeout.
89+
return "", fmt.Errorf("api_key_helper command timeout after %v", HelperTimeout)
90+
91+
case <-time.After(2 * time.Second):
92+
// Grace period expired: send SIGKILL to force termination
93+
_ = syscall.Kill(-pgid, syscall.SIGKILL)
94+
<-done // Wait for cleanup
95+
return "", fmt.Errorf("api_key_helper command timeout after %v", HelperTimeout)
96+
}
97+
}
98+
}

‎util/api_key_helper_windows.go‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
//go:build windows
2+
3+
package util
4+
5+
import (
6+
"bytes"
7+
"context"
8+
"fmt"
9+
"os/exec"
10+
"strings"
11+
"syscall"
12+
"unsafe"
13+
14+
"golang.org/x/sys/windows"
15+
)
16+
17+
// createKillOnCloseJob creates a Windows Job Object with KILL_ON_JOB_CLOSE flag.
18+
// When the job handle is closed, all processes in the job will be terminated.
19+
func createKillOnCloseJob() (windows.Handle, error) {
20+
job, err := windows.CreateJobObject(nil, nil)
21+
if err != nil {
22+
return 0, err
23+
}
24+
25+
var info windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION
26+
// Enable KILL_ON_JOB_CLOSE flag
27+
info.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
28+
29+
_, err = windows.SetInformationJobObject(
30+
job,
31+
windows.JobObjectExtendedLimitInformation,
32+
uintptr(unsafe.Pointer(&info)),
33+
uint32(unsafe.Sizeof(info)),
34+
)
35+
if err != nil {
36+
_ = windows.CloseHandle(job)
37+
return 0, err
38+
}
39+
return job, nil
40+
}
41+
42+
// assignProcessToJob assigns a process to a Job Object.
43+
// Returns the process handle which should be closed by the caller.
44+
func assignProcessToJob(job windows.Handle, pid int) (windows.Handle, error) {
45+
// Validate PID range to prevent overflow
46+
if pid < 0 || pid > 0x7FFFFFFF {
47+
return 0, fmt.Errorf("invalid process ID: %d", pid)
48+
}
49+
50+
// Get child process handle (requires PROCESS_ALL_ACCESS)
51+
// #nosec G115 -- PID validated above to prevent overflow
52+
hProc, err := windows.OpenProcess(windows.PROCESS_ALL_ACCESS, false, uint32(pid))
53+
if err != nil {
54+
return 0, err
55+
}
56+
// Assign to Job
57+
if err = windows.AssignProcessToJobObject(job, hProc); err != nil {
58+
_ = windows.CloseHandle(hProc)
59+
return 0, err
60+
}
61+
return hProc, nil
62+
}
63+
64+
// GetAPIKeyFromHelper executes a shell command to dynamically generate an API key.
65+
// The command is executed in cmd.exe with a timeout controlled by the provided context.
66+
// It returns the trimmed output from stdout, or an error if the command fails.
67+
//
68+
// On timeout, it terminates the entire Job Object (cmd.exe and all descendants).
69+
//
70+
// Security note: The returned API key is sensitive and should not be logged.
71+
func GetAPIKeyFromHelper(ctx context.Context, helperCmd string) (string, error) {
72+
if helperCmd == "" {
73+
return "", fmt.Errorf("api_key_helper command is empty")
74+
}
75+
76+
// Create context with timeout if not already set
77+
if _, hasDeadline := ctx.Deadline(); !hasDeadline {
78+
var cancel context.CancelFunc
79+
ctx, cancel = context.WithTimeout(ctx, HelperTimeout)
80+
defer cancel()
81+
}
82+
83+
// Execute command in cmd.exe
84+
cmd := exec.CommandContext(ctx, "cmd.exe", "/c", helperCmd)
85+
86+
// Use CREATE_NEW_PROCESS_GROUP and CREATE_BREAKAWAY_FROM_JOB flags
87+
// This allows the child process to be assigned to a new Job,
88+
// even if the parent process is already in a Job
89+
cmd.SysProcAttr = &syscall.SysProcAttr{
90+
CreationFlags: windows.CREATE_NEW_PROCESS_GROUP | windows.CREATE_BREAKAWAY_FROM_JOB,
91+
}
92+
93+
var stdout, stderr bytes.Buffer
94+
cmd.Stdout = &stdout
95+
cmd.Stderr = &stderr
96+
97+
// Create Job Object first
98+
job, err := createKillOnCloseJob()
99+
if err != nil {
100+
return "", fmt.Errorf("create job failed: %w", err)
101+
}
102+
// With KILL_ON_JOB_CLOSE, closing the job will kill all processes
103+
defer func() {
104+
_ = windows.CloseHandle(job)
105+
}()
106+
107+
// Start the child process
108+
if err = cmd.Start(); err != nil {
109+
return "", fmt.Errorf("api_key_helper start failed: %w", err)
110+
}
111+
112+
// Assign child process to Job
113+
hProc, err := assignProcessToJob(job, cmd.Process.Pid)
114+
if err != nil {
115+
// If unable to breakaway due to policy, fall back to just killing the process
116+
// (but this won't guarantee killing grandchild processes)
117+
_ = cmd.Process.Kill()
118+
_ = cmd.Wait()
119+
return "", fmt.Errorf("assign process to job failed: %w", err)
120+
}
121+
defer func() {
122+
_ = windows.CloseHandle(hProc)
123+
}()
124+
125+
done := make(chan error, 1)
126+
go func() {
127+
done <- cmd.Wait()
128+
}()
129+
130+
select {
131+
case err := <-done:
132+
if err != nil {
133+
// Don't include stderr in error message as it might contain sensitive info
134+
return "", fmt.Errorf("api_key_helper command failed: %w", err)
135+
}
136+
apiKey := strings.TrimSpace(stdout.String())
137+
if apiKey == "" {
138+
return "", fmt.Errorf("api_key_helper command returned empty output")
139+
}
140+
return apiKey, nil
141+
142+
case <-ctx.Done():
143+
// Timeout: terminate the entire Job (all descendants)
144+
_ = windows.TerminateJobObject(job, 1)
145+
<-done // Wait for cleanup
146+
return "", fmt.Errorf("api_key_helper command timeout after %v", HelperTimeout)
147+
}
148+
}

0 commit comments

Comments
 (0)