Skip to content

tests: Harden OAuth device, PKCE, and refresh-token flows#9732

Open
nishantmonu51 wants to merge 2 commits into
mainfrom
nishant/test-audit-oauth-device-pkce
Open

tests: Harden OAuth device, PKCE, and refresh-token flows#9732
nishantmonu51 wants to merge 2 commits into
mainfrom
nishant/test-audit-oauth-device-pkce

Conversation

@nishantmonu51

Copy link
Copy Markdown
Collaborator

Harden OAuth protocol boundaries across the admin server and CLI clients.

  • constrain dynamically registered clients and privileged scopes
  • validate exact redirect and PKCE exchange contracts
  • rotate refresh tokens atomically and reject replayed token families
  • enforce pending device user-code uniqueness in PostgreSQL
  • handle pending, denied, expired, throttled, and concurrent device polling
  • make CLI device and PKCE clients propagate cancellation and protocol failures
  • replace the probabilistic code-collision test with deterministic format and database-boundary coverage

Validation:

  • go test -short -race -count=1 -timeout=15m ./admin ./admin/server/auth ./cli/pkg/deviceauth ./cli/pkg/pkce

Checklist:

  • Covered by tests
  • Ran it and it works as intended
  • Reviewed the diff before requesting a review
  • Checked for unhandled edge cases
  • Linked the issues it closes
  • Checked if the docs need to be updated. No documentation update is required.
  • Intend to cherry-pick into the release branch
  • I am proud of this work!

@nishantmonu51 nishantmonu51 added the tests Changes to tests and test infrastructure label Jul 23, 2026
@nishantmonu51 nishantmonu51 changed the title Harden OAuth device, PKCE, and refresh-token flows tests: Harden OAuth device, PKCE, and refresh-token flows Jul 23, 2026
Comment on lines +66 to +71
scope := strings.TrimSpace(values.Get("scope"))
if scope == "" {
http.Error(w, "scope is required", http.StatusBadRequest)
return
}
scopes := strings.Fields(scope)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: I believe Fields trims leading/trailing whitespace, so no need for the separate TrimSpace

"go.uber.org/zap"
)

const maxOAuthRegistrationBodyBytes = 1 << 20

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider adding // 1 MB comment

Comment on lines +87 to +95
// Bound this public endpoint before decoding so a registration request cannot
// consume an arbitrary amount of server memory.
body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, maxOAuthRegistrationBodyBytes))
if err != nil {
var maxBytesErr *http.MaxBytesError
if errors.As(err, &maxBytesErr) {
http.Error(w, "request body too large", http.StatusRequestEntityTooLarge)
return
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any particular reason to enforce the body limit for this particular endpoint? Consider adding a middleware on all auth HTTP endpoints instead.

Also, isn't 1 MB a really large body for this endpoint? Would just have expected a few KB max.

Comment thread admin/server/auth/pkce.go
Comment on lines +39 to +42
if !hasGrantType(authClient.GrantTypes, authorizationCodeGrantType) {
http.Error(w, "client is not permitted to use authorization codes", http.StatusBadRequest)
return
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Worth double checking that all clients (including old ones) send this to ensure backwards compatible logins

Comment thread admin/server/auth/pkce.go
Comment on lines +338 to +343
// ValidateAuthToken has a short cache, so enforce expiry again at the
// rotation boundary in case the token expired after it was cached.
if userToken.ExpiresOn != nil && !userToken.ExpiresOn.After(time.Now()) {
http.Error(w, "invalid refresh token", http.StatusUnauthorized)
return
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't this be done inside ValidateAuthToken instead of here so it's applied for all callers?

Comment thread admin/server/auth/pkce.go
internalServerError(w, fmt.Errorf("failed to commit refresh token rotation, %w", err))
return
}
a.admin.PurgeAuthTokenCache()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this will purge all cached tokens, not just the one that was refreshed.

Might be better to only remove the refreshed one? Or perhaps just keep it simple and let the old token stay valid until the 10 second cache TTL expires? (Since it's an in-memory cache, other nodes may anyway still have it cached, so there's no security win by purging early.)

Comment thread admin/server/auth/pkce.go
Comment on lines +369 to +374
// Deleting inside the transaction atomically claims this token. Concurrent
// rotations can both validate it, but only one can delete and commit it.
err = a.admin.DB.DeleteUserAuthToken(txCtx, userToken.ID)
if err != nil {
if errors.Is(err, database.ErrNotFound) {
http.Error(w, "invalid refresh token", http.StatusUnauthorized)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure, but seems like the new logic has more error cases to handle than the previous logic. Have you checked the frontend and CLI handle these cases?

"client_id": []string{d.ClientID},
"scope": []string{"full_account"},
"redirect": []string{url.QueryEscape(redirectURL)},
"redirect": []string{redirectURL},

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure why we had url.QueryEscape on this previously, but did you confirm it's no longer needed in any cases?

Comment on lines +150 to +151
// Poll only after the server-provided interval. The wait is capped at the
// device-code expiry so a long interval cannot keep login alive past it.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new retry logic seems overly complex to me (now requiring both a separate enum type and a helper function for a simple poll loop). Consider if the complexity is necessary or if it can be simplified

Comment on lines 82 to +88
func (a *Authenticator) ExchangeCodeForToken(code string) (string, error) {
return a.ExchangeCodeForTokenContext(context.Background(), code)
}

// ExchangeCodeForTokenContext lets a caller tie the token request to its callback
// lifecycle, so cancelling the callback also cancels an in-flight exchange.
func (a *Authenticator) ExchangeCodeForTokenContext(ctx context.Context, code string) (string, error) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider making ExchangeCodeForToken accept a ctx instead of creating a separate function (ExchangeCodeForToken is anyway only used one place, so easy to refactor).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

tests Changes to tests and test infrastructure

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants