-
-
Notifications
You must be signed in to change notification settings - Fork 132
refactor: isolate platform-specific process termination logic #259
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
ff045e1
8548e44
6394674
03a377e
99a5a84
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,98 @@ | ||
| //go:build !windows | ||
|
|
||
| package util | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "context" | ||
| "fmt" | ||
| "os/exec" | ||
| "strings" | ||
| "syscall" | ||
| "time" | ||
| ) | ||
|
|
||
| // GetAPIKeyFromHelper executes a shell command to dynamically generate an API key. | ||
| // The command is executed in /bin/sh with a timeout controlled by the provided context. | ||
| // It returns the trimmed output from stdout, or an error if the command fails. | ||
| // | ||
| // On timeout, it kills the entire process group (shell and all descendants) using | ||
| // a two-phase approach: SIGTERM for graceful termination, then SIGKILL if needed. | ||
| // | ||
| // Security note: The returned API key is sensitive and should not be logged. | ||
| func GetAPIKeyFromHelper(ctx context.Context, helperCmd string) (string, error) { | ||
| if helperCmd == "" { | ||
| return "", fmt.Errorf("api_key_helper command is empty") | ||
| } | ||
|
|
||
| // Create context with timeout if not already set | ||
| if _, hasDeadline := ctx.Deadline(); !hasDeadline { | ||
| var cancel context.CancelFunc | ||
| ctx, cancel = context.WithTimeout(ctx, HelperTimeout) | ||
| defer cancel() | ||
| } | ||
|
|
||
| // Execute command in /bin/sh | ||
| cmd := exec.CommandContext(ctx, "/bin/sh", "-c", helperCmd) | ||
|
|
||
| // Create a new process group so we can kill all descendants on timeout | ||
| cmd.SysProcAttr = &syscall.SysProcAttr{ | ||
| Setpgid: true, | ||
| } | ||
|
|
||
| var stdout, stderr bytes.Buffer | ||
| cmd.Stdout = &stdout | ||
| cmd.Stderr = &stderr | ||
|
|
||
| // Start the command | ||
| if err := cmd.Start(); err != nil { | ||
| return "", fmt.Errorf("api_key_helper start failed: %w", err) | ||
| } | ||
|
|
||
| // Wait for command completion in a goroutine | ||
| done := make(chan error, 1) | ||
| go func() { | ||
| // Always Wait to avoid zombie processes | ||
| done <- cmd.Wait() | ||
| }() | ||
|
|
||
| select { | ||
| case err := <-done: | ||
| // Command completed normally | ||
| if err != nil { | ||
| // Don't include stderr in error message as it might contain sensitive info | ||
| return "", fmt.Errorf("api_key_helper command failed: %w", err) | ||
| } | ||
| apiKey := strings.TrimSpace(stdout.String()) | ||
| if apiKey == "" { | ||
| return "", fmt.Errorf("api_key_helper command returned empty output") | ||
| } | ||
| return apiKey, nil | ||
|
|
||
| case <-ctx.Done(): | ||
| // Timeout or cancellation: terminate the process group gracefully, then forcefully | ||
appleboy marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| if cmd.Process == nil { | ||
| // Process handle not initialized; wait for cleanup and report timeout | ||
| <-done | ||
| return "", fmt.Errorf("api_key_helper command timeout after %v", HelperTimeout) | ||
| } | ||
| pgid := cmd.Process.Pid | ||
|
|
||
| // First attempt: send SIGTERM to the entire process group for graceful shutdown | ||
| _ = syscall.Kill(-pgid, syscall.SIGTERM) | ||
|
|
||
| // Wait for graceful termination with a grace period | ||
| select { | ||
| case <-done: | ||
| // Process exited after timeout was reached; treat as timeout regardless of exit status. | ||
| // We intentionally ignore stdout/stderr here to avoid returning a key after a timeout. | ||
| return "", fmt.Errorf("api_key_helper command timeout after %v", HelperTimeout) | ||
|
|
||
| case <-time.After(2 * time.Second): | ||
|
||
| // Grace period expired: send SIGKILL to force termination | ||
| _ = syscall.Kill(-pgid, syscall.SIGKILL) | ||
| <-done // Wait for cleanup | ||
| return "", fmt.Errorf("api_key_helper command timeout after %v", HelperTimeout) | ||
| } | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,148 @@ | ||
| //go:build windows | ||
|
|
||
| package util | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "context" | ||
| "fmt" | ||
| "os/exec" | ||
| "strings" | ||
| "syscall" | ||
| "unsafe" | ||
|
|
||
| "golang.org/x/sys/windows" | ||
| ) | ||
|
|
||
| // createKillOnCloseJob creates a Windows Job Object with KILL_ON_JOB_CLOSE flag. | ||
| // When the job handle is closed, all processes in the job will be terminated. | ||
| func createKillOnCloseJob() (windows.Handle, error) { | ||
| job, err := windows.CreateJobObject(nil, nil) | ||
| if err != nil { | ||
| return 0, err | ||
| } | ||
|
|
||
| var info windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION | ||
| // Enable KILL_ON_JOB_CLOSE flag | ||
| info.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE | ||
|
|
||
| _, err = windows.SetInformationJobObject( | ||
| job, | ||
| windows.JobObjectExtendedLimitInformation, | ||
| uintptr(unsafe.Pointer(&info)), | ||
| uint32(unsafe.Sizeof(info)), | ||
| ) | ||
| if err != nil { | ||
| _ = windows.CloseHandle(job) | ||
| return 0, err | ||
| } | ||
| return job, nil | ||
| } | ||
|
|
||
| // assignProcessToJob assigns a process to a Job Object. | ||
| // Returns the process handle which should be closed by the caller. | ||
| func assignProcessToJob(job windows.Handle, pid int) (windows.Handle, error) { | ||
| // Validate PID range to prevent overflow | ||
| if pid < 0 || pid > 0x7FFFFFFF { | ||
| return 0, fmt.Errorf("invalid process ID: %d", pid) | ||
| } | ||
|
|
||
| // Get child process handle (requires PROCESS_ALL_ACCESS) | ||
| // #nosec G115 -- PID validated above to prevent overflow | ||
| hProc, err := windows.OpenProcess(windows.PROCESS_ALL_ACCESS, false, uint32(pid)) | ||
| if err != nil { | ||
| return 0, err | ||
| } | ||
| // Assign to Job | ||
| if err = windows.AssignProcessToJobObject(job, hProc); err != nil { | ||
| _ = windows.CloseHandle(hProc) | ||
| return 0, err | ||
| } | ||
| return hProc, nil | ||
| } | ||
|
|
||
| // GetAPIKeyFromHelper executes a shell command to dynamically generate an API key. | ||
| // The command is executed in cmd.exe with a timeout controlled by the provided context. | ||
| // It returns the trimmed output from stdout, or an error if the command fails. | ||
| // | ||
| // On timeout, it terminates the entire Job Object (cmd.exe and all descendants). | ||
| // | ||
| // Security note: The returned API key is sensitive and should not be logged. | ||
| func GetAPIKeyFromHelper(ctx context.Context, helperCmd string) (string, error) { | ||
| if helperCmd == "" { | ||
| return "", fmt.Errorf("api_key_helper command is empty") | ||
| } | ||
|
|
||
| // Create context with timeout if not already set | ||
| if _, hasDeadline := ctx.Deadline(); !hasDeadline { | ||
| var cancel context.CancelFunc | ||
| ctx, cancel = context.WithTimeout(ctx, HelperTimeout) | ||
| defer cancel() | ||
| } | ||
|
|
||
| // Execute command in cmd.exe | ||
| cmd := exec.CommandContext(ctx, "cmd.exe", "/c", helperCmd) | ||
|
|
||
| // Use CREATE_NEW_PROCESS_GROUP and CREATE_BREAKAWAY_FROM_JOB flags | ||
| // This allows the child process to be assigned to a new Job, | ||
| // even if the parent process is already in a Job | ||
| cmd.SysProcAttr = &syscall.SysProcAttr{ | ||
| CreationFlags: windows.CREATE_NEW_PROCESS_GROUP | windows.CREATE_BREAKAWAY_FROM_JOB, | ||
| } | ||
|
|
||
| var stdout, stderr bytes.Buffer | ||
| cmd.Stdout = &stdout | ||
| cmd.Stderr = &stderr | ||
|
|
||
| // Create Job Object first | ||
| job, err := createKillOnCloseJob() | ||
| if err != nil { | ||
| return "", fmt.Errorf("create job failed: %w", err) | ||
| } | ||
| // With KILL_ON_JOB_CLOSE, closing the job will kill all processes | ||
| defer func() { | ||
| _ = windows.CloseHandle(job) | ||
| }() | ||
|
|
||
| // Start the child process | ||
| if err = cmd.Start(); err != nil { | ||
| return "", fmt.Errorf("api_key_helper start failed: %w", err) | ||
| } | ||
|
|
||
| // Assign child process to Job | ||
| hProc, err := assignProcessToJob(job, cmd.Process.Pid) | ||
| if err != nil { | ||
| // If unable to breakaway due to policy, fall back to just killing the process | ||
| // (but this won't guarantee killing grandchild processes) | ||
| _ = cmd.Process.Kill() | ||
| _ = cmd.Wait() | ||
|
Comment on lines
+117
to
+118
|
||
| return "", fmt.Errorf("assign process to job failed: %w", err) | ||
| } | ||
| defer func() { | ||
| _ = windows.CloseHandle(hProc) | ||
| }() | ||
|
|
||
| done := make(chan error, 1) | ||
| go func() { | ||
| done <- cmd.Wait() | ||
| }() | ||
|
|
||
| select { | ||
| case err := <-done: | ||
| if err != nil { | ||
| // Don't include stderr in error message as it might contain sensitive info | ||
| return "", fmt.Errorf("api_key_helper command failed: %w", err) | ||
| } | ||
| apiKey := strings.TrimSpace(stdout.String()) | ||
| if apiKey == "" { | ||
| return "", fmt.Errorf("api_key_helper command returned empty output") | ||
| } | ||
| return apiKey, nil | ||
|
|
||
| case <-ctx.Done(): | ||
| // Timeout: terminate the entire Job (all descendants) | ||
| _ = windows.TerminateJobObject(job, 1) | ||
| <-done // Wait for cleanup | ||
| return "", fmt.Errorf("api_key_helper command timeout after %v", HelperTimeout) | ||
| } | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This documentation comment block for GetAPIKeyFromHelper is not attached to any function declaration in this file. The actual platform-specific implementations in api_key_helper_unix.go and api_key_helper_windows.go have their own complete documentation. Consider either removing this orphaned comment block or converting it to a package-level comment that explains the platform-specific implementations exist elsewhere.