diff --git a/cmd/login.go b/cmd/login.go
index 7555035f..0cedf4a9 100644
--- a/cmd/login.go
+++ b/cmd/login.go
@@ -64,7 +64,7 @@ func runLogin(cmd *cobra.Command, args []string) error {
// Handle common error cases with helpful messages
if ctx.Err() == context.Canceled {
pterm.Info.Println("Authentication cancelled by user")
- return nil
+ return ctx.Err()
}
return fmt.Errorf("authentication failed: %w", err)
diff --git a/pkg/auth/oauth.go b/pkg/auth/oauth.go
index ebf6a9f3..f72dc47e 100644
--- a/pkg/auth/oauth.go
+++ b/pkg/auth/oauth.go
@@ -22,14 +22,14 @@ import (
)
//go:embed success.html
-var successHTMLTemplate string
+var authorizationReceivedHTMLTemplate string
//go:embed favicon.svg
var faviconSVG string
// Inlined as a data URI because the callback server shuts down before the browser could fetch a served icon.
-var successHTML = strings.Replace(
- successHTMLTemplate,
+var authorizationReceivedHTML = strings.Replace(
+ authorizationReceivedHTMLTemplate,
"__KERNEL_FAVICON__",
"data:image/svg+xml;base64,"+base64.StdEncoding.EncodeToString([]byte(faviconSVG)),
1,
@@ -45,6 +45,8 @@ const (
// OAuth scopes - openid for the MCP server flow
DefaultScope = "openid email"
+
+ defaultTokenExchangeTimeout = 30 * time.Second
)
// OAuthConfig represents the OAuth2 configuration
@@ -250,7 +252,7 @@ func (oc *OAuthConfig) StartOAuthFlow(ctx context.Context) (*TokenStorage, error
// Success page
w.Header().Set("Content-Type", "text/html")
w.WriteHeader(http.StatusOK)
- w.Write([]byte(successHTML))
+ w.Write([]byte(authorizationReceivedHTML))
// Pass both code and org_id to the channel using JSON encoding
result := AuthResult{
@@ -298,12 +300,14 @@ func (oc *OAuthConfig) StartOAuthFlow(ctx context.Context) (*TokenStorage, error
return nil, ctx.Err()
}
+ pterm.Info.Println("Browser authorization received; completing authentication...")
+
// Exchange authorization code for tokens
- return oc.exchangeCodeForTokens(ctx, authCode, orgID)
+ return oc.exchangeCodeForTokens(ctx, authCode, orgID, defaultTokenExchangeTimeout)
}
// exchangeCodeForTokens exchanges the authorization code for access and refresh tokens
-func (oc *OAuthConfig) exchangeCodeForTokens(ctx context.Context, code, orgID string) (*TokenStorage, error) {
+func (oc *OAuthConfig) exchangeCodeForTokens(ctx context.Context, code, orgID string, timeout time.Duration) (*TokenStorage, error) {
// Use PKCE verifier in token exchange, and include org_id if available
var opts []oauth2.AuthCodeOption
opts = append(opts, oauth2.SetAuthURLParam("code_verifier", oc.Verifier))
@@ -311,8 +315,14 @@ func (oc *OAuthConfig) exchangeCodeForTokens(ctx context.Context, code, orgID st
opts = append(opts, oauth2.SetAuthURLParam("org_id", orgID))
}
- token, err := oc.Config.Exchange(ctx, code, opts...)
+ exchangeCtx, cancel := context.WithTimeout(ctx, timeout)
+ defer cancel()
+
+ token, err := oc.Config.Exchange(exchangeCtx, code, opts...)
if err != nil {
+ if exchangeCtx.Err() == context.DeadlineExceeded {
+ return nil, fmt.Errorf("token exchange timed out after %s: %w", timeout, context.DeadlineExceeded)
+ }
return nil, fmt.Errorf("failed to exchange code for token: %w", err)
}
diff --git a/pkg/auth/oauth_test.go b/pkg/auth/oauth_test.go
index 059034e0..1a6864ef 100644
--- a/pkg/auth/oauth_test.go
+++ b/pkg/auth/oauth_test.go
@@ -1,6 +1,58 @@
package auth
-import "testing"
+import (
+ "context"
+ "errors"
+ "net/http"
+ "net/http/httptest"
+ "strings"
+ "testing"
+ "time"
+)
+
+func TestCallbackPageDoesNotClaimAuthenticationComplete(t *testing.T) {
+ t.Parallel()
+
+ if strings.Contains(authorizationReceivedHTML, "authentication successful") {
+ t.Fatal("callback page must not claim authentication succeeded before token exchange completes")
+ }
+ if !strings.Contains(authorizationReceivedHTML, "authorization received") {
+ t.Fatal("callback page should tell the user browser authorization was received")
+ }
+}
+
+func TestTokenExchangeTimesOut(t *testing.T) {
+ release := make(chan struct{})
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ select {
+ case <-r.Context().Done():
+ case <-release:
+ }
+ }))
+ defer server.Close()
+ defer close(release)
+
+ cfg, err := NewOAuthConfig()
+ if err != nil {
+ t.Fatalf("NewOAuthConfig() error = %v", err)
+ }
+ cfg.Config.Endpoint.TokenURL = server.URL
+
+ done := make(chan error, 1)
+ go func() {
+ _, err := cfg.exchangeCodeForTokens(context.Background(), "code", "org", 25*time.Millisecond)
+ done <- err
+ }()
+
+ select {
+ case err := <-done:
+ if err == nil || !strings.Contains(err.Error(), "token exchange timed out") || !errors.Is(err, context.DeadlineExceeded) {
+ t.Fatalf("exchange error = %v, want token exchange timeout", err)
+ }
+ case <-time.After(250 * time.Millisecond):
+ t.Fatal("token exchange did not honor its timeout")
+ }
+}
func TestNewOAuthConfigUsesAuthOverrides(t *testing.T) {
t.Setenv("KERNEL_AUTH_BASE_URL", "https://auth.dev.onkernel.com/")
diff --git a/pkg/auth/success.html b/pkg/auth/success.html
index 1fb33602..bdf8cd3b 100644
--- a/pkg/auth/success.html
+++ b/pkg/auth/success.html
@@ -1,7 +1,7 @@
- kernel cli - authenticated
+ kernel cli - authorization received