cognito: a Cognito PKCE login flow for CLIs

[lei-wego]

Adds github.com/wego/pkg/cognito: the Cognito authorization-code-with-PKCE login a command-line tool uses to obtain a human operator's tokens, plus a cognito/storage subpackage that caches the token set in the OS keychain.

This is the CLI-side counterpart to http/jwt. That package verifies a token arriving at a service; this one obtains a token at a terminal. Nothing in the repo covered the second half before.

Extracted from the payments repo's pay-admin CLI (wego/payments#2300), where it was written against these constraints from the start, so the move needed no redesign — package rename and import paths only.

API

type Config struct {
    AuthorizeURL, TokenURL, ClientID, RedirectURI string
    Scopes, AllowedDomain, CallbackAddr           string
    IdentityProvider                              string             // optional: skip Cognito's IdP chooser
    OpenBrowser                                   func(string) error // nil => OS default
    NoBrowser                                     bool               // suppress the launch, report the URL instead
    PromptURL                                     func(string) error // receives the URL when NoBrowser is set
    HTTPClient                                    *http.Client       // nil => sane default
    Now                                           func() time.Time   // nil => time.Now
}

func Login(ctx context.Context, cfg Config) (*TokenSet, error)
func Refresh(ctx context.Context, cfg Config, refreshToken string) (*TokenSet, error)

type TokenSet struct{ AccessToken, IDToken, RefreshToken string; ExpiresAt time.Time }
func (t *TokenSet) IsExpired(now time.Time) bool
func (t *TokenSet) Email() (string, error)

// cognito/storage
type Store interface{ Load(ns string) (*cognito.TokenSet, error); Save(...) error; Delete(...) error }
func NewKeyring(service string) Store
func NewMemory() Store

Headless sign-in

NoBrowser suppresses the launch and hands the authorize URL to PromptURL instead, for a headless shell, a terminal on a remote host, or a caller that wants to surface the URL its own way. Everything else is unchanged: same PKCE challenge, same state, code still delivered to the loopback listener.

PromptURL is a func rather than a bool-plus-stdout so the library never writes to a stream the caller did not choose — it can print, render a QR code, or hand the URL to another process. Setting NoBrowser without it is refused at validation.

One limit is documented on the field rather than papered over: the redirect still lands on CallbackAddr, so a browser on a different machine than the CLI needs that port forwarded (ssh -L). Suppressing the launch does not move where the code is delivered.

Two properties worth preserving

• No package-level mutable state. Every dependency — clock, HTTP client, browser opener, identity provider — arrives through Config, so one process can hold several environments live at once. http/jwt keeps its JWKS URL and header in package globals and can therefore serve exactly one issuer; payments' bo-refunds plan records that as the reason a second issuer became cross-team work. This package deliberately does not repeat the shape.
• Stdlib-only OAuth (bar wego/pkg/strings). Hand-rolling the exchange keeps every wire parameter visible and auditable, which matters more here than the convenience an OAuth library buys.

Security-relevant behaviour

Verifier and state both come from crypto/rand; state is compared with subtle.ConstantTimeCompare and a blank value on either side is a non-match. The challenge is S256. The callback listener binds loopback only and is single-use; neither the code nor error_description is interpolated into the served HTML. Email() parses the id_token without verifying its signature — correct here because the token came from the token endpoint over TLS or the caller's own keychain, and it is documented as such — and it cannot panic on malformed input. No token, verifier, or state appears in any error string.

Kept as one module

Splitting storage into its own module would force cognito to be tagged before storage could require it, and every other module here requires tagged siblings with no replace. Import paths are identical either way, so the only cost is that an OAuth-only consumer also pulls go-keyring.

Testing

go test ./... — cognito 95.4%, cognito/storage 94.7%. The uncovered lines are the ones that touch the OS: the browser exec and the real keychain adapter, both behind unexported seams so the logic around them is covered. NewMemory carries the full Store contract tests so nothing pops a keychain prompt in CI.

After merge

Tag cognito/v0.1.0 via ./auto_version, then wego/payments#2300 drops its in-repo copy and requires the tag. sdc-cli has the same duplicated flow and can adopt it in a follow-up.

[yanyi-wego]

Review withdrawn.

Withdrawn because internal review context was posted to this public repository in error.

[lei-wego]

Ready for another look. Since the last review this branch has three fixes and one addition, all with tests.

Token endpoint no longer follows redirects (ff2fbcb). Config.httpClient returned a client with no CheckRedirect, so Go's default applied and a 307/308 would re-send the request body. That body carries the authorization code and PKCE verifier on sign-in and the refresh token on renewal. It now returns a copy with redirects refused, so a caller-supplied client keeps its transport and timeout but cannot reinstate following, and the caller's own client is not mutated.

Token sets commit atomically (dda4783). Save wrote four separate keychain entries and Load gated on one of them, so a write that failed partway could leave a new refresh token beside an old access token, id token and expiry, which Load then returned as a live session. A single entry would be the obvious fix but does not fit: zalando/go-keyring shells out to security on macOS and rejects any command over 4096 bytes, which after its base64 expansion leaves roughly 3 KB of secret per entry. So fields keep their own entries and a small pointer entry names which of two slots is live. Save fills the inactive slot and then moves the pointer, which is one write and the only one that changes what Load sees.

Concurrent reads survive a commit (f338c57). The two-slot design introduced a race: Load chose a slot, a parallel process committed the replacement and cleared the old one, and the read failed on a missing entry. Load now re-reads the pointer on a failed field read and retries on the live slot if it moved, while still reporting a genuine inconsistency unchanged. Bounded at three attempts.

Paste-back sign-in (f308520). NoBrowser only moved where the authorize URL is shown; the redirect still had to reach the local callback port, so a host with no browser and no route to that port could not sign in at all. Config.ReadRedirect lets the caller return the URL the browser was redirected to instead of a listener receiving it. The device authorization grant would be the conventional answer and this provider does not implement it, so this is the substitute; it is still authorization code with PKCE, with only the redirect carried by hand. State is still checked and matters more here, since nothing else ties a pasted code to the attempt that started it.

Coverage is 95.3% on cognito and 94.4% on cognito/storage, with -race clean. The remaining uncovered statements are the three thin wrappers over the real OS keychain.

[lei-wego]

@yanyi-wego review bump — this is now the only thing blocking the payments chain.

Everything downstream is done: payments#2317 is approved and merged into #2300's branch, pennyworth#1525 is approved, merged and deployed to staging. payments#2300 carries both and is green with zero open threads. None of it can reach main until the cognito package lands here and gets tagged, because #2300's require points at a pseudo-version of this branch.

State here: build green on both jobs, zero open threads, 17 files, MERGEABLE but BEHIND main. It reads REVIEW_REQUIRED because your review was withdrawn from GitHub rather than never given — I was mis-describing that as "never reviewed" for a while, which was wrong.

Every finding from that review is fixed, including round 4's storage races, which were the real ones:

• Save/Save hybrid — two concurrent Save calls could commit a mixed credential set (one call's access token with another's ID and refresh tokens).
• ABA on the pointer — an A→B→A generation cycle could make Load accept a stale set as current.

Both were genuine holes in my own design, not test artifacts. I reproduced each one first and matched your reported outputs exactly before changing anything. The fix is commit-pointer storage where a generation name is never reused, which removes both failure modes rather than narrowing the window, and both of your reproductions are committed as regression tests. That replaced the earlier two-reusable-slot scheme.

Also in since your review: the paste-back sign-in path for an unreachable callback port (Config.ReadRedirect), after I confirmed three ways that Cognito has no RFC 8628 device grant — no device_authorization_endpoint in discovery, /oauth2/device_authorization answers 404, and a device_code grant request returns unsupported_grant_type.

Merge order once you approve: this → tag cognito/v0.1.0 → repoint the require in payments#2300 → merge #2300. I'll bring this up to main right before merge rather than now, so it doesn't churn CI while it waits.

[yanyi-wego]

Requesting changes for the credential transport blocker. Two non-blocking defects are noted inline.

[lei-wego]

@yanyi-wego round 5 is answered in b0738ab — all three fixed, each with the tests you asked for.

Finding	Fix	Tests
Blocking: cleartext endpoints, non-loopback bind	requireHTTPS on TokenURL (via validateClient, so Refresh too) and AuthorizeURL; requireLoopbackAddr / requireLoopbackRedirect on the local pair; new exported ErrInsecureEndpoint	9 rejection cases + a 3-host acceptance test for the loopback-http redirect
Callback delivered before the state check	expectedState passed into callbackHandler, compared before anything reaches the channel	invalid-then-valid regression + 3 no-delivery cases
ReadRedirect not cancellation-aware	Config.readRedirect(ctx) — goroutine plus select on ctx.Done()	blocked-reader cancellation test

All three are mutation-checked — reverting any one fails its own tests, not just the suite. -race is clean, and the full pkg pre-push suite passed.

I've resolved the blocking thread. I left the other two open on purpose, not because work is pending but because each carries a question back to you:

• On the callback: a tampered state now surfaces as the wait expiring rather than "state mismatch", since it is never delivered. That also means an error= redirect with no echoed state would time out instead of reporting. RFC 6749 requires state to be echoed when sent, so I accepted it — flagging in case you would rather the error branch stayed reachable.
• On cancellation: I kept ReadRedirect's signature as func() (string, error) rather than adding ctx. Adding it would be the better API and this is the last moment before cognito/v0.1.0 makes it breaking, but it would force a matching edit in payments#2300, which is green and waiting on approval. Your call — it is free right now and not later.

Two notes on things that changed beyond the fixes. The suite moved to httptest.NewTLSServer, because every existing test pointed at a plaintext loopback server and would otherwise have exercised the new rejection path instead of the flow; exempting loopback to keep the tests unchanged felt like the wrong trade. And I verified the payments CLI still passes the new rules rather than assuming — it binds 127.0.0.1:8100, redirects to http://localhost:8100/callback, and builds both endpoints as https://<domain>/oauth2/*.

Build is green. This is still the head of the chain: payments#2317 is merged into #2300's branch and pennyworth#1525 is merged and on staging, so #245 plus #2300 are the last two approvals before the tag and the merge.

[yanyi-wego]

One blocking cancellation-contract issue remains inline.

[yanyi-wego]

LGTM