fix: make cancellation interrupt auth and plugin waits

Author: danielgtaylor

internal/auth/context.go

@@ -0,0 +1,13 @@
+package auth
+
+import (
+	"context"
+	"net/http"
+)
+
+func requestWithContext(req *http.Request, ctx context.Context) *http.Request {
+	if req == nil || ctx == nil || req.Context() == ctx {
+		return req
+	}
+	return req.WithContext(ctx)
+}

internal/auth/external_tool.go

@@ -73,6 +73,9 @@ func (a *ExternalTool) OnRequest(req *http.Request, params map[string]string) er
 }
 
 func (a *ExternalTool) run(ctx context.Context, req *http.Request, params map[string]string) error {
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	commandLine := params["commandline"]
 	if commandLine == "" {
 		return fmt.Errorf("external-tool auth: missing required param \"commandline\"")
@@ -180,7 +183,7 @@ func (a *ExternalTool) Authenticate(ctx context.Context, req *http.Request, ac A
 	if stderr == nil {
 		stderr = ac.Stderr
 	}
-	if req.Context() != nil {
+	if ctx == nil && req.Context() != nil {
 		ctx = req.Context()
 	}
 	return (&ExternalTool{Stderr: stderr, Timeout: a.Timeout}).run(ctx, req, ac.Params)

internal/auth/oauth_authcode.go

@@ -5,6 +5,7 @@ import (
 	"crypto/rand"
 	"crypto/sha256"
 	"encoding/base64"
+	"errors"
 	"fmt"
 	"html"
 	"io"
@@ -152,6 +153,7 @@ func (h *AuthorizationCode) Authenticate(ctx context.Context, req *http.Request,
 		h2.Prompt = ac.Prompter.Prompt
 		h2.CanPrompt = true
 	}
+	req = requestWithContext(req, ctx)
 	return h2.authenticateRequest(req, authParams(ac), ac.Force)
 }
 
@@ -373,6 +375,9 @@ func (h *AuthorizationCode) doBrowserFlow(ctx context.Context, params map[string
 		case err = <-errCh:
 			return CachedToken{}, fmt.Errorf("callback error: %w", err)
 		case <-ctx2.Done():
+			if errors.Is(ctx2.Err(), context.Canceled) {
+				return CachedToken{}, ctx2.Err()
+			}
 			return CachedToken{}, fmt.Errorf("timed out waiting for authorization callback")
 		}
 	}

internal/auth/oauth_authcode_test.go

@@ -296,6 +296,42 @@ func TestAuthCode_BrowserFlow_ImmediateCallbackDuringOpenBrowser(t *testing.T) {
 	}
 }
 
+func TestAuthCodeAuthenticateUsesExplicitContext(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	h := &AuthorizationCode{
+		HTTPClient: testHTTPClient(func(r *http.Request) (*http.Response, error) {
+			if err := r.Context().Err(); err != nil {
+				return nil, err
+			}
+			return testResponse(200, "application/json", `{"access_token":"unexpected-token","token_type":"bearer","expires_in":3600}`), nil
+		}),
+		OpenBrowser: func(raw string) error {
+			callbackURL, state := mustCallbackURL(t, raw)
+			resp, err := http.Get(fmt.Sprintf("%s/?state=%s&code=late-code", callbackURL, url.QueryEscape(state)))
+			if err != nil {
+				return err
+			}
+			defer resp.Body.Close()
+			return nil
+		},
+	}
+
+	req, _ := http.NewRequest("GET", "https://api.example.com", nil)
+	err := h.Authenticate(ctx, req, AuthContext{Params: map[string]string{
+		"client_id":     "id1",
+		"authorize_url": "https://auth.example.com/authorize",
+		"token_url":     "https://auth.example.com/token",
+		"redirect_port": availablePort(t),
+	}})
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected context cancellation, got %v", err)
+	}
+	if got := req.Header.Get("Authorization"); got != "" {
+		t.Fatalf("Authorization = %q, want empty", got)
+	}
+}
+
 func TestAuthCode_BrowserFlow_TwoStrayPreflightsDoNotDeadlock(t *testing.T) {
 	h := &AuthorizationCode{
 		HTTPClient: testHTTPClient(func(r *http.Request) (*http.Response, error) {

internal/auth/oauth_client_creds.go

@@ -52,6 +52,7 @@ func (h *ClientCredentials) Authenticate(ctx context.Context, req *http.Request,
 	if ac.HTTPClient != nil {
 		h2.HTTPClient = ac.HTTPClient
 	}
+	req = requestWithContext(req, ctx)
 	return h2.authenticateRequest(req, authParams(ac), ac.Force)
 }
 

internal/auth/oauth_device_code.go

@@ -221,6 +221,9 @@ func (h *DeviceCode) runFlow(ctx context.Context, params map[string]string, devi
 		select {
 		case <-pollCtx.Done():
 			timer.Stop()
+			if errors.Is(pollCtx.Err(), context.Canceled) {
+				return CachedToken{}, pollCtx.Err()
+			}
 			return CachedToken{}, fmt.Errorf("timed out waiting for device authorization")
 		case <-timer.C:
 		}
@@ -303,6 +306,7 @@ func (h *DeviceCode) Authenticate(ctx context.Context, req *http.Request, ac Aut
 	if ac.Stderr != nil {
 		h2.Stderr = ac.Stderr
 	}
+	req = requestWithContext(req, ctx)
 	return h2.authenticateRequest(req, authParams(ac), ac.Force)
 }
 

internal/cli/api_auth.go

@@ -92,7 +92,7 @@ func (c *CLI) runAPIAuthList(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
-	set, hasOps := c.cachedOperationSetForAPI(apiName, apiCfg, profileName)
+	set, hasOps := c.cachedOperationSetForAPI(requestContext(cmd), apiName, apiCfg, profileName)
 	_, profileReady, profileErr := c.profileAuthReadiness(apiName, profileName, prof)
 	coverage := operationAuthCoverage{}
 	if hasOps {
@@ -204,7 +204,7 @@ func (c *CLI) runAPIAuthAdd(cmd *cobra.Command, args []string) error {
 	if prof.Credentials[credentialID] == nil {
 		prof.Credentials[credentialID] = &config.CredentialConfig{}
 	}
-	defaultNeeds := c.cachedCredentialDefaultNeeds(apiName, apiCfg, profileName, credentialID)
+	defaultNeeds := c.cachedCredentialDefaultNeeds(requestContext(cmd), apiName, apiCfg, profileName, credentialID)
 	if prof.Credentials[credentialID].Auth == nil {
 		if authCfg, ok, err := c.cachedAuthConfigForCredential(apiName, apiCfg, credentialID); err != nil {
 			return err
@@ -298,7 +298,7 @@ func (c *CLI) runAPIAuthInspect(cmd *cobra.Command, args []string) error {
 }
 
 func (c *CLI) runAPIAuthInspectOperation(cmd *cobra.Command, apiName, profileName string, apiCfg *config.APIConfig, prof *config.ProfileConfig, operationName, rawHeader string, redact bool) error {
-	op, ok, err := c.cachedOperationForAPI(apiName, apiCfg, profileName, operationName)
+	op, ok, err := c.cachedOperationForAPI(requestContext(cmd), apiName, apiCfg, profileName, operationName)
 	if err != nil {
 		return err
 	}
@@ -471,7 +471,7 @@ func (c *CLI) operationAuthInspectionRequest(cmd *cobra.Command, apiName, profil
 	if err != nil {
 		return nil, err
 	}
-	req, _ := http.NewRequest("GET", "http://example.com", nil)
+	req, _ := http.NewRequestWithContext(requestContext(cmd), "GET", "http://example.com", nil)
 	for _, item := range selected {
 		step, err := c.operationAuthStep(apiName, profileName, item, authOpts)
 		if err != nil {
@@ -510,7 +510,7 @@ func (c *CLI) authInspectionRequest(cmd *cobra.Command, apiName, profileName str
 	if err != nil {
 		return nil, err
 	}
-	req, _ := http.NewRequest("GET", "http://example.com", nil)
+	req, _ := http.NewRequestWithContext(requestContext(cmd), "GET", "http://example.com", nil)
 	if err := handler.Authenticate(requestContext(cmd), req, c.authContext(requestContext(cmd), apiName, profileName, params, resolved.CacheKey, false)); err != nil {
 		return nil, fmt.Errorf("building auth inspection: %w", err)
 	}
@@ -636,11 +636,11 @@ func (c *CLI) apiProfileForAuth(apiName, profileName string, create bool) (*conf
 	return apiCfg, apiCfg.Profiles[profileName], nil
 }
 
-func (c *CLI) cachedOperationSetForAPI(apiName string, apiCfg *config.APIConfig, profileName string) (spec.OperationSet, bool) {
+func (c *CLI) cachedOperationSetForAPI(ctx context.Context, apiName string, apiCfg *config.APIConfig, profileName string) (spec.OperationSet, bool) {
 	if set, _, ok := c.cachedOperationSetStatusForAPI(apiName, apiCfg, profileName); ok {
 		return set, true
 	}
-	set, ok, _ := c.operationSetForAPI(context.Background(), apiName, apiCfg, profileName, false)
+	set, ok, _ := c.operationSetForAPI(ctx, apiName, apiCfg, profileName, false)
 	return set, ok
 }
 
@@ -708,8 +708,8 @@ func (c *CLI) operationSetForAPI(ctx context.Context, apiName string, apiCfg *co
 	return set, true, nil
 }
 
-func (c *CLI) cachedOperationForAPI(apiName string, apiCfg *config.APIConfig, profileName, value string) (spec.Operation, bool, error) {
-	set, ok := c.cachedOperationSetForAPI(apiName, apiCfg, profileName)
+func (c *CLI) cachedOperationForAPI(ctx context.Context, apiName string, apiCfg *config.APIConfig, profileName, value string) (spec.Operation, bool, error) {
+	set, ok := c.cachedOperationSetForAPI(ctx, apiName, apiCfg, profileName)
 	if !ok {
 		return spec.Operation{}, false, nil
 	}
@@ -858,8 +858,8 @@ func authRequirementKindSupported(kind string) bool {
 	}
 }
 
-func (c *CLI) cachedCredentialDefaultNeeds(apiName string, apiCfg *config.APIConfig, profileName, credentialID string) []string {
-	set, ok := c.cachedOperationSetForAPI(apiName, apiCfg, profileName)
+func (c *CLI) cachedCredentialDefaultNeeds(ctx context.Context, apiName string, apiCfg *config.APIConfig, profileName, credentialID string) []string {
+	set, ok := c.cachedOperationSetForAPI(ctx, apiName, apiCfg, profileName)
 	if !ok {
 		return nil
 	}

internal/cli/auth.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"crypto/sha256"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"log"
 	"net/http"
@@ -312,8 +313,11 @@ func (c *CLI) runSecretCommand(commandLine string) (string, error) {
 	var stderr bytes.Buffer
 	cmd.Stderr = &limitedWriter{w: &stderr, limit: 4096}
 	out, err := cmd.Output()
-	if ctx.Err() != nil {
-		return "", fmt.Errorf("secret command timed out or was canceled")
+	if ctxErr := ctx.Err(); ctxErr != nil {
+		if errors.Is(ctxErr, context.Canceled) {
+			return "", ctxErr
+		}
+		return "", fmt.Errorf("secret command timed out: %w", ctxErr)
 	}
 	if err != nil {
 		excerpt := strings.TrimSpace(stderr.String())

internal/cli/auth_internal_test.go

@@ -67,7 +67,7 @@ func TestCachedOperationSetForAPIUsesEffectiveProfileBase(t *testing.T) {
 		t.Fatalf("Discover: %v", err)
 	}
 
-	got, ok := c.cachedOperationSetForAPI("svc", apiCfg, "default")
+	got, ok := c.cachedOperationSetForAPI(context.Background(), "svc", apiCfg, "default")
 	if !ok {
 		t.Fatal("expected cached profile operation set")
 	}

internal/cli/cli.go

@@ -14,6 +14,7 @@ import (
 	"path/filepath"
 	"sort"
 	"strings"
+	"sync"
 	"syscall"
 	"time"
 
@@ -830,17 +831,24 @@ func signalAwareContext() (context.Context, context.CancelFunc) {
 	ctx, cancelCause := context.WithCancelCause(context.Background())
 	sigCh := make(chan os.Signal, 1)
 	done := make(chan struct{})
+	var stopOnce sync.Once
+	stopSignals := func() {
+		stopOnce.Do(func() {
+			signal.Stop(sigCh)
+			close(done)
+		})
+	}
 	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
 	go func() {
 		select {
 		case sig := <-sigCh:
 			cancelCause(signalCancelError{signal: sig})
+			stopSignals()
 		case <-done:
 		}
 	}()
 	cancel := func() {
-		signal.Stop(sigCh)
-		close(done)
+		stopSignals()
 		cancelCause(nil)
 	}
 	return ctx, cancel