refactor: isolate platform-specific process termination logic

util/api_key_helper.go

@@ -1,17 +1,13 @@
 package util
 
 import (
-	"bytes"
 	"context"
 	"crypto/sha256"
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"os"
-	"os/exec"
 	"path/filepath"
-	"strings"
-	"syscall"
 	"time"
 )
 
@@ -127,91 +123,16 @@ func needsRefresh(cache *apiKeyCache, refreshInterval time.Duration) bool {
 }
 
 // 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.
+// Platform-specific implementations are in api_key_helper_unix.go and api_key_helper_windows.go.
+//
+// The command is executed 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.
+// On timeout:
+//   - Unix/Linux/macOS: kills the entire process group (shell and all descendants)
+//   - Windows: terminates the 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 /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
-		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 err := <-done:
-			if err != nil {
-				return "", fmt.Errorf("api_key_helper terminated after timeout: %w", err)
-			}
-			apiKey := strings.TrimSpace(stdout.String())
-			if apiKey == "" {
-				return "", fmt.Errorf(
-					"api_key_helper command returned empty output after timeout termination",
-				)
-			}
-			return apiKey, nil
-
-		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 timed out after %v", HelperTimeout)
-		}
-	}
-}
 
 // GetAPIKeyFromHelperWithCache executes a shell command to dynamically generate an API key,
 // with file-based caching support. The API key is cached for the specified refresh interval.

util/api_key_helper_unix.go

@@ -0,0 +1,100 @@
+//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
+		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 err := <-done:
+			if err != nil {
+				return "", fmt.Errorf("api_key_helper terminated after timeout: %w", err)
+			}
+			apiKey := strings.TrimSpace(stdout.String())
+			if apiKey == "" {
+				return "", fmt.Errorf(
+					"api_key_helper command returned empty output after timeout termination",
+				)
+			}
+			return apiKey, nil
+
+		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 timed out after %v", HelperTimeout)
+		}
+	}
+}

util/api_key_helper_windows.go

@@ -0,0 +1,137 @@
+//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) {
+	// Get child process handle (requires PROCESS_ALL_ACCESS)
+	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)
+	}
+	defer windows.CloseHandle(job) // With KILL_ON_JOB_CLOSE, closing will kill all processes
+
+	// 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()
+		return "", fmt.Errorf("assign process to job failed: %w", err)
+	}
+	defer 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 timed out after %v", HelperTimeout)
+	}
+}