refactor: refactor OAuth callback server for safe, complete token exchange

- Add .authgate-tokens.json to .gitignore for sensitive file exclusion
- Refactor callback server to handle the token exchange internally and return TokenStorage, not just the code
- Ensure the browser receives the true exchange result by holding the HTTP response open until exchange completion
- Update callback server timeout handling for write deadline with a dedicated constant
- Make callback server concurrency-safe and idempotent for browser retries by using sync.Once
- Update tests to validate TokenStorage and error handling, and add test coverage for exchange failures
- Change performAuthCodeFlow to delegate code exchange to the callback handler, streamlining error handling and response

Signed-off-by: appleboy <appleboy.tw@gmail.com>

Author: appleboy

.gitignore

@@ -31,3 +31,4 @@ go.work.sum
 # .idea/
 # .vscode/
 bin
+.authgate-tokens.json

callback.go

@@ -13,21 +13,32 @@ import (
 const (
 	// callbackTimeout is how long we wait for the browser to deliver the code.
 	callbackTimeout = 5 * time.Minute
+
+	// callbackWriteTimeout is the HTTP write deadline for the callback handler.
+	// It must exceed tokenExchangeTimeout to ensure the exchange result can be
+	// written back to the browser before the connection times out.
+	callbackWriteTimeout = 30 * time.Second
 )
 
 // callbackResult holds the outcome of the local callback round-trip.
 type callbackResult struct {
-	Code  string
-	Error string
-	Desc  string
+	Storage *TokenStorage
+	Error   string
+	Desc    string
 }
 
 // startCallbackServer starts a local HTTP server on the given port and waits
-// for the OAuth callback. It validates the returned state against expectedState
-// and returns the authorization code (or an error).
+// for the OAuth callback. It validates the returned state against expectedState,
+// then calls exchangeFn with the received authorization code. The HTTP response
+// is held open until exchangeFn returns so the browser reflects the true outcome.
 //
 // The server shuts itself down after the first request or when ctx is cancelled.
-func startCallbackServer(ctx context.Context, port int, expectedState string) (string, error) {
+func startCallbackServer(
+	ctx context.Context,
+	port int,
+	expectedState string,
+	exchangeFn func(ctx context.Context, code string) (*TokenStorage, error),
+) (*TokenStorage, error) {
 	resultCh := make(chan callbackResult, 1)
 
 	// sendResult delivers the result exactly once. Any concurrent or subsequent
@@ -38,6 +49,14 @@ func startCallbackServer(ctx context.Context, port int, expectedState string) (s
 		once.Do(func() { resultCh <- r })
 	}
 
+	// exchangeOnce ensures the token exchange runs at most once even when the
+	// browser retries the callback request.
+	var (
+		exchangeOnce    sync.Once
+		exchangeStorage *TokenStorage
+		exchangeErr     error
+	)
+
 	mux := http.NewServeMux()
 	mux.HandleFunc("/callback", func(w http.ResponseWriter, r *http.Request) {
 		q := r.URL.Query()
@@ -69,21 +88,32 @@ func startCallbackServer(ctx context.Context, port int, expectedState string) (s
 			return
 		}
 
+		// Hold the HTTP response open while exchanging the code for tokens so
+		// the browser reflects the true outcome (success or failure).
+		exchangeOnce.Do(func() {
+			exchangeStorage, exchangeErr = exchangeFn(r.Context(), code)
+		})
+		if exchangeErr != nil {
+			writeCallbackPage(w, false, "token_exchange_failed", exchangeErr.Error())
+			sendResult(callbackResult{Error: "token_exchange_failed", Desc: exchangeErr.Error()})
+			return
+		}
+
 		writeCallbackPage(w, true, "", "")
-		sendResult(callbackResult{Code: code})
+		sendResult(callbackResult{Storage: exchangeStorage})
 	})
 
 	srv := &http.Server{
 		Addr:         fmt.Sprintf("127.0.0.1:%d", port),
 		Handler:      mux,
 		ReadTimeout:  10 * time.Second,
-		WriteTimeout: 10 * time.Second,
+		WriteTimeout: callbackWriteTimeout,
 	}
 
 	// Use a listener so we can report the actual bound port.
 	ln, err := (&net.ListenConfig{}).Listen(ctx, "tcp", srv.Addr)
 	if err != nil {
-		return "", fmt.Errorf("failed to start callback server on port %d: %w", port, err)
+		return nil, fmt.Errorf("failed to start callback server on port %d: %w", port, err)
 	}
 
 	// Serve in background; shut down after receiving the result.
@@ -102,17 +132,17 @@ func startCallbackServer(ctx context.Context, port int, expectedState string) (s
 	case result := <-resultCh:
 		if result.Error != "" {
 			if result.Desc != "" {
-				return "", fmt.Errorf("%s: %s", result.Error, result.Desc)
+				return nil, fmt.Errorf("%s: %s", result.Error, result.Desc)
 			}
-			return "", fmt.Errorf("%s", result.Error)
+			return nil, fmt.Errorf("%s", result.Error)
 		}
-		return result.Code, nil
+		return result.Storage, nil
 
 	case <-ctx.Done():
-		return "", ctx.Err()
+		return nil, ctx.Err()
 
 	case <-time.After(callbackTimeout):
-		return "", fmt.Errorf("timed out waiting for browser authorization (%s)", callbackTimeout)
+		return nil, fmt.Errorf("timed out waiting for browser authorization (%s)", callbackTimeout)
 	}
 }
 

callback_test.go

@@ -10,29 +10,48 @@ import (
 	"time"
 )
 
+type serverResult struct {
+	storage *TokenStorage
+	err     error
+}
+
 // startCallbackServerAsync starts the callback server in a goroutine and
-// returns a channel that will receive the authorization code (or error string).
-func startCallbackServerAsync(t *testing.T, port int, state string) chan string {
+// returns a channel that will receive the final result (storage or error).
+func startCallbackServerAsync(
+	t *testing.T,
+	port int,
+	state string,
+	exchangeFn func(ctx context.Context, code string) (*TokenStorage, error),
+) chan serverResult {
 	t.Helper()
-	ch := make(chan string, 1)
+	ch := make(chan serverResult, 1)
 	go func() {
-		code, err := startCallbackServer(context.Background(), port, state)
-		if err != nil {
-			ch <- "ERROR:" + err.Error()
-		} else {
-			ch <- code
-		}
+		storage, err := startCallbackServer(context.Background(), port, state, exchangeFn)
+		ch <- serverResult{storage: storage, err: err}
 	}()
 	// Give the server a moment to bind.
 	time.Sleep(50 * time.Millisecond)
 	return ch
 }
 
+// mockExchangeFn returns an exchangeFn that succeeds with a stub TokenStorage.
+func mockExchangeFn(t *testing.T) func(ctx context.Context, code string) (*TokenStorage, error) {
+	t.Helper()
+	return func(_ context.Context, _ string) (*TokenStorage, error) {
+		return &TokenStorage{
+			AccessToken:  "mock-access-token",
+			RefreshToken: "mock-refresh-token",
+			TokenType:    "Bearer",
+			ExpiresAt:    time.Now().Add(time.Hour),
+		}, nil
+	}
+}
+
 func TestCallbackServer_Success(t *testing.T) {
 	const port = 19001
 	state := "test-state-success"
 
-	ch := startCallbackServerAsync(t, port, state)
+	ch := startCallbackServerAsync(t, port, state, mockExchangeFn(t))
 
 	// Simulate the browser redirect.
 	callbackURL := fmt.Sprintf(
@@ -53,11 +72,57 @@ func TestCallbackServer_Success(t *testing.T) {
 		t.Errorf("expected success page, got: %s", string(body))
 	}
 
-	// Check code returned to CLI.
+	// Check that storage is returned to the CLI.
+	select {
+	case result := <-ch:
+		if result.err != nil {
+			t.Errorf("expected no error, got: %v", result.err)
+		}
+		if result.storage == nil || result.storage.AccessToken != "mock-access-token" {
+			t.Errorf("unexpected storage: %+v", result.storage)
+		}
+	case <-time.After(3 * time.Second):
+		t.Fatal("timed out waiting for callback result")
+	}
+}
+
+func TestCallbackServer_ExchangeFailure(t *testing.T) {
+	const port = 19006
+	state := "test-state-exchange-fail"
+
+	failFn := func(_ context.Context, _ string) (*TokenStorage, error) {
+		return nil, fmt.Errorf("server returned status 400: invalid_grant")
+	}
+	ch := startCallbackServerAsync(t, port, state, failFn)
+
+	callbackURL := fmt.Sprintf(
+		"http://127.0.0.1:%d/callback?code=badcode&state=%s",
+		port, state,
+	)
+	resp, err := http.Get(callbackURL) //nolint:noctx,gosec
+	if err != nil {
+		t.Fatalf("GET callback failed: %v", err)
+	}
+	defer resp.Body.Close()
+
+	body, _ := io.ReadAll(resp.Body)
+	if resp.StatusCode != http.StatusOK {
+		t.Errorf("unexpected status %d", resp.StatusCode)
+	}
+	if !strings.Contains(string(body), "Authorization Failed") {
+		t.Errorf("expected failure page, got: %s", string(body))
+	}
+	if !strings.Contains(string(body), "invalid_grant") {
+		t.Errorf("expected error detail in page, got: %s", string(body))
+	}
+
 	select {
 	case result := <-ch:
-		if result != "mycode123" {
-			t.Errorf("expected code mycode123, got: %s", result)
+		if result.err == nil {
+			t.Error("expected an error, got nil")
+		}
+		if result.storage != nil {
+			t.Errorf("expected nil storage, got: %+v", result.storage)
 		}
 	case <-time.After(3 * time.Second):
 		t.Fatal("timed out waiting for callback result")
@@ -68,7 +133,7 @@ func TestCallbackServer_StateMismatch(t *testing.T) {
 	const port = 19002
 	state := "expected-state"
 
-	ch := startCallbackServerAsync(t, port, state)
+	ch := startCallbackServerAsync(t, port, state, nil)
 
 	callbackURL := fmt.Sprintf(
 		"http://127.0.0.1:%d/callback?code=mycode&state=wrong-state",
@@ -87,8 +152,8 @@ func TestCallbackServer_StateMismatch(t *testing.T) {
 
 	select {
 	case result := <-ch:
-		if !strings.HasPrefix(result, "ERROR:") {
-			t.Errorf("expected error for state mismatch, got: %s", result)
+		if result.err == nil {
+			t.Errorf("expected error for state mismatch, got nil")
 		}
 	case <-time.After(3 * time.Second):
 		t.Fatal("timed out waiting for callback result")
@@ -99,7 +164,7 @@ func TestCallbackServer_OAuthError(t *testing.T) {
 	const port = 19003
 	state := "state-for-error"
 
-	ch := startCallbackServerAsync(t, port, state)
+	ch := startCallbackServerAsync(t, port, state, nil)
 
 	callbackURL := fmt.Sprintf(
 		"http://127.0.0.1:%d/callback?error=access_denied&error_description=User+denied&state=%s",
@@ -118,11 +183,11 @@ func TestCallbackServer_OAuthError(t *testing.T) {
 
 	select {
 	case result := <-ch:
-		if !strings.HasPrefix(result, "ERROR:") {
-			t.Errorf("expected error for access_denied, got: %s", result)
+		if result.err == nil {
+			t.Errorf("expected error for access_denied, got nil")
 		}
-		if !strings.Contains(result, "access_denied") {
-			t.Errorf("expected error to mention access_denied, got: %s", result)
+		if !strings.Contains(result.err.Error(), "access_denied") {
+			t.Errorf("expected error to mention access_denied, got: %v", result.err)
 		}
 	case <-time.After(3 * time.Second):
 		t.Fatal("timed out waiting for callback result")
@@ -137,7 +202,7 @@ func TestCallbackServer_DoubleCallback(t *testing.T) {
 	const port = 19005
 	state := "test-state-double"
 
-	ch := startCallbackServerAsync(t, port, state)
+	ch := startCallbackServerAsync(t, port, state, mockExchangeFn(t))
 
 	url := fmt.Sprintf("http://127.0.0.1:%d/callback?code=mycode&state=%s", port, state)
 
@@ -163,11 +228,14 @@ func TestCallbackServer_DoubleCallback(t *testing.T) {
 		}
 	}
 
-	// startCallbackServer must also return promptly.
+	// startCallbackServer must also return promptly with a valid storage.
 	select {
 	case result := <-ch:
-		if result != "mycode" {
-			t.Errorf("expected mycode, got: %s", result)
+		if result.err != nil {
+			t.Errorf("expected no error, got: %v", result.err)
+		}
+		if result.storage == nil {
+			t.Error("expected non-nil storage")
 		}
 	case <-time.After(3 * time.Second):
 		t.Fatal("timed out waiting for callback result")
@@ -178,7 +246,7 @@ func TestCallbackServer_MissingCode(t *testing.T) {
 	const port = 19004
 	state := "state-for-missing-code"
 
-	ch := startCallbackServerAsync(t, port, state)
+	ch := startCallbackServerAsync(t, port, state, nil)
 
 	// Correct state but no code parameter.
 	callbackURL := fmt.Sprintf(
@@ -193,8 +261,8 @@ func TestCallbackServer_MissingCode(t *testing.T) {
 
 	select {
 	case result := <-ch:
-		if !strings.HasPrefix(result, "ERROR:") {
-			t.Errorf("expected error for missing code, got: %s", result)
+		if result.err == nil {
+			t.Errorf("expected error for missing code, got nil")
 		}
 	case <-time.After(3 * time.Second):
 		t.Fatal("timed out waiting for callback result")

main.go

@@ -339,20 +339,20 @@ func performAuthCodeFlow(ctx context.Context) (*TokenStorage, error) {
 		fmt.Println("Browser opened. Please complete authorization in your browser.")
 	}
 
-	// Start local callback server and wait for the code.
+	// Start local callback server. The exchange runs inside the handler so the
+	// browser sees the true outcome (success or failure) rather than a premature
+	// success page.
 	fmt.Printf("Step 2: Waiting for callback on http://localhost:%d/callback ...\n", callbackPort)
-	code, err := startCallbackServer(ctx, callbackPort, state)
+	storage, err := startCallbackServer(ctx, callbackPort, state,
+		func(cbCtx context.Context, code string) (*TokenStorage, error) {
+			fmt.Println("Authorization code received!")
+			fmt.Println("Step 3: Exchanging authorization code for tokens...")
+			return exchangeCode(cbCtx, code, pkce.Verifier)
+		},
+	)
 	if err != nil {
 		return nil, fmt.Errorf("authorization failed: %w", err)
 	}
-	fmt.Println("Authorization code received!")
-
-	// Exchange code for tokens.
-	fmt.Println("Step 3: Exchanging authorization code for tokens...")
-	storage, err := exchangeCode(ctx, code, pkce.Verifier)
-	if err != nil {
-		return nil, fmt.Errorf("token exchange failed: %w", err)
-	}
 
 	if err := saveTokens(storage); err != nil {
 		fmt.Printf("Warning: Failed to save tokens: %v\n", err)