refactor: refactor callback server for in-server token exchange and testing

- Refactor callback handling to exchange the authorization code for tokens within the callback server, passing an exchange function as a parameter.
- Adjust error handling and response types to propagate token exchange errors and results instead of raw codes.
- Update related tests to use new exchange function injection, improving test coverage for error cases and exchange failures.
- Add comprehensive unit tests for callback server scenarios, including token exchange failure and better validation of test outcomes.
- Increase HTTP server write timeout from 10 to 15 seconds during the callback.

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

Author: appleboy

browser_flow.go

@@ -43,7 +43,11 @@ func performBrowserFlow(ctx context.Context) (*TokenStorage, bool, error) {
 	fmt.Println("Browser opened. Please complete authorization in your browser.")
 	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(callbackCtx context.Context, code string) (*TokenStorage, error) {
+			fmt.Println("Step 3: Exchanging authorization code for tokens...")
+			return exchangeCode(callbackCtx, code, pkce.Verifier)
+		})
 	if err != nil {
 		if errors.Is(err, ErrCallbackTimeout) {
 			// User opened the browser but didn't complete authorization in time.
@@ -54,14 +58,7 @@ func performBrowserFlow(ctx context.Context) (*TokenStorage, bool, error) {
 			)
 			return nil, false, nil
 		}
-		return nil, false, fmt.Errorf("authorization failed: %w", err)
-	}
-	fmt.Println("Authorization code received!")
-
-	fmt.Println("Step 3: Exchanging authorization code for tokens...")
-	storage, err := exchangeCode(ctx, code, pkce.Verifier)
-	if err != nil {
-		return nil, false, fmt.Errorf("token exchange failed: %w", err)
+		return nil, false, fmt.Errorf("authentication failed: %w", err)
 	}
 	storage.Flow = "browser"
 

callback.go

@@ -22,17 +22,20 @@ var ErrCallbackTimeout = fmt.Errorf("browser authorization timed out")
 
 // 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,
+// calls exchangeFn to exchange the code for tokens, and returns the resulting
+// TokenStorage (or an error).
 //
 // The server shuts itself down after the first request.
-func startCallbackServer(ctx context.Context, port int, expectedState string) (string, error) {
+func startCallbackServer(ctx context.Context, port int, expectedState string,
+	exchangeFn func(context.Context, string) (*TokenStorage, error),
+) (*TokenStorage, error) {
 	resultCh := make(chan callbackResult, 1)
 
 	var once sync.Once
@@ -69,20 +72,26 @@ func startCallbackServer(ctx context.Context, port int, expectedState string) (s
 			return
 		}
 
+		storage, 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: storage})
 	})
 
 	srv := &http.Server{
 		Addr:         fmt.Sprintf("127.0.0.1:%d", port),
 		Handler:      mux,
 		ReadTimeout:  10 * time.Second,
-		WriteTimeout: 10 * time.Second,
+		WriteTimeout: 15 * time.Second,
 	}
 
 	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)
 	}
 
 	go func() {
@@ -99,14 +108,14 @@ 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 <-time.After(callbackTimeout):
-		return "", fmt.Errorf("%w after %s", ErrCallbackTimeout, callbackTimeout)
+		return nil, fmt.Errorf("%w after %s", ErrCallbackTimeout, callbackTimeout)
 	}
 }
 

callback_test.go

@@ -10,31 +10,53 @@ import (
 	"time"
 )
 
+type callbackServerResult 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).
+// returns a channel that will receive the result (storage or error).
 func startCallbackServerAsync(
 	t *testing.T, ctx context.Context, port int, state string,
-) chan string {
+	exchangeFn func(context.Context, string) (*TokenStorage, error),
+) chan callbackServerResult {
 	t.Helper()
-	ch := make(chan string, 1)
+	ch := make(chan callbackServerResult, 1)
 	go func() {
-		code, err := startCallbackServer(ctx, port, state)
-		if err != nil {
-			ch <- "ERROR:" + err.Error()
-		} else {
-			ch <- code
-		}
+		storage, err := startCallbackServer(ctx, port, state, exchangeFn)
+		ch <- callbackServerResult{storage, err}
 	}()
 	// Give the server a moment to bind.
 	time.Sleep(50 * time.Millisecond)
 	return ch
 }
 
+// noExchangeFn returns an exchange function that fails the test if called.
+func noExchangeFn(t *testing.T) func(context.Context, string) (*TokenStorage, error) {
+	t.Helper()
+	return func(_ context.Context, _ string) (*TokenStorage, error) {
+		t.Error("exchangeFn should not be called")
+		return nil, fmt.Errorf("should not be called")
+	}
+}
+
+// stubExchangeFn returns an exchange function that validates the received code
+// and returns a minimal TokenStorage on success.
+func stubExchangeFn(wantCode string) func(context.Context, string) (*TokenStorage, error) {
+	return func(_ context.Context, gotCode string) (*TokenStorage, error) {
+		if gotCode != wantCode {
+			return nil, fmt.Errorf("unexpected code: got %q, want %q", gotCode, wantCode)
+		}
+		return &TokenStorage{AccessToken: "test-token"}, nil
+	}
+}
+
 func TestCallbackServer_Success(t *testing.T) {
 	const port = 19101
 	state := "test-state-success"
 
-	ch := startCallbackServerAsync(t, context.Background(), port, state)
+	ch := startCallbackServerAsync(t, context.Background(), port, state, stubExchangeFn("mycode123"))
 
 	callbackURL := fmt.Sprintf(
 		"http://127.0.0.1:%d/callback?code=mycode123&state=%s",
@@ -56,8 +78,11 @@ func TestCallbackServer_Success(t *testing.T) {
 
 	select {
 	case result := <-ch:
-		if result != "mycode123" {
-			t.Errorf("expected code mycode123, got: %s", result)
+		if result.err != nil {
+			t.Errorf("expected success, got error: %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")
@@ -68,7 +93,7 @@ func TestCallbackServer_StateMismatch(t *testing.T) {
 	const port = 19102
 	state := "expected-state"
 
-	ch := startCallbackServerAsync(t, context.Background(), port, state)
+	ch := startCallbackServerAsync(t, context.Background(), port, state, noExchangeFn(t))
 
 	callbackURL := fmt.Sprintf(
 		"http://127.0.0.1:%d/callback?code=mycode&state=wrong-state",
@@ -87,8 +112,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.Error("expected error for state mismatch, got nil")
 		}
 	case <-time.After(3 * time.Second):
 		t.Fatal("timed out waiting for callback result")
@@ -99,7 +124,7 @@ func TestCallbackServer_OAuthError(t *testing.T) {
 	const port = 19103
 	state := "state-for-error"
 
-	ch := startCallbackServerAsync(t, context.Background(), port, state)
+	ch := startCallbackServerAsync(t, context.Background(), port, state, noExchangeFn(t))
 
 	callbackURL := fmt.Sprintf(
 		"http://127.0.0.1:%d/callback?error=access_denied&error_description=User+denied&state=%s",
@@ -118,11 +143,48 @@ 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.Error("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")
+	}
+}
+
+func TestCallbackServer_ExchangeFailure(t *testing.T) {
+	const port = 19106
+	state := "state-for-exchange-failure"
+
+	ch := startCallbackServerAsync(t, context.Background(), port, state,
+		func(_ context.Context, _ string) (*TokenStorage, error) {
+			return nil, fmt.Errorf("unauthorized_client: unauthorized_client")
+		})
+
+	callbackURL := fmt.Sprintf(
+		"http://127.0.0.1:%d/callback?code=mycode&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 !strings.Contains(string(body), "Authorization Failed") {
+		t.Errorf("expected failure page for exchange error, got: %s", string(body))
+	}
+
+	select {
+	case result := <-ch:
+		if result.err == nil {
+			t.Error("expected error for exchange failure, got nil")
+		}
+		if !strings.Contains(result.err.Error(), "unauthorized_client") {
+			t.Errorf("expected error to mention unauthorized_client, got: %v", result.err)
 		}
 	case <-time.After(3 * time.Second):
 		t.Fatal("timed out waiting for callback result")
@@ -133,7 +195,7 @@ func TestCallbackServer_DoubleCallback(t *testing.T) {
 	const port = 19105
 	state := "test-state-double"
 
-	ch := startCallbackServerAsync(t, context.Background(), port, state)
+	ch := startCallbackServerAsync(t, context.Background(), port, state, stubExchangeFn("mycode"))
 
 	url := fmt.Sprintf("http://127.0.0.1:%d/callback?code=mycode&state=%s", port, state)
 
@@ -158,8 +220,11 @@ func TestCallbackServer_DoubleCallback(t *testing.T) {
 
 	select {
 	case result := <-ch:
-		if result != "mycode" {
-			t.Errorf("expected mycode, got: %s", result)
+		if result.err != nil {
+			t.Errorf("expected success, got error: %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")
@@ -170,7 +235,7 @@ func TestCallbackServer_MissingCode(t *testing.T) {
 	const port = 19104
 	state := "state-for-missing-code"
 
-	ch := startCallbackServerAsync(t, context.Background(), port, state)
+	ch := startCallbackServerAsync(t, context.Background(), port, state, noExchangeFn(t))
 
 	callbackURL := fmt.Sprintf(
 		"http://127.0.0.1:%d/callback?state=%s",
@@ -184,8 +249,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.Error("expected error for missing code, got nil")
 		}
 	case <-time.After(3 * time.Second):
 		t.Fatal("timed out waiting for callback result")