diff --git a/cognito/browser.go b/cognito/browser.go
new file mode 100644
index 0000000..3473419
--- /dev/null
+++ b/cognito/browser.go
@@ -0,0 +1,28 @@
+package cognito
+
+import (
+ "os/exec"
+ "runtime"
+)
+
+// browserCommand returns the command and arguments that open rawURL in the
+// platform's default browser.
+func browserCommand(goos, rawURL string) (string, []string) {
+ switch goos {
+ case "darwin":
+ return "open", []string{rawURL}
+ case "windows":
+ return "rundll32", []string{"url.dll,FileProtocolHandler", rawURL}
+ default:
+ return "xdg-open", []string{rawURL}
+ }
+}
+
+// openBrowser launches the platform's default browser. This is the one
+// genuinely untestable function in the package - a unit test must not spawn a
+// real browser - so it is kept to a single statement over browserCommand,
+// which is tested. Callers that need a seam set Config.OpenBrowser instead.
+func openBrowser(rawURL string) error {
+ name, args := browserCommand(runtime.GOOS, rawURL)
+ return exec.Command(name, args...).Start()
+}
diff --git a/cognito/browser_internal_test.go b/cognito/browser_internal_test.go
new file mode 100644
index 0000000..4957a56
--- /dev/null
+++ b/cognito/browser_internal_test.go
@@ -0,0 +1,46 @@
+package cognito
+
+import (
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+)
+
+func TestBrowserCommand(t *testing.T) {
+ const authorizeURL = "https://cognito.test/oauth2/authorize?client_id=x"
+
+ tests := []struct {
+ name string
+ givenGOOS string
+ wantName string
+ wantArgs []string
+ }{
+ {
+ name: "darwin uses open",
+ givenGOOS: "darwin",
+ wantName: "open",
+ wantArgs: []string{authorizeURL},
+ },
+ {
+ name: "windows goes through the protocol handler",
+ givenGOOS: "windows",
+ wantName: "rundll32",
+ wantArgs: []string{"url.dll,FileProtocolHandler", authorizeURL},
+ },
+ {
+ name: "other platforms fall back to xdg-open",
+ givenGOOS: "linux",
+ wantName: "xdg-open",
+ wantArgs: []string{authorizeURL},
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ gotName, gotArgs := browserCommand(tt.givenGOOS, authorizeURL)
+
+ assert.Equal(t, tt.wantName, gotName)
+ assert.Equal(t, tt.wantArgs, gotArgs)
+ })
+ }
+}
diff --git a/cognito/callback.go b/cognito/callback.go
new file mode 100644
index 0000000..ae9d29f
--- /dev/null
+++ b/cognito/callback.go
@@ -0,0 +1,163 @@
+package cognito
+
+import (
+ "context"
+ "errors"
+ "fmt"
+ "net"
+ "net/http"
+ "net/url"
+ "time"
+
+ wegostrings "github.com/wego/pkg/strings"
+)
+
+const (
+ // defaultCallbackTimeout bounds how long we hold the loopback listener
+ // open waiting for the operator to finish signing in.
+ defaultCallbackTimeout = 5 * time.Minute
+
+ callbackReadHeaderTimeout = 5 * time.Second
+ callbackShutdownTimeout = 2 * time.Second
+)
+
+const successHTML = `
+Signed in
+
+Signed in
+You can close this window and return to your terminal.
`
+
+// failureHTML deliberately carries no detail from the request. The operator
+// reads the actual reason in their terminal, where it cannot be reflected back
+// into a page, so nothing from the redirect is ever interpolated into HTML.
+const failureHTML = `
+Sign-in failed
+
+Sign-in failed
+Return to your terminal for the reason, and try again.
`
+
+// callbackResult is the single outcome a callback server reports.
+type callbackResult struct {
+ code string
+ state string
+ err error
+}
+
+// callbackServer is a single-use loopback listener for the OAuth redirect.
+type callbackServer struct {
+ results <-chan callbackResult
+ srv *http.Server
+}
+
+// callbackPath is the request path the redirect URI points at, falling back to
+// root when the URI carries no path (or cannot be parsed at all).
+func callbackPath(redirectURI string) string {
+ parsed, err := url.Parse(redirectURI)
+ if err != nil || wegostrings.IsBlank(parsed.Path) {
+ return "/"
+ }
+ return parsed.Path
+}
+
+// startCallbackServer binds addr and serves the OAuth redirect at path.
+//
+// A bind failure is terminal on purpose: the Cognito app client registers
+// exactly one callback URL, so listening somewhere else would produce a
+// redirect the authorization server refuses. Fail loudly instead.
+func startCallbackServer(addr, path string) (*callbackServer, error) {
+ listener, err := net.Listen("tcp", addr)
+ if err != nil {
+ return nil, fmt.Errorf(
+ "bind the sign-in callback listener on %s: %w (that address must be free - the Cognito app client registers exactly one callback URL, so another port is not an option; close whatever holds it and retry)",
+ addr, err)
+ }
+
+ results := make(chan callbackResult, 1)
+ mux := http.NewServeMux()
+ mux.HandleFunc(path, callbackHandler(results))
+
+ srv := &http.Server{
+ Handler: mux,
+ ReadHeaderTimeout: callbackReadHeaderTimeout,
+ }
+ go func() {
+ _ = srv.Serve(listener)
+ }()
+
+ return &callbackServer{results: results, srv: srv}, nil
+}
+
+// callbackHandler serves the redirect, renders a minimal page for the
+// operator, and reports the outcome exactly once.
+func callbackHandler(results chan<- callbackResult) http.HandlerFunc {
+ return func(w http.ResponseWriter, r *http.Request) {
+ query := r.URL.Query()
+ w.Header().Set("Content-Type", "text/html; charset=utf-8")
+
+ if authErr := query.Get("error"); wegostrings.IsNotBlank(authErr) {
+ message := authErr
+ if desc := query.Get("error_description"); wegostrings.IsNotBlank(desc) {
+ message += ": " + desc
+ }
+ w.WriteHeader(http.StatusBadRequest)
+ _, _ = fmt.Fprint(w, failureHTML)
+ deliver(results, callbackResult{err: fmt.Errorf("authorization server rejected the sign-in: %s", message)})
+ return
+ }
+
+ code := query.Get("code")
+ if wegostrings.IsBlank(code) {
+ w.WriteHeader(http.StatusBadRequest)
+ _, _ = fmt.Fprint(w, failureHTML)
+ deliver(results, callbackResult{err: errors.New("the sign-in redirect carried no authorization code")})
+ return
+ }
+
+ _, _ = fmt.Fprint(w, successHTML)
+ deliver(results, callbackResult{code: code, state: query.Get("state")})
+ }
+}
+
+// deliver reports the first result and silently drops later ones. The flow is
+// single-use: a reloaded browser tab must neither block the handler nor
+// overwrite the outcome we already acted on.
+func deliver(results chan<- callbackResult, result callbackResult) {
+ select {
+ case results <- result:
+ default:
+ }
+}
+
+// wait blocks for the callback, honouring both ctx cancellation and timeout.
+func (c *callbackServer) wait(ctx context.Context, timeout time.Duration) (code, state string, err error) {
+ if timeout <= 0 {
+ timeout = defaultCallbackTimeout
+ }
+
+ timer := time.NewTimer(timeout)
+ defer timer.Stop()
+
+ select {
+ case <-ctx.Done():
+ return "", "", fmt.Errorf("waiting for the sign-in callback: %w", ctx.Err())
+ case <-timer.C:
+ return "", "", fmt.Errorf("timed out after %s waiting for the sign-in callback", timeout)
+ case result := <-c.results:
+ if result.err != nil {
+ return "", "", result.err
+ }
+ return result.code, result.state, nil
+ }
+}
+
+// shutdown releases the loopback listener.
+func (c *callbackServer) shutdown() {
+ if c.srv == nil {
+ return
+ }
+ ctx, cancel := context.WithTimeout(context.Background(), callbackShutdownTimeout)
+ defer cancel()
+ _ = c.srv.Shutdown(ctx)
+}
diff --git a/cognito/callback_internal_test.go b/cognito/callback_internal_test.go
new file mode 100644
index 0000000..930c3cf
--- /dev/null
+++ b/cognito/callback_internal_test.go
@@ -0,0 +1,228 @@
+package cognito
+
+import (
+ "context"
+ "net"
+ "net/http"
+ "net/http/httptest"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+ wegostrings "github.com/wego/pkg/strings"
+)
+
+func TestCallbackPath(t *testing.T) {
+ tests := []struct {
+ name string
+ givenURI string
+ want string
+ }{
+ {name: "unparseable falls back to root", givenURI: "http://[::1", want: "/"},
+ {name: "no path falls back to root", givenURI: "http://localhost:8110", want: "/"},
+ {name: "blank falls back to root", givenURI: "", want: "/"},
+ {name: "explicit path is used", givenURI: "http://localhost:8110/callback", want: "/callback"},
+ {name: "nested path is used", givenURI: "http://127.0.0.1:8110/oauth/cb", want: "/oauth/cb"},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ assert.Equal(t, tt.want, callbackPath(tt.givenURI))
+ })
+ }
+}
+
+func TestCallbackHandler(t *testing.T) {
+ tests := []struct {
+ name string
+ givenQuery string
+ wantStatus int
+ wantCode string
+ wantState string
+ wantErrContains string
+ }{
+ {
+ name: "authorize error is surfaced with its description",
+ givenQuery: "error=access_denied&error_description=user+said+no",
+ wantStatus: http.StatusBadRequest,
+ wantErrContains: "user said no",
+ },
+ {
+ name: "authorize error without a description still reports",
+ givenQuery: "error=server_error",
+ wantStatus: http.StatusBadRequest,
+ wantErrContains: "server_error",
+ },
+ {
+ name: "missing code is rejected",
+ givenQuery: "state=abc",
+ wantStatus: http.StatusBadRequest,
+ wantErrContains: "no authorization code",
+ },
+ {
+ name: "code and state are captured",
+ givenQuery: "code=the-code&state=the-state",
+ wantStatus: http.StatusOK,
+ wantCode: "the-code",
+ wantState: "the-state",
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ ch := make(chan callbackResult, 1)
+ rec := httptest.NewRecorder()
+ req := httptest.NewRequest(http.MethodGet, "/callback?"+tt.givenQuery, nil)
+
+ callbackHandler(ch)(rec, req)
+
+ assert.Equal(t, tt.wantStatus, rec.Code)
+ assert.Contains(t, rec.Header().Get("Content-Type"), "text/html")
+
+ select {
+ case got := <-ch:
+ if wegostrings.IsNotEmpty(tt.wantErrContains) {
+ require.Error(t, got.err)
+ assert.Contains(t, got.err.Error(), tt.wantErrContains)
+ return
+ }
+ require.NoError(t, got.err)
+ assert.Equal(t, tt.wantCode, got.code)
+ assert.Equal(t, tt.wantState, got.state)
+ default:
+ t.Fatal("handler did not emit a callbackResult")
+ }
+ })
+ }
+}
+
+func TestCallbackHandler_IsSingleUse(t *testing.T) {
+ ch := make(chan callbackResult, 1)
+ handler := callbackHandler(ch)
+
+ for range 3 {
+ rec := httptest.NewRecorder()
+ handler(rec, httptest.NewRequest(http.MethodGet, "/callback?code=c&state=s", nil))
+ }
+
+ require.Len(t, ch, 1, "a redelivered callback must neither block nor enqueue a second result")
+}
+
+func TestStartCallbackServer_PortAlreadyBoundFailsFast(t *testing.T) {
+ ln, err := net.Listen("tcp", "127.0.0.1:0")
+ require.NoError(t, err)
+ t.Cleanup(func() { _ = ln.Close() })
+
+ addr := ln.Addr().String()
+ cs, err := startCallbackServer(addr, "/callback")
+
+ require.Error(t, err, "binding an occupied port must fail rather than silently pick another")
+ assert.Nil(t, cs)
+ assert.Contains(t, err.Error(), addr, "the error must name the address the operator has to free")
+}
+
+func TestStartCallbackServer_ServesTheCallback(t *testing.T) {
+ addr := mustFreeAddr(t)
+ cs, err := startCallbackServer(addr, "/callback")
+ require.NoError(t, err)
+ t.Cleanup(cs.shutdown)
+
+ mustGet(t, "http://"+addr+"/callback?code=abc&state=xyz")
+
+ code, state, err := cs.wait(context.Background(), 2*time.Second)
+ require.NoError(t, err)
+ assert.Equal(t, "abc", code)
+ assert.Equal(t, "xyz", state)
+}
+
+func TestCallbackServer_Wait(t *testing.T) {
+ tests := []struct {
+ name string
+ givenResult *callbackResult
+ givenTimeout time.Duration
+ givenCancel bool
+ wantCode string
+ wantErrContains string
+ }{
+ {
+ name: "timeout is reported",
+ givenTimeout: 10 * time.Millisecond,
+ wantErrContains: "timed out",
+ },
+ {
+ name: "cancellation is reported",
+ givenTimeout: time.Minute,
+ givenCancel: true,
+ wantErrContains: "context canceled",
+ },
+ {
+ name: "callback error is propagated",
+ givenResult: &callbackResult{err: stubError("authorize: access_denied")},
+ givenTimeout: time.Minute,
+ wantErrContains: "access_denied",
+ },
+ {
+ name: "code and state are returned",
+ givenResult: &callbackResult{code: "c", state: "s"},
+ givenTimeout: time.Minute,
+ wantCode: "c",
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ ch := make(chan callbackResult, 1)
+ if tt.givenResult != nil {
+ ch <- *tt.givenResult
+ }
+ cs := &callbackServer{results: ch}
+
+ ctx, cancel := context.WithCancel(context.Background())
+ defer cancel()
+ if tt.givenCancel {
+ cancel()
+ }
+
+ code, _, err := cs.wait(ctx, tt.givenTimeout)
+ if wegostrings.IsNotEmpty(tt.wantErrContains) {
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), tt.wantErrContains)
+ return
+ }
+ require.NoError(t, err)
+ assert.Equal(t, tt.wantCode, code)
+ })
+ }
+}
+
+func TestCallbackServer_ShutdownIsSafeWithoutAServer(t *testing.T) {
+ cs := &callbackServer{results: make(chan callbackResult, 1)}
+ assert.NotPanics(t, cs.shutdown)
+}
+
+// mustFreeAddr reserves and releases a loopback port, returning its address.
+func mustFreeAddr(t *testing.T) string {
+ t.Helper()
+ ln, err := net.Listen("tcp", "127.0.0.1:0")
+ require.NoError(t, err)
+ addr := ln.Addr().String()
+ require.NoError(t, ln.Close())
+ return addr
+}
+
+// mustGet issues a throwaway GET, standing in for the operator's browser.
+func mustGet(t *testing.T, rawURL string) {
+ t.Helper()
+ req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, rawURL, nil)
+ require.NoError(t, err)
+ //nolint:gosec // G704: a loopback URL this test built itself.
+ resp, err := http.DefaultClient.Do(req)
+ require.NoError(t, err)
+ require.NoError(t, resp.Body.Close())
+}
+
+// stubError is a minimal error for table cases that need a pre-baked failure.
+type stubError string
+
+func (e stubError) Error() string { return string(e) }
diff --git a/cognito/go.mod b/cognito/go.mod
new file mode 100644
index 0000000..5f53e7e
--- /dev/null
+++ b/cognito/go.mod
@@ -0,0 +1,18 @@
+module github.com/wego/pkg/cognito
+
+go 1.25.0
+
+require (
+ github.com/stretchr/testify v1.11.1
+ github.com/wego/pkg/strings v0.1.2
+ github.com/zalando/go-keyring v0.2.8
+)
+
+require (
+ github.com/danieljoos/wincred v1.2.3 // indirect
+ github.com/davecgh/go-spew v1.1.1 // indirect
+ github.com/godbus/dbus/v5 v5.2.2 // indirect
+ github.com/pmezard/go-difflib v1.0.0 // indirect
+ golang.org/x/sys v0.27.0 // indirect
+ gopkg.in/yaml.v3 v3.0.1 // indirect
+)
diff --git a/cognito/go.sum b/cognito/go.sum
new file mode 100644
index 0000000..3bcf314
--- /dev/null
+++ b/cognito/go.sum
@@ -0,0 +1,24 @@
+github.com/danieljoos/wincred v1.2.3 h1:v7dZC2x32Ut3nEfRH+vhoZGvN72+dQ/snVXo/vMFLdQ=
+github.com/danieljoos/wincred v1.2.3/go.mod h1:6qqX0WNrS4RzPZ1tnroDzq9kY3fu1KwE7MRLQK4X0bs=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/godbus/dbus/v5 v5.2.2 h1:TUR3TgtSVDmjiXOgAAyaZbYmIeP3DPkld3jgKGV8mXQ=
+github.com/godbus/dbus/v5 v5.2.2/go.mod h1:3AAv2+hPq5rdnr5txxxRwiGjPXamgoIHgz9FPBfOp3c=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
+github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
+github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
+github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
+github.com/wego/pkg/pointer v0.1.2 h1:KghXP86aWukvpSVPQ+Fg7YOkW8p8kyXcuOAvWVX1RUk=
+github.com/wego/pkg/pointer v0.1.2/go.mod h1:TincAjFVHSyuZ05qnSP4APqs+eg+adjOfZV6VH0+CUA=
+github.com/wego/pkg/strings v0.1.2 h1:sFfYDrC90JI43UCs4fnlxtuG/kXm6WJPIfiND1aqrQE=
+github.com/wego/pkg/strings v0.1.2/go.mod h1:nS3SS/em72lPIJYtgaLE5XgboOLOS8lnsYRAD8U2Ptc=
+github.com/zalando/go-keyring v0.2.8 h1:6sD/Ucpl7jNq10rM2pgqTs0sZ9V3qMrqfIIy5YPccHs=
+github.com/zalando/go-keyring v0.2.8/go.mod h1:tsMo+VpRq5NGyKfxoBVjCuMrG47yj8cmakZDO5QGii0=
+golang.org/x/sys v0.27.0 h1:wBqf8DvsY9Y/2P8gAfPDEYNuS30J4lPHJxXSb/nJZ+s=
+golang.org/x/sys v0.27.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
diff --git a/cognito/oauth.go b/cognito/oauth.go
new file mode 100644
index 0000000..8827fdb
--- /dev/null
+++ b/cognito/oauth.go
@@ -0,0 +1,418 @@
+// Package cognito implements the Cognito authorization-code-with-PKCE login
+// flow a command-line tool uses to obtain a human operator's tokens.
+//
+// It is built for CLIs rather than services: the flow opens a browser, receives
+// the authorization code on a loopback listener, and hands back the token set
+// for the caller to cache. A service verifying an incoming token wants
+// github.com/wego/pkg/http/jwt instead.
+//
+// Two deliberate constraints shape this package:
+//
+// - It depends on nothing outside the standard library (bar Wego's string
+// helpers). Hand-rolling the OAuth exchange keeps every wire parameter
+// visible and auditable, which matters more here than the convenience an
+// OAuth library would buy.
+// - It holds no package-level mutable state. Every dependency - the clock,
+// the HTTP client, the browser opener - arrives through Config, so tests
+// and callers never race over shared globals.
+package cognito
+
+import (
+ "context"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "io"
+ "net/http"
+ "net/url"
+ "strings"
+ "time"
+
+ wegostrings "github.com/wego/pkg/strings"
+)
+
+const (
+ // placeholderPrefix marks a Cognito value that ops has not filled in yet.
+ placeholderPrefix = "REPLACE_WITH_"
+
+ defaultHTTPTimeout = 30 * time.Second
+ maxTokenResponseBytes = 1 << 20
+)
+
+// Config carries everything the login flow needs. Nothing is read from the
+// environment or from package state, so a caller holds the whole contract.
+type Config struct {
+ // AuthorizeURL is the Cognito hosted-UI authorize endpoint.
+ AuthorizeURL string
+
+ // TokenURL is the Cognito token endpoint.
+ TokenURL string
+
+ // ClientID is the Cognito app client id.
+ ClientID string
+
+ // RedirectURI is the callback URL registered on the app client. Its path
+ // determines where the local callback server listens.
+ RedirectURI string
+
+ // Scopes is the space-separated scope list to request.
+ Scopes string
+
+ // AllowedDomain, when set, is the email suffix an operator must sign in
+ // with, e.g. "@wego.com".
+ AllowedDomain string
+
+ // CallbackAddr is the host:port the local callback server binds, e.g.
+ // "127.0.0.1:8110". It must agree with RedirectURI.
+ CallbackAddr string
+
+ // IdentityProvider, when set, is forwarded as identity_provider so Cognito
+ // jumps straight to that IdP instead of showing its own chooser.
+ IdentityProvider string
+
+ // OpenBrowser launches the authorize URL. Nil uses the OS default opener.
+ OpenBrowser func(string) error
+
+ // NoBrowser suppresses the browser launch. Login reports the authorize URL
+ // through PromptURL instead and then waits on the loopback listener exactly
+ // as it otherwise would, for a headless shell, a terminal on a remote host,
+ // or an operator who would rather open the URL themselves.
+ //
+ // The redirect still lands on CallbackAddr, so when the browser runs on a
+ // different machine than the CLI that port has to be reachable from it —
+ // usually `ssh -L :localhost:`. Suppressing the launch does not
+ // move where the code is delivered.
+ NoBrowser bool
+
+ // PromptURL receives the authorize URL in place of a browser launch, so the
+ // caller decides how to surface it: print it, render a QR code, hand it to
+ // another process. Required when NoBrowser is set — a sign-in whose URL the
+ // operator never sees cannot complete — and ignored otherwise.
+ PromptURL func(url string) error
+
+ // HTTPClient calls the token endpoint. Nil uses a client with a timeout.
+ HTTPClient *http.Client
+
+ // Now supplies the current time, for computing token expiry. Nil uses
+ // time.Now; tests inject a fixed clock.
+ Now func() time.Time
+}
+
+// Login runs the browser-based authorization-code-with-PKCE flow and returns
+// the operator's tokens.
+func Login(ctx context.Context, cfg Config) (*TokenSet, error) {
+ if err := cfg.validateForLogin(); err != nil {
+ return nil, err
+ }
+
+ verifier, err := generateVerifier()
+ if err != nil {
+ return nil, err
+ }
+ state, err := generateState()
+ if err != nil {
+ return nil, err
+ }
+
+ // Bind the callback port BEFORE sending the operator to Cognito. If the
+ // port is unavailable the redirect could never land, and there is no
+ // fallback port to try, so failing here saves a pointless round trip.
+ server, err := startCallbackServer(cfg.CallbackAddr, callbackPath(cfg.RedirectURI))
+ if err != nil {
+ return nil, err
+ }
+ defer server.shutdown()
+
+ if err := cfg.presentAuthorizeURL(cfg.buildAuthorizeURL(state, generateChallenge(verifier))); err != nil {
+ return nil, err
+ }
+
+ code, callbackState, err := server.wait(ctx, defaultCallbackTimeout)
+ if err != nil {
+ return nil, err
+ }
+
+ // CSRF protection, not decoration: a callback whose state is not the value
+ // we minted did not come from the authorize request we started, so the
+ // code it carries is not ours to redeem.
+ if !stateMatches(state, callbackState) {
+ return nil, errors.New("oauth state mismatch: the sign-in callback did not come from this login attempt")
+ }
+
+ tokens, err := cfg.exchangeCode(ctx, code, verifier)
+ if err != nil {
+ return nil, err
+ }
+
+ if err := cfg.checkAllowedDomain(tokens); err != nil {
+ return nil, err
+ }
+
+ return tokens, nil
+}
+
+// Refresh exchanges a refresh token for a fresh access and id token.
+func Refresh(ctx context.Context, cfg Config, refreshToken string) (*TokenSet, error) {
+ if wegostrings.IsBlank(refreshToken) {
+ return nil, errors.New("no refresh token available: sign in again")
+ }
+ if err := cfg.validateClient(); err != nil {
+ return nil, err
+ }
+
+ form := url.Values{}
+ form.Set("grant_type", "refresh_token")
+ form.Set("client_id", cfg.ClientID)
+ form.Set("refresh_token", refreshToken)
+
+ return cfg.postToken(ctx, form, refreshToken)
+}
+
+// exchangeCode redeems an authorization code together with its PKCE verifier.
+func (c Config) exchangeCode(ctx context.Context, code, verifier string) (*TokenSet, error) {
+ form := url.Values{}
+ form.Set("grant_type", "authorization_code")
+ form.Set("client_id", c.ClientID)
+ form.Set("code", code)
+ form.Set("redirect_uri", c.RedirectURI)
+ form.Set("code_verifier", verifier)
+
+ return c.postToken(ctx, form, "")
+}
+
+// tokenResponse is the wire shape of a Cognito token response. The field names
+// are fixed by the OAuth spec, so gosec's secret-field warning is expected.
+type tokenResponse struct {
+ AccessToken string `json:"access_token"` //nolint:gosec // G117: OAuth wire field; decoded in-process and never logged.
+ IDToken string `json:"id_token"`
+ RefreshToken string `json:"refresh_token"` //nolint:gosec // G117: OAuth wire field, as above.
+ ExpiresIn int `json:"expires_in"`
+}
+
+// postToken performs a token-endpoint call. fallbackRefresh is the refresh
+// token to keep if the response omits one.
+func (c Config) postToken(ctx context.Context, form url.Values, fallbackRefresh string) (*TokenSet, error) {
+ req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.TokenURL, strings.NewReader(form.Encode()))
+ if err != nil {
+ return nil, fmt.Errorf("build token request: %w", err)
+ }
+ req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+
+ //nolint:gosec // G704: TokenURL is operator-controlled CLI config (the Cognito domain), not request input.
+ resp, err := c.httpClient().Do(req)
+ if err != nil {
+ return nil, fmt.Errorf("call the token endpoint: %w", err)
+ }
+ defer func() { _ = resp.Body.Close() }()
+
+ body, err := io.ReadAll(io.LimitReader(resp.Body, maxTokenResponseBytes))
+ if err != nil {
+ return nil, fmt.Errorf("read token response: %w", err)
+ }
+
+ if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
+ return nil, tokenEndpointError(resp.StatusCode, body)
+ }
+
+ var parsed tokenResponse
+ if err := json.Unmarshal(body, &parsed); err != nil {
+ return nil, fmt.Errorf("decode token response: %w", err)
+ }
+
+ // Cognito omits refresh_token from a refresh response, so the original has
+ // to be carried forward. This is load-bearing: dropping it would silently
+ // sign the operator out on their next command. Do not add validation that
+ // rejects this fallback.
+ refresh := parsed.RefreshToken
+ if wegostrings.IsBlank(refresh) {
+ refresh = fallbackRefresh
+ }
+
+ switch {
+ case wegostrings.IsBlank(parsed.AccessToken):
+ return nil, errors.New("token response is missing access_token")
+ case wegostrings.IsBlank(parsed.IDToken):
+ return nil, errors.New("token response is missing id_token")
+ case wegostrings.IsBlank(refresh):
+ return nil, errors.New("token response is missing refresh_token")
+ case parsed.ExpiresIn <= 0:
+ // Silently accepting this is worse than failing. ExpiresAt would land on
+ // exactly now(), and IsExpired subtracts a leeway on top, so a login that
+ // just succeeded would read as already expired and the next command would
+ // refresh or send the operator back through sign-in. Refusing here names
+ // the real problem instead.
+ return nil, errors.New("token response has no usable expires_in")
+ }
+
+ return &TokenSet{
+ AccessToken: parsed.AccessToken,
+ IDToken: parsed.IDToken,
+ RefreshToken: refresh,
+ ExpiresAt: c.now().Add(time.Duration(parsed.ExpiresIn) * time.Second),
+ }, nil
+}
+
+// tokenEndpointError reports a non-2xx token response.
+//
+// Only the standard OAuth error fields are echoed. An arbitrary response body
+// is deliberately withheld, so no token can ever ride out inside an error
+// message that ends up in a log or a bug report.
+func tokenEndpointError(status int, body []byte) error {
+ var parsed struct {
+ Error string `json:"error"`
+ Description string `json:"error_description"`
+ }
+
+ if err := json.Unmarshal(body, &parsed); err == nil && wegostrings.IsNotBlank(parsed.Error) {
+ if wegostrings.IsNotBlank(parsed.Description) {
+ return fmt.Errorf("token endpoint returned %d: %s: %s", status, parsed.Error, parsed.Description)
+ }
+ return fmt.Errorf("token endpoint returned %d: %s", status, parsed.Error)
+ }
+
+ return fmt.Errorf("token endpoint returned %d (response body withheld: it may carry credentials)", status)
+}
+
+// buildAuthorizeURL assembles the hosted-UI URL the operator is sent to.
+func (c Config) buildAuthorizeURL(state, challenge string) string {
+ query := url.Values{}
+ query.Set("response_type", "code")
+ query.Set("client_id", c.ClientID)
+ query.Set("redirect_uri", c.RedirectURI)
+ query.Set("scope", c.Scopes)
+ query.Set("state", state)
+ query.Set("code_challenge", challenge)
+ query.Set("code_challenge_method", "S256")
+ if wegostrings.IsNotBlank(c.IdentityProvider) {
+ query.Set("identity_provider", c.IdentityProvider)
+ }
+
+ separator := "?"
+ if strings.Contains(c.AuthorizeURL, "?") {
+ separator = "&"
+ }
+ return c.AuthorizeURL + separator + query.Encode()
+}
+
+// checkAllowedDomain rejects an operator signed in outside AllowedDomain.
+//
+// This is client-side UX, NOT a security boundary: it catches "you signed in
+// with your personal Google account" before the CLI starts issuing calls that
+// would fail confusingly. The server does not enforce it, so nothing may rely
+// on it for authorization.
+func (c Config) checkAllowedDomain(tokens *TokenSet) error {
+ if wegostrings.IsBlank(c.AllowedDomain) {
+ return nil
+ }
+
+ email, err := tokens.Email()
+ if err != nil {
+ return fmt.Errorf("check the signed-in email: %w", err)
+ }
+
+ if !strings.HasSuffix(strings.ToLower(email), strings.ToLower(c.AllowedDomain)) {
+ return fmt.Errorf("signed in as %s, but a %s account is required", email, c.AllowedDomain)
+ }
+ return nil
+}
+
+// validateClient checks the values every token call needs.
+func (c Config) validateClient() error {
+ if wegostrings.IsBlank(c.ClientID) || strings.HasPrefix(c.ClientID, placeholderPrefix) {
+ return errors.New("cognito app client is not provisioned yet: Config.ClientID is blank or still a placeholder, so the app client needs to be created and its id wired into the caller's config")
+ }
+ if wegostrings.IsBlank(c.TokenURL) {
+ return errors.New("cognito token url is not configured")
+ }
+ return nil
+}
+
+// validateForLogin checks everything the interactive flow additionally needs.
+func (c Config) validateForLogin() error {
+ if err := c.validateClient(); err != nil {
+ return err
+ }
+ if wegostrings.IsBlank(c.AuthorizeURL) {
+ return errors.New("cognito authorize url is not configured")
+ }
+ if wegostrings.IsBlank(c.RedirectURI) {
+ return errors.New("cognito redirect uri is not configured")
+ }
+ if wegostrings.IsBlank(c.CallbackAddr) {
+ return errors.New("local callback address is not configured")
+ }
+ if c.NoBrowser && c.PromptURL == nil {
+ return errors.New("no-browser sign-in needs Config.PromptURL: with no browser launched and no way to report the url, the operator has nothing to open")
+ }
+ return nil
+}
+
+// httpClient is the client to call the token endpoint with. It never follows
+// redirects, whoever supplied it.
+//
+// The returned client is a COPY, so a caller-supplied HTTPClient is neither
+// mutated nor able to reinstate redirect-following. That override is
+// deliberate: the token request carries the authorization code, the PKCE
+// verifier and the client id on sign-in and the refresh token on renewal, and
+// Go's default policy follows up to ten redirects, re-sending the body
+// verbatim on a 307 or 308. One redirect would therefore hand a complete
+// credential set to whatever host the response named. It is also an injection
+// route inwards, since the body that came back would be parsed as the session
+// to use.
+//
+// A redirect from the token endpoint has no legitimate meaning here: TokenURL
+// is a Cognito domain the operator configured, and Cognito answers it
+// directly. Refusing turns the redirect into the error it should be.
+func (c Config) httpClient() *http.Client {
+ base := c.HTTPClient
+ if base == nil {
+ base = &http.Client{Timeout: defaultHTTPTimeout}
+ }
+
+ refusing := *base
+ refusing.CheckRedirect = refuseRedirect
+
+ return &refusing
+}
+
+// refuseRedirect stops the client at the redirect response instead of
+// following it. Returning ErrUseLastResponse rather than an error of our own
+// hands postToken the 3xx itself, which its status check then reports with the
+// endpoint and status an operator needs to debug the misconfiguration.
+func refuseRedirect(_ *http.Request, _ []*http.Request) error {
+ return http.ErrUseLastResponse
+}
+
+// now is the current time, from the injected clock when there is one.
+func (c Config) now() time.Time {
+ if c.Now != nil {
+ return c.Now()
+ }
+ return time.Now()
+}
+
+// openBrowserAt sends the operator to rawURL.
+// presentAuthorizeURL gets the operator to the authorize URL, by launching a
+// browser or, under NoBrowser, by handing the URL to PromptURL.
+func (c Config) presentAuthorizeURL(rawURL string) error {
+ if c.NoBrowser {
+ if err := c.PromptURL(rawURL); err != nil {
+ return fmt.Errorf("present the sign-in url: %w", err)
+ }
+
+ return nil
+ }
+
+ if err := c.openBrowserAt(rawURL); err != nil {
+ return fmt.Errorf("open browser for sign-in: %w", err)
+ }
+
+ return nil
+}
+
+func (c Config) openBrowserAt(rawURL string) error {
+ if c.OpenBrowser != nil {
+ return c.OpenBrowser(rawURL)
+ }
+ return openBrowser(rawURL)
+}
diff --git a/cognito/oauth_internal_test.go b/cognito/oauth_internal_test.go
new file mode 100644
index 0000000..6b84acf
--- /dev/null
+++ b/cognito/oauth_internal_test.go
@@ -0,0 +1,220 @@
+package cognito
+
+import (
+ "context"
+ "encoding/json"
+ "net/http"
+ "net/http/httptest"
+ "net/url"
+ "sync/atomic"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestTokenEndpointError(t *testing.T) {
+ tests := []struct {
+ name string
+ givenStatus int
+ givenBody string
+ wantContains []string
+ wantNotContains []string
+ }{
+ {
+ name: "non-json body is withheld entirely",
+ givenStatus: http.StatusInternalServerError,
+ givenBody: "access_token=super-secret-value",
+ // An unparseable body could be anything, so we report only the
+ // status: a token must never ride out inside an error string.
+ wantContains: []string{"500", "withheld"},
+ wantNotContains: []string{"super-secret-value"},
+ },
+ {
+ name: "json without an error field is withheld",
+ givenStatus: http.StatusBadGateway,
+ givenBody: `{"id_token":"secret-jwt-value"}`,
+ wantContains: []string{"502", "withheld"},
+ wantNotContains: []string{"secret-jwt-value"},
+ },
+ {
+ name: "oauth error is echoed",
+ givenStatus: http.StatusBadRequest,
+ givenBody: `{"error":"invalid_grant"}`,
+ wantContains: []string{"400", "invalid_grant"},
+ },
+ {
+ name: "oauth error and description are echoed",
+ givenStatus: http.StatusBadRequest,
+ givenBody: `{"error":"invalid_grant","error_description":"code expired"}`,
+ wantContains: []string{"400", "invalid_grant", "code expired"},
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ err := tokenEndpointError(tt.givenStatus, []byte(tt.givenBody))
+
+ require.Error(t, err)
+ for _, want := range tt.wantContains {
+ assert.Contains(t, err.Error(), want)
+ }
+ for _, notWant := range tt.wantNotContains {
+ assert.NotContains(t, err.Error(), notWant)
+ }
+ })
+ }
+}
+
+func TestConfig_BuildAuthorizeURL(t *testing.T) {
+ tests := []struct {
+ name string
+ givenAuthorizeURL string
+ givenIdentityProvider string
+ wantQueryHasProvider bool
+ }{
+ {
+ name: "an existing query string is appended to",
+ givenAuthorizeURL: "https://cognito.test/oauth2/authorize?foo=bar",
+ },
+ {
+ name: "no identity provider leaves the parameter out",
+ givenAuthorizeURL: "https://cognito.test/oauth2/authorize",
+ },
+ {
+ name: "identity provider is forwarded when set",
+ givenAuthorizeURL: "https://cognito.test/oauth2/authorize",
+ givenIdentityProvider: "Google",
+ wantQueryHasProvider: true,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ cfg := Config{
+ AuthorizeURL: tt.givenAuthorizeURL,
+ ClientID: "client",
+ RedirectURI: "http://127.0.0.1:8110/callback",
+ Scopes: "openid email",
+ IdentityProvider: tt.givenIdentityProvider,
+ }
+
+ got, err := url.Parse(cfg.buildAuthorizeURL("the-state", "the-challenge"))
+ require.NoError(t, err)
+
+ query := got.Query()
+ assert.Equal(t, "the-state", query.Get("state"))
+ assert.Equal(t, "the-challenge", query.Get("code_challenge"))
+ assert.Equal(t, "S256", query.Get("code_challenge_method"))
+ assert.Equal(t, tt.givenIdentityProvider, query.Get("identity_provider"))
+ assert.Equal(t, tt.wantQueryHasProvider, query.Has("identity_provider"))
+
+ if tt.givenAuthorizeURL == "https://cognito.test/oauth2/authorize?foo=bar" {
+ assert.Equal(t, "bar", query.Get("foo"), "an existing query parameter must survive")
+ }
+ })
+ }
+}
+
+func TestConfig_Defaults(t *testing.T) {
+ var zero Config
+
+ assert.Equal(t, defaultHTTPTimeout, zero.httpClient().Timeout, "a nil HTTPClient must get a timeout")
+ assert.WithinDuration(t, time.Now(), zero.now(), time.Minute, "a nil Now must fall back to the system clock")
+
+ // An injected client is honoured for its transport and timeout but is
+ // COPIED, never handed back: httpClient has to own CheckRedirect, and it
+ // must not mutate a client the caller may use elsewhere. See
+ // TestPostToken_RefusesRedirects for why the override exists.
+ transport := &http.Transport{}
+ custom := &http.Client{Timeout: time.Second, Transport: transport}
+ got := Config{HTTPClient: custom}.httpClient()
+
+ assert.NotSame(t, custom, got, "the caller's client must not be handed back")
+ assert.Equal(t, time.Second, got.Timeout, "the caller's timeout must survive")
+ assert.Same(t, transport, got.Transport, "the caller's transport must survive")
+ assert.NotNil(t, got.CheckRedirect, "the returned client must refuse redirects")
+ assert.Nil(t, custom.CheckRedirect, "the caller's client must not be mutated")
+
+ fixed := time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC)
+ assert.Equal(t, fixed, Config{Now: func() time.Time { return fixed }}.now())
+}
+
+func TestConfig_OpenBrowserAtUsesTheInjectedOpener(t *testing.T) {
+ var got string
+ cfg := Config{OpenBrowser: func(rawURL string) error {
+ got = rawURL
+ return nil
+ }}
+
+ require.NoError(t, cfg.openBrowserAt("https://example.test/authorize"))
+ assert.Equal(t, "https://example.test/authorize", got)
+}
+
+// TestPostToken_RefusesRedirects proves a token request is never re-sent to a
+// host other than the one configured.
+//
+// The form carries the authorization code, the PKCE verifier and the client id
+// on sign-in, and the refresh token on renewal. Go's default client follows up
+// to ten redirects and re-sends the body verbatim on a 307 or 308, so
+// following one would hand a complete credential set to whatever host the
+// response named. A redirect is also an injection route in the other
+// direction: the body that comes back would be parsed as a token set, letting
+// an unintended host choose the session the CLI then uses.
+func TestPostToken_RefusesRedirects(t *testing.T) {
+ tests := []struct {
+ name string
+ givenStatus int
+ givenClient *http.Client
+ }{
+ {
+ name: "the default client refuses a 307, which would re-send the body",
+ givenStatus: http.StatusTemporaryRedirect,
+ },
+ {
+ name: "the default client refuses a 302",
+ givenStatus: http.StatusFound,
+ },
+ {
+ // A caller supplying its own client must not be able to reinstate
+ // following, deliberately or by copying a client from elsewhere.
+ name: "an injected permissive client cannot opt back into following",
+ givenStatus: http.StatusTemporaryRedirect,
+ givenClient: &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return nil }},
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ var elsewhereHits atomic.Int32
+
+ // Answers with a usable token set, so following the redirect would
+ // look like a successful sign-in rather than an error.
+ elsewhere := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+ elsewhereHits.Add(1)
+ w.Header().Set("Content-Type", "application/json")
+ _ = json.NewEncoder(w).Encode(map[string]any{
+ "access_token": "attacker-access",
+ "id_token": "attacker-id",
+ "expires_in": 3600,
+ })
+ }))
+ t.Cleanup(elsewhere.Close)
+
+ origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ http.Redirect(w, r, elsewhere.URL, tt.givenStatus)
+ }))
+ t.Cleanup(origin.Close)
+
+ cfg := Config{TokenURL: origin.URL, HTTPClient: tt.givenClient}
+ form := url.Values{"code": {"secret-code"}, "code_verifier": {"secret-verifier"}}
+
+ tokens, err := cfg.postToken(context.Background(), form, "")
+
+ require.Error(t, err, "a redirected token endpoint must not produce a session")
+ assert.Nil(t, tokens)
+ assert.Zero(t, elsewhereHits.Load(), "credentials must never reach the redirect target")
+ })
+ }
+}
diff --git a/cognito/oauth_test.go b/cognito/oauth_test.go
new file mode 100644
index 0000000..b7c70bc
--- /dev/null
+++ b/cognito/oauth_test.go
@@ -0,0 +1,765 @@
+package cognito_test
+
+import (
+ "context"
+ "crypto/sha256"
+ "encoding/base64"
+ "encoding/json"
+ "errors"
+ "net"
+ "net/http"
+ "net/http/httptest"
+ "net/url"
+ "sync"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+ wegostrings "github.com/wego/pkg/strings"
+
+ "github.com/wego/pkg/cognito"
+)
+
+const (
+ testAuthCode = "the-authorization-code"
+ testAccessValue = "access-value"
+ testIDTokenLabel = "id_token"
+ testRefreshValue = "refresh-value"
+)
+
+func TestLogin(t *testing.T) {
+ tests := []struct {
+ name string
+ givenConfig func(*cognito.Config)
+ givenBrowser func(*fakeBrowser)
+ givenHandler func(*testing.T) http.HandlerFunc
+ givenCancel bool
+ wantEmail string
+ wantErrContains string
+ }{
+ {
+ name: "blank client id tells the operator it is unprovisioned",
+ givenConfig: func(c *cognito.Config) { c.ClientID = "" },
+ wantErrContains: "not provisioned",
+ },
+ {
+ name: "placeholder client id tells the operator it is unprovisioned",
+ givenConfig: func(c *cognito.Config) { c.ClientID = "REPLACE_WITH_COGNITO_CLIENT_ID" },
+ wantErrContains: "not provisioned",
+ },
+ {
+ name: "missing authorize url is rejected",
+ givenConfig: func(c *cognito.Config) { c.AuthorizeURL = "" },
+ wantErrContains: "authorize url",
+ },
+ {
+ name: "missing token url is rejected",
+ givenConfig: func(c *cognito.Config) { c.TokenURL = "" },
+ wantErrContains: "token url",
+ },
+ {
+ name: "missing redirect uri is rejected",
+ givenConfig: func(c *cognito.Config) { c.RedirectURI = "" },
+ wantErrContains: "redirect uri",
+ },
+ {
+ name: "missing callback address is rejected",
+ givenConfig: func(c *cognito.Config) { c.CallbackAddr = "" },
+ wantErrContains: "callback address",
+ },
+ {
+ name: "browser launch failure is surfaced",
+ givenBrowser: func(f *fakeBrowser) { f.openErr = errors.New("no browser here") },
+ wantErrContains: "open browser",
+ },
+ {
+ name: "tampered state is rejected",
+ givenBrowser: func(f *fakeBrowser) {
+ f.tamper = func(q url.Values) { q.Set("state", "tampered-state") }
+ },
+ wantErrContains: "state mismatch",
+ },
+ {
+ name: "missing state is rejected",
+ givenBrowser: func(f *fakeBrowser) {
+ f.tamper = func(q url.Values) { q.Del("state") }
+ },
+ wantErrContains: "state mismatch",
+ },
+ {
+ name: "authorize server error is surfaced",
+ givenBrowser: func(f *fakeBrowser) {
+ f.tamper = func(q url.Values) {
+ q.Del("code")
+ q.Set("error", "access_denied")
+ q.Set("error_description", "operator declined")
+ }
+ },
+ wantErrContains: "access_denied",
+ },
+ {
+ name: "missing code is rejected",
+ givenBrowser: func(f *fakeBrowser) {
+ f.tamper = func(q url.Values) { q.Del("code") }
+ },
+ wantErrContains: "no authorization code",
+ },
+ {
+ name: "token endpoint rejection is surfaced",
+ givenHandler: func(t *testing.T) http.HandlerFunc {
+ t.Helper()
+ return jsonHandler(http.StatusBadRequest, map[string]any{
+ "error": "invalid_grant",
+ "error_description": "code already used",
+ })
+ },
+ wantErrContains: "invalid_grant",
+ },
+ {
+ name: "unparseable token response is rejected",
+ givenHandler: func(t *testing.T) http.HandlerFunc {
+ t.Helper()
+ return rawHandler(http.StatusOK, "not json")
+ },
+ wantErrContains: "decode token response",
+ },
+ {
+ name: "token response without an access token is rejected",
+ givenHandler: func(t *testing.T) http.HandlerFunc {
+ t.Helper()
+ body := goodTokenBody(t, testOperatorEmail)
+ delete(body, "access_token")
+ return jsonHandler(http.StatusOK, body)
+ },
+ wantErrContains: "missing access_token",
+ },
+ {
+ name: "token response without an id token is rejected",
+ givenHandler: func(t *testing.T) http.HandlerFunc {
+ t.Helper()
+ body := goodTokenBody(t, testOperatorEmail)
+ delete(body, testIDTokenLabel)
+ return jsonHandler(http.StatusOK, body)
+ },
+ wantErrContains: "missing id_token",
+ },
+ {
+ name: "token response without a refresh token is rejected",
+ givenHandler: func(t *testing.T) http.HandlerFunc {
+ t.Helper()
+ body := goodTokenBody(t, testOperatorEmail)
+ delete(body, "refresh_token")
+ return jsonHandler(http.StatusOK, body)
+ },
+ wantErrContains: "missing refresh_token",
+ },
+ {
+ name: "id token without an email claim is rejected",
+ givenHandler: func(t *testing.T) http.HandlerFunc {
+ t.Helper()
+ body := goodTokenBody(t, testOperatorEmail)
+ body[testIDTokenLabel] = mustIDToken(t, map[string]any{"sub": "abc"})
+ return jsonHandler(http.StatusOK, body)
+ },
+ wantErrContains: "email",
+ },
+ {
+ name: "email outside the allowed domain is rejected",
+ givenHandler: func(t *testing.T) http.HandlerFunc {
+ t.Helper()
+ return jsonHandler(http.StatusOK, goodTokenBody(t, "someone@gmail.com"))
+ },
+ wantErrContains: "@wego.com",
+ },
+ {
+ name: "cancelled context stops waiting",
+ givenBrowser: func(f *fakeBrowser) { f.suppress = true },
+ givenCancel: true,
+ wantErrContains: "context canceled",
+ },
+ {
+ name: "successful login returns a token set",
+ wantEmail: testOperatorEmail,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ var handler http.HandlerFunc
+ if tt.givenHandler != nil {
+ handler = tt.givenHandler(t)
+ }
+ ts := newTokenServer(t, handler)
+
+ browser := newFakeBrowser()
+ if tt.givenBrowser != nil {
+ tt.givenBrowser(browser)
+ }
+
+ cfg := baseConfig(t, ts.URL, mustFreeAddr(t))
+ cfg.OpenBrowser = browser.open
+ if tt.givenConfig != nil {
+ tt.givenConfig(&cfg)
+ }
+
+ ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+ defer cancel()
+ if tt.givenCancel {
+ cancel()
+ }
+
+ got, err := cognito.Login(ctx, cfg)
+
+ if wegostrings.IsNotEmpty(tt.wantErrContains) {
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), tt.wantErrContains)
+ assert.Nil(t, got)
+ return
+ }
+ require.NoError(t, err)
+ require.NotNil(t, got)
+
+ email, err := got.Email()
+ require.NoError(t, err)
+ assert.Equal(t, tt.wantEmail, email)
+ })
+ }
+}
+
+func TestLogin_BindsTheCallbackPortBeforeOpeningTheBrowser(t *testing.T) {
+ ln, err := net.Listen("tcp", "127.0.0.1:0")
+ require.NoError(t, err)
+ t.Cleanup(func() { _ = ln.Close() })
+ addr := ln.Addr().String()
+
+ ts := newTokenServer(t, nil)
+ browser := newFakeBrowser()
+ cfg := baseConfig(t, ts.URL, addr)
+ cfg.OpenBrowser = browser.open
+
+ got, err := cognito.Login(context.Background(), cfg)
+
+ require.Error(t, err, "an occupied callback port must fail fast, not fall back to another port")
+ assert.Nil(t, got)
+ assert.Contains(t, err.Error(), addr, "the error must name the address the operator has to free")
+ assert.Empty(t, browser.capturedURL(), "the browser must not be launched when the callback port is unavailable")
+}
+
+func TestLogin_SendsPKCEAndState(t *testing.T) {
+ ts := newTokenServer(t, nil)
+ browser := newFakeBrowser()
+ cfg := baseConfig(t, ts.URL, mustFreeAddr(t))
+ cfg.OpenBrowser = browser.open
+
+ got, err := cognito.Login(context.Background(), cfg)
+ require.NoError(t, err)
+
+ authorizeURL, err := url.Parse(browser.capturedURL())
+ require.NoError(t, err)
+ aq := authorizeURL.Query()
+
+ assert.Equal(t, "https://cognito.test/oauth2/authorize", authorizeURL.Scheme+"://"+authorizeURL.Host+authorizeURL.Path)
+ assert.Equal(t, "code", aq.Get("response_type"))
+ assert.Equal(t, cfg.ClientID, aq.Get("client_id"))
+ assert.Equal(t, cfg.RedirectURI, aq.Get("redirect_uri"))
+ assert.Equal(t, cfg.Scopes, aq.Get("scope"))
+ assert.Equal(t, "S256", aq.Get("code_challenge_method"), "plain PKCE must never be used")
+ assert.NotEmpty(t, aq.Get("state"), "state is CSRF protection and must always be sent")
+
+ form := ts.lastForm(t)
+ assert.Equal(t, "authorization_code", form.Get("grant_type"))
+ assert.Equal(t, cfg.ClientID, form.Get("client_id"))
+ assert.Equal(t, testAuthCode, form.Get("code"))
+ assert.Equal(t, cfg.RedirectURI, form.Get("redirect_uri"))
+
+ verifier := form.Get("code_verifier")
+ require.NotEmpty(t, verifier, "the exchange must carry the PKCE verifier")
+ assert.GreaterOrEqual(t, len(verifier), 43)
+ assert.LessOrEqual(t, len(verifier), 128)
+
+ sum := sha256.Sum256([]byte(verifier))
+ assert.Equal(t, base64.RawURLEncoding.EncodeToString(sum[:]), aq.Get("code_challenge"),
+ "code_challenge must be the S256 hash of the verifier actually redeemed")
+
+ assert.Equal(t, fixedNow.Add(time.Hour), got.ExpiresAt, "ExpiresAt must be derived from the injected clock")
+ assert.Equal(t, testAccessValue, got.AccessToken)
+ assert.Equal(t, testRefreshValue, got.RefreshToken)
+}
+
+func TestLogin_ErrorsDoNotLeakVerifierOrState(t *testing.T) {
+ ts := newTokenServer(t, jsonHandler(http.StatusBadRequest, map[string]any{"error": "invalid_grant"}))
+ browser := newFakeBrowser()
+ cfg := baseConfig(t, ts.URL, mustFreeAddr(t))
+ cfg.OpenBrowser = browser.open
+
+ _, err := cognito.Login(context.Background(), cfg)
+ require.Error(t, err)
+
+ authorizeURL, parseErr := url.Parse(browser.capturedURL())
+ require.NoError(t, parseErr)
+ state := authorizeURL.Query().Get("state")
+ require.NotEmpty(t, state)
+ verifier := ts.lastForm(t).Get("code_verifier")
+ require.NotEmpty(t, verifier)
+
+ assert.NotContains(t, err.Error(), state, "state must never appear in an error message")
+ assert.NotContains(t, err.Error(), verifier, "the PKCE verifier must never appear in an error message")
+}
+
+func TestLogin_FallsBackToTheSystemClockWhenNowIsNil(t *testing.T) {
+ ts := newTokenServer(t, nil)
+ browser := newFakeBrowser()
+ cfg := baseConfig(t, ts.URL, mustFreeAddr(t))
+ cfg.OpenBrowser = browser.open
+ cfg.Now = nil
+
+ before := time.Now()
+ got, err := cognito.Login(context.Background(), cfg)
+ require.NoError(t, err)
+
+ assert.WithinDuration(t, before.Add(time.Hour), got.ExpiresAt, time.Minute)
+}
+
+func TestLogin_AllowsAnyEmailWhenNoDomainIsConfigured(t *testing.T) {
+ ts := newTokenServer(t, jsonHandler(http.StatusOK, goodTokenBody(t, "contractor@example.test")))
+ browser := newFakeBrowser()
+ cfg := baseConfig(t, ts.URL, mustFreeAddr(t))
+ cfg.OpenBrowser = browser.open
+ cfg.AllowedDomain = ""
+
+ got, err := cognito.Login(context.Background(), cfg)
+ require.NoError(t, err)
+ require.NotNil(t, got)
+}
+
+func TestRefresh(t *testing.T) {
+ tests := []struct {
+ name string
+ givenConfig func(*cognito.Config)
+ givenRefresh string
+ givenHandler func(*testing.T) http.HandlerFunc
+ wantRefresh string
+ wantErrContains string
+ }{
+ {
+ name: "blank refresh token is rejected",
+ givenRefresh: "",
+ wantErrContains: "refresh token",
+ },
+ {
+ name: "blank client id tells the operator it is unprovisioned",
+ givenRefresh: testRefreshValue,
+ givenConfig: func(c *cognito.Config) { c.ClientID = "" },
+ wantErrContains: "not provisioned",
+ },
+ {
+ name: "placeholder client id tells the operator it is unprovisioned",
+ givenRefresh: testRefreshValue,
+ givenConfig: func(c *cognito.Config) { c.ClientID = "REPLACE_WITH_COGNITO_CLIENT_ID" },
+ wantErrContains: "not provisioned",
+ },
+ {
+ name: "missing token url is rejected",
+ givenRefresh: testRefreshValue,
+ givenConfig: func(c *cognito.Config) { c.TokenURL = "" },
+ wantErrContains: "token url",
+ },
+ {
+ name: "server failure is surfaced",
+ givenRefresh: testRefreshValue,
+ givenHandler: func(t *testing.T) http.HandlerFunc {
+ t.Helper()
+ return rawHandler(http.StatusInternalServerError, "upstream exploded")
+ },
+ wantErrContains: "500",
+ },
+ {
+ name: "unparseable response is rejected",
+ givenRefresh: testRefreshValue,
+ givenHandler: func(t *testing.T) http.HandlerFunc {
+ t.Helper()
+ return rawHandler(http.StatusOK, "not json")
+ },
+ wantErrContains: "decode token response",
+ },
+ {
+ name: "response without an access token is rejected",
+ givenRefresh: testRefreshValue,
+ givenHandler: func(t *testing.T) http.HandlerFunc {
+ t.Helper()
+ body := goodTokenBody(t, testOperatorEmail)
+ delete(body, "access_token")
+ return jsonHandler(http.StatusOK, body)
+ },
+ wantErrContains: "missing access_token",
+ },
+ {
+ name: "a rotated refresh token is adopted",
+ givenRefresh: "original-refresh",
+ givenHandler: func(t *testing.T) http.HandlerFunc {
+ t.Helper()
+ body := goodTokenBody(t, testOperatorEmail)
+ body["refresh_token"] = "rotated-refresh"
+ return jsonHandler(http.StatusOK, body)
+ },
+ wantRefresh: "rotated-refresh",
+ },
+ {
+ name: "an omitted refresh token falls back to the original",
+ givenRefresh: "original-refresh",
+ givenHandler: func(t *testing.T) http.HandlerFunc {
+ t.Helper()
+ body := goodTokenBody(t, testOperatorEmail)
+ delete(body, "refresh_token")
+ return jsonHandler(http.StatusOK, body)
+ },
+ wantRefresh: "original-refresh",
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ var handler http.HandlerFunc
+ if tt.givenHandler != nil {
+ handler = tt.givenHandler(t)
+ }
+ ts := newTokenServer(t, handler)
+
+ cfg := baseConfig(t, ts.URL, "127.0.0.1:0")
+ if tt.givenConfig != nil {
+ tt.givenConfig(&cfg)
+ }
+
+ got, err := cognito.Refresh(context.Background(), cfg, tt.givenRefresh)
+
+ if wegostrings.IsNotEmpty(tt.wantErrContains) {
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), tt.wantErrContains)
+ assert.Nil(t, got)
+ return
+ }
+ require.NoError(t, err)
+ require.NotNil(t, got)
+ assert.Equal(t, tt.wantRefresh, got.RefreshToken)
+ assert.Equal(t, "refresh_token", ts.lastForm(t).Get("grant_type"))
+ })
+ }
+}
+
+// TestRefresh_PreservesTheOriginalRefreshToken pins the load-bearing Cognito
+// invariant: a refresh response omits refresh_token, and dropping the original
+// would log the operator out on the very next command.
+func TestRefresh_PreservesTheOriginalRefreshToken(t *testing.T) {
+ body := goodTokenBody(t, testOperatorEmail)
+ delete(body, "refresh_token")
+ ts := newTokenServer(t, jsonHandler(http.StatusOK, body))
+
+ cfg := baseConfig(t, ts.URL, "127.0.0.1:0")
+
+ got, err := cognito.Refresh(context.Background(), cfg, "the-long-lived-refresh-token")
+ require.NoError(t, err)
+ assert.Equal(t, "the-long-lived-refresh-token", got.RefreshToken)
+ assert.Equal(t, "the-long-lived-refresh-token", ts.lastForm(t).Get("refresh_token"),
+ "the original refresh token must be what we present to Cognito")
+ assert.Equal(t, fixedNow.Add(time.Hour), got.ExpiresAt)
+}
+
+// TestRefresh_DoesNotApplyTheDomainGate documents that the allowed-domain check
+// is a login-time UX affordance, not something to re-run on every refresh.
+func TestRefresh_DoesNotApplyTheDomainGate(t *testing.T) {
+ ts := newTokenServer(t, jsonHandler(http.StatusOK, goodTokenBody(t, "someone@gmail.com")))
+ cfg := baseConfig(t, ts.URL, "127.0.0.1:0")
+
+ got, err := cognito.Refresh(context.Background(), cfg, testRefreshValue)
+ require.NoError(t, err)
+ require.NotNil(t, got)
+}
+
+func TestRefresh_HonoursContextCancellation(t *testing.T) {
+ ts := newTokenServer(t, nil)
+ cfg := baseConfig(t, ts.URL, "127.0.0.1:0")
+
+ ctx, cancel := context.WithCancel(context.Background())
+ cancel()
+
+ got, err := cognito.Refresh(ctx, cfg, testRefreshValue)
+ require.Error(t, err)
+ assert.Nil(t, got)
+}
+
+// baseConfig is a working staging-shaped config aimed at a test token server.
+func baseConfig(t *testing.T, tokenURL, callbackAddr string) cognito.Config {
+ t.Helper()
+ return cognito.Config{
+ AuthorizeURL: "https://cognito.test/oauth2/authorize",
+ TokenURL: tokenURL,
+ ClientID: "test-client-id",
+ RedirectURI: "http://" + callbackAddr + "/callback",
+ Scopes: "openid email profile",
+ AllowedDomain: "@wego.com",
+ CallbackAddr: callbackAddr,
+ Now: func() time.Time { return fixedNow },
+ }
+}
+
+// goodTokenBody is a well-formed Cognito token response.
+func goodTokenBody(t *testing.T, email string) map[string]any {
+ t.Helper()
+ return map[string]any{
+ "access_token": testAccessValue,
+ testIDTokenLabel: mustIDToken(t, map[string]any{"email": email}),
+ "refresh_token": testRefreshValue,
+ "expires_in": 3600,
+ "token_type": "Bearer",
+ }
+}
+
+func jsonHandler(status int, body any) http.HandlerFunc {
+ return func(w http.ResponseWriter, _ *http.Request) {
+ enc, err := json.Marshal(body)
+ if err != nil {
+ http.Error(w, "marshal failed", http.StatusInternalServerError)
+ return
+ }
+ w.Header().Set("Content-Type", "application/json")
+ w.WriteHeader(status)
+ _, _ = w.Write(enc)
+ }
+}
+
+func rawHandler(status int, body string) http.HandlerFunc {
+ return func(w http.ResponseWriter, _ *http.Request) {
+ w.WriteHeader(status)
+ _, _ = w.Write([]byte(body))
+ }
+}
+
+// tokenServer records every form it is posted so tests can assert on the
+// exchange itself, not just its result.
+type tokenServer struct {
+ *httptest.Server
+
+ mu sync.Mutex
+ forms []url.Values
+}
+
+func newTokenServer(t *testing.T, handler http.HandlerFunc) *tokenServer {
+ t.Helper()
+
+ ts := &tokenServer{}
+ if handler == nil {
+ handler = jsonHandler(http.StatusOK, goodTokenBody(t, testOperatorEmail))
+ }
+
+ ts.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if err := r.ParseForm(); err != nil {
+ http.Error(w, "bad form", http.StatusBadRequest)
+ return
+ }
+ ts.mu.Lock()
+ ts.forms = append(ts.forms, r.PostForm)
+ ts.mu.Unlock()
+ handler(w, r)
+ }))
+ t.Cleanup(ts.Close)
+
+ return ts
+}
+
+func (ts *tokenServer) lastForm(t *testing.T) url.Values {
+ t.Helper()
+ ts.mu.Lock()
+ defer ts.mu.Unlock()
+ require.NotEmpty(t, ts.forms, "the token endpoint was never called")
+ return ts.forms[len(ts.forms)-1]
+}
+
+// fakeBrowser stands in for the operator's browser: it reads the authorize URL
+// and drives the loopback callback the way a real redirect would.
+type fakeBrowser struct {
+ openErr error
+ suppress bool
+ tamper func(url.Values)
+
+ mu sync.Mutex
+ authorizeURL string
+}
+
+func newFakeBrowser() *fakeBrowser {
+ return &fakeBrowser{}
+}
+
+func (f *fakeBrowser) open(rawURL string) error {
+ f.mu.Lock()
+ f.authorizeURL = rawURL
+ f.mu.Unlock()
+
+ if f.openErr != nil {
+ return f.openErr
+ }
+ if f.suppress {
+ return nil
+ }
+
+ authorizeURL, err := url.Parse(rawURL)
+ if err != nil {
+ return err
+ }
+ callback, err := url.Parse(authorizeURL.Query().Get("redirect_uri"))
+ if err != nil {
+ return err
+ }
+
+ q := url.Values{}
+ q.Set("code", testAuthCode)
+ q.Set("state", authorizeURL.Query().Get("state"))
+ if f.tamper != nil {
+ f.tamper(q)
+ }
+ callback.RawQuery = q.Encode()
+
+ go func() {
+ req, reqErr := http.NewRequestWithContext(context.Background(), http.MethodGet, callback.String(), nil)
+ if reqErr != nil {
+ return
+ }
+ //nolint:gosec // G704: a loopback callback URL derived from the authorize URL this test built.
+ resp, doErr := http.DefaultClient.Do(req)
+ if doErr != nil {
+ return
+ }
+ _ = resp.Body.Close()
+ }()
+
+ return nil
+}
+
+func (f *fakeBrowser) capturedURL() string {
+ f.mu.Lock()
+ defer f.mu.Unlock()
+ return f.authorizeURL
+}
+
+// mustFreeAddr reserves and releases a loopback port, returning its address.
+func mustFreeAddr(t *testing.T) string {
+ t.Helper()
+ ln, err := net.Listen("tcp", "127.0.0.1:0")
+ require.NoError(t, err)
+ addr := ln.Addr().String()
+ require.NoError(t, ln.Close())
+ return addr
+}
+
+// TestLogin_NoBrowser covers the headless path: the caller suppresses the
+// launch and surfaces the URL itself, and the sign-in still completes on the
+// loopback listener. Suppressing the launch must not change anything else —
+// the same PKCE and state parameters have to be sent, because a URL an
+// operator pastes by hand is the same URL a browser would have been given.
+func TestLogin_NoBrowser(t *testing.T) {
+ ts := newTokenServer(t, nil)
+
+ // The fake browser doubles as the prompt: it records the URL and drives the
+ // callback, which is exactly what an operator pasting the URL would cause.
+ prompt := newFakeBrowser()
+
+ cfg := baseConfig(t, ts.URL, mustFreeAddr(t))
+ cfg.NoBrowser = true
+ cfg.PromptURL = prompt.open
+ cfg.OpenBrowser = func(string) error {
+ t.Error("NoBrowser must not launch a browser")
+
+ return nil
+ }
+
+ got, err := cognito.Login(context.Background(), cfg)
+ require.NoError(t, err)
+ require.Equal(t, testAccessValue, got.AccessToken)
+
+ authorizeURL, err := url.Parse(prompt.capturedURL())
+ require.NoError(t, err, "PromptURL must receive a usable authorize url")
+
+ aq := authorizeURL.Query()
+ assert.Equal(t, cfg.ClientID, aq.Get("client_id"))
+ assert.Equal(t, cfg.RedirectURI, aq.Get("redirect_uri"))
+ assert.Equal(t, "S256", aq.Get("code_challenge_method"), "plain PKCE must never be used")
+ assert.NotEmpty(t, aq.Get("code_challenge"))
+ assert.NotEmpty(t, aq.Get("state"), "state is CSRF protection and must always be sent")
+}
+
+func TestLogin_NoBrowserRequiresPromptURL(t *testing.T) {
+ cfg := baseConfig(t, "https://unused.test/oauth2/token", mustFreeAddr(t))
+ cfg.NoBrowser = true
+ cfg.PromptURL = nil
+
+ _, err := cognito.Login(context.Background(), cfg)
+
+ require.Error(t, err, "a login with no browser and no way to report the url cannot complete")
+ assert.Contains(t, err.Error(), "PromptURL",
+ "the error must name the field that is missing")
+}
+
+func TestLogin_NoBrowserPromptFailurePropagates(t *testing.T) {
+ wantErr := errors.New("no tty to print to")
+
+ cfg := baseConfig(t, "https://unused.test/oauth2/token", mustFreeAddr(t))
+ cfg.NoBrowser = true
+ cfg.PromptURL = func(string) error { return wantErr }
+
+ _, err := cognito.Login(context.Background(), cfg)
+
+ require.Error(t, err)
+ assert.ErrorIs(t, err, wantErr, "a prompt failure must reach the caller unwrapped in meaning")
+ assert.NotContains(t, err.Error(), "open browser", "the browser path was not taken")
+}
+
+// TestLogin_RejectsAnUnusableExpiresIn pins that a token response with no
+// usable expires_in fails loudly.
+//
+// Accepting it silently is worse than failing: ExpiresAt would land on exactly
+// now(), IsExpired subtracts a leeway on top, and a login that had just
+// succeeded would read as already expired — sending the operator back through
+// sign-in on their very next command with no indication why.
+func TestLogin_RejectsAnUnusableExpiresIn(t *testing.T) {
+ tests := []struct {
+ name string
+ givenExpiresIn any
+ wantErrContains string
+ }{
+ {
+ name: "expires_in omitted entirely",
+ givenExpiresIn: nil,
+ wantErrContains: "expires_in",
+ },
+ {
+ name: "expires_in zero",
+ givenExpiresIn: 0,
+ wantErrContains: "expires_in",
+ },
+ {
+ name: "expires_in negative",
+ givenExpiresIn: -1,
+ wantErrContains: "expires_in",
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ body := goodTokenBody(t, testOperatorEmail)
+ delete(body, "expires_in")
+ if tt.givenExpiresIn != nil {
+ body["expires_in"] = tt.givenExpiresIn
+ }
+
+ ts := newTokenServer(t, jsonHandler(http.StatusOK, body))
+ browser := newFakeBrowser()
+ cfg := baseConfig(t, ts.URL, mustFreeAddr(t))
+ cfg.OpenBrowser = browser.open
+
+ _, err := cognito.Login(context.Background(), cfg)
+
+ require.Error(t, err, "an unusable expires_in must not produce a session")
+ assert.Contains(t, err.Error(), tt.wantErrContains)
+ })
+ }
+}
diff --git a/cognito/pkce.go b/cognito/pkce.go
new file mode 100644
index 0000000..63572c6
--- /dev/null
+++ b/cognito/pkce.go
@@ -0,0 +1,56 @@
+package cognito
+
+import (
+ "crypto/rand"
+ "crypto/sha256"
+ "crypto/subtle"
+ "encoding/base64"
+ "fmt"
+
+ wegostrings "github.com/wego/pkg/strings"
+)
+
+const (
+ // verifierBytes yields a 128-character unpadded base64url verifier, the
+ // longest RFC 7636 permits for code_verifier.
+ verifierBytes = 96
+
+ // stateBytes yields a 32-character unpadded base64url state value.
+ stateBytes = 24
+)
+
+// generateVerifier draws a fresh PKCE code_verifier from crypto/rand.
+func generateVerifier() (string, error) {
+ buf := make([]byte, verifierBytes)
+ if _, err := rand.Read(buf); err != nil {
+ return "", fmt.Errorf("generate pkce verifier: %w", err)
+ }
+ return base64.RawURLEncoding.EncodeToString(buf), nil
+}
+
+// generateChallenge derives the S256 code_challenge for a verifier. Both the
+// verifier and the challenge are unpadded base64url, per RFC 7636.
+func generateChallenge(verifier string) string {
+ sum := sha256.Sum256([]byte(verifier))
+ return base64.RawURLEncoding.EncodeToString(sum[:])
+}
+
+// generateState draws a fresh OAuth state value from crypto/rand.
+func generateState() (string, error) {
+ buf := make([]byte, stateBytes)
+ if _, err := rand.Read(buf); err != nil {
+ return "", fmt.Errorf("generate oauth state: %w", err)
+ }
+ return base64.RawURLEncoding.EncodeToString(buf), nil
+}
+
+// stateMatches compares the state a callback presented against the one we
+// minted, in constant time. A blank value on either side never matches: a
+// callback that omits state is precisely the CSRF case the parameter exists to
+// catch, so it must be rejected rather than waved through.
+func stateMatches(want, got string) bool {
+ if wegostrings.IsBlank(want) || wegostrings.IsBlank(got) {
+ return false
+ }
+ return subtle.ConstantTimeCompare([]byte(want), []byte(got)) == 1
+}
diff --git a/cognito/pkce_internal_test.go b/cognito/pkce_internal_test.go
new file mode 100644
index 0000000..7cda37a
--- /dev/null
+++ b/cognito/pkce_internal_test.go
@@ -0,0 +1,104 @@
+package cognito
+
+import (
+ "crypto/sha256"
+ "encoding/base64"
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestGenerateVerifier(t *testing.T) {
+ verifier, err := generateVerifier()
+ require.NoError(t, err)
+
+ assert.GreaterOrEqual(t, len(verifier), 43, "RFC 7636 requires a code_verifier of at least 43 chars, got %q", verifier)
+ assert.LessOrEqual(t, len(verifier), 128, "RFC 7636 caps code_verifier at 128 chars, got %d chars", len(verifier))
+ assert.NotContains(t, verifier, "=", "code_verifier must be base64url WITHOUT padding")
+ assert.NotContains(t, verifier, "+", "code_verifier must be base64url, not standard base64")
+ assert.NotContains(t, verifier, "/", "code_verifier must be base64url, not standard base64")
+
+ _, err = base64.RawURLEncoding.DecodeString(verifier)
+ require.NoError(t, err, "code_verifier must decode as unpadded base64url")
+}
+
+func TestGenerateVerifier_IsUniquePerCall(t *testing.T) {
+ seen := make(map[string]struct{}, 32)
+ for range 32 {
+ verifier, err := generateVerifier()
+ require.NoError(t, err)
+ _, dup := seen[verifier]
+ require.False(t, dup, "generateVerifier returned a duplicate; it must be drawn from crypto/rand")
+ seen[verifier] = struct{}{}
+ }
+}
+
+func TestGenerateChallenge(t *testing.T) {
+ tests := []struct {
+ name string
+ givenVerifier string
+ want string
+ }{
+ {
+ // RFC 7636 appendix B test vector - pins the S256 transformation.
+ name: "rfc 7636 appendix b vector",
+ givenVerifier: "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk",
+ want: "E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM",
+ },
+ {
+ name: "empty verifier still hashes",
+ givenVerifier: "",
+ want: base64.RawURLEncoding.EncodeToString(sha256Sum("")),
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got := generateChallenge(tt.givenVerifier)
+ assert.Equal(t, tt.want, got, "challenge for verifier %q", tt.givenVerifier)
+ assert.NotContains(t, got, "=", "code_challenge must be unpadded base64url")
+ })
+ }
+}
+
+func TestGenerateState(t *testing.T) {
+ state, err := generateState()
+ require.NoError(t, err)
+
+ require.NotEmpty(t, state)
+ assert.NotContains(t, state, "=", "state must be unpadded base64url")
+ _, err = base64.RawURLEncoding.DecodeString(state)
+ require.NoError(t, err, "state must decode as unpadded base64url")
+
+ other, err := generateState()
+ require.NoError(t, err)
+ assert.NotEqual(t, state, other, "state must be freshly drawn from crypto/rand on every call")
+}
+
+func TestStateMatches(t *testing.T) {
+ tests := []struct {
+ name string
+ givenWant string
+ givenGot string
+ want bool
+ }{
+ {name: "mismatch", givenWant: "abc", givenGot: "abd", want: false},
+ {name: "missing callback state", givenWant: "abc", givenGot: "", want: false},
+ {name: "both blank is still a mismatch", givenWant: "", givenGot: "", want: false},
+ {name: "different length", givenWant: "abc", givenGot: "abcdef", want: false},
+ {name: "match", givenWant: "abc", givenGot: "abc", want: true},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ assert.Equal(t, tt.want, stateMatches(tt.givenWant, tt.givenGot))
+ })
+ }
+}
+
+// sha256Sum computes the expectation independently of the code under test.
+func sha256Sum(s string) []byte {
+ sum := sha256.Sum256([]byte(s))
+ return sum[:]
+}
diff --git a/cognito/storage/keyring.go b/cognito/storage/keyring.go
new file mode 100644
index 0000000..cfc3481
--- /dev/null
+++ b/cognito/storage/keyring.go
@@ -0,0 +1,309 @@
+package storage
+
+import (
+ "encoding/json"
+ "errors"
+ "fmt"
+ "time"
+
+ wegostrings "github.com/wego/pkg/strings"
+ "github.com/zalando/go-keyring"
+
+ "github.com/wego/pkg/cognito"
+)
+
+// Layout.
+//
+// A Cognito JWT - the access token especially - can exceed the 4 KiB
+// command-line cap zalando/go-keyring runs into on macOS, where it shells out
+// to /usr/bin/security (see keyring_darwin.go, which refuses any command over
+// 4096 bytes). After the library's base64 expansion that leaves roughly 3 KB
+// of secret per entry, and a whole token set does not reliably fit. So each
+// field keeps its own entry.
+//
+// That rules out getting atomicity by writing one entry, and a keychain has no
+// transaction. Instead each token set is written into one of two SLOTS, and a
+// separate pointer entry names the slot that counts. Save fills the inactive
+// slot and then moves the pointer, which is one small write and the only write
+// that changes what Load sees. A failure anywhere before it leaves the pointer
+// and the live slot untouched, so a torn write costs the new session, never
+// the old one.
+//
+// Two slots rather than a counter keeps the entry count fixed and means a
+// re-login never writes over the entries the current session is read from.
+//
+// Entries are accounted as "/current" for the pointer and
+// "//" for the token fields.
+const (
+ fieldAccess = "access"
+ fieldID = "id"
+ fieldRefresh = "refresh"
+ fieldMeta = "meta"
+ // fieldCurrent is the pointer entry naming the live slot. Committing a
+ // token set is exactly one write of this entry.
+ fieldCurrent = "current"
+)
+
+// The two slots a token set alternates between.
+const (
+ slotA = "a"
+ slotB = "b"
+)
+
+// tokenFields are the per-field entries that make up one stored token set.
+var tokenFields = []string{fieldAccess, fieldID, fieldRefresh, fieldMeta}
+
+// keyringBackend is the slice of zalando/go-keyring this package depends on.
+// Naming it lets tests substitute a fake instead of prompting a real keychain.
+type keyringBackend interface {
+ Set(service, user, password string) error
+ Get(service, user string) (string, error)
+ Delete(service, user string) error
+}
+
+// systemKeyring is the real OS keychain.
+type systemKeyring struct{}
+
+func (systemKeyring) Set(service, user, password string) error {
+ return keyring.Set(service, user, password)
+}
+
+func (systemKeyring) Get(service, user string) (string, error) {
+ return keyring.Get(service, user)
+}
+
+func (systemKeyring) Delete(service, user string) error {
+ return keyring.Delete(service, user)
+}
+
+// keyringMeta is the non-secret metadata, kept in one small entry beside the
+// tokens themselves.
+type keyringMeta struct {
+ ExpiresAt time.Time `json:"expires_at"`
+}
+
+// keyringStore persists tokens in the OS keychain.
+type keyringStore struct {
+ service string
+ backend keyringBackend
+}
+
+// NewKeyring returns a Store backed by the operating system keychain, with
+// every entry filed under service.
+func NewKeyring(service string) Store {
+ return &keyringStore{service: service, backend: systemKeyring{}}
+}
+
+// Load returns the tokens stored under namespace, or (nil, nil) if the
+// operator is not signed in.
+func (k *keyringStore) Load(namespace string) (*cognito.TokenSet, error) {
+ if err := k.validate(namespace); err != nil {
+ return nil, err
+ }
+
+ // Gate on the pointer: its absence means "not signed in", which is a normal
+ // state rather than a failure. Once it is present, anything the slot it
+ // names is missing is a real inconsistency and must be reported.
+ slot, err := k.backend.Get(k.service, pointerAccount(namespace))
+ if errors.Is(err, keyring.ErrNotFound) {
+ return nil, nil
+ }
+ if err != nil {
+ return nil, fmt.Errorf("read the keychain: %w (is the system keychain unlocked?)", err)
+ }
+
+ if slot != slotA && slot != slotB {
+ // Nothing here writes any other value, so this entry was tampered with
+ // or written by a version that stored something else. Guessing which
+ // slot was meant would be worse than refusing.
+ return nil, fmt.Errorf("keychain names an unknown token slot %q for %q", slot, namespace)
+ }
+
+ access, err := k.read(namespace, slot, fieldAccess)
+ if err != nil {
+ return nil, err
+ }
+ idToken, err := k.read(namespace, slot, fieldID)
+ if err != nil {
+ return nil, err
+ }
+ refresh, err := k.read(namespace, slot, fieldRefresh)
+ if err != nil {
+ return nil, err
+ }
+ rawMeta, err := k.read(namespace, slot, fieldMeta)
+ if err != nil {
+ return nil, err
+ }
+
+ var meta keyringMeta
+ if err := json.Unmarshal([]byte(rawMeta), &meta); err != nil {
+ return nil, fmt.Errorf("parse token metadata from the keychain: %w", err)
+ }
+
+ return &cognito.TokenSet{
+ AccessToken: access,
+ IDToken: idToken,
+ RefreshToken: refresh,
+ ExpiresAt: meta.ExpiresAt,
+ }, nil
+}
+
+// Save writes tokens under namespace, replacing anything already there.
+//
+// It is atomic from Load's point of view: the fields go into the slot that is
+// not live, and only the final pointer write makes them the session. If any
+// write fails, the previous session is still whole and still what Load
+// returns.
+func (k *keyringStore) Save(namespace string, tokens *cognito.TokenSet) error {
+ if err := k.validate(namespace); err != nil {
+ return err
+ }
+ if tokens == nil {
+ return errNoTokens
+ }
+
+ meta, err := json.Marshal(keyringMeta{ExpiresAt: tokens.ExpiresAt})
+ if err != nil {
+ return fmt.Errorf("encode token metadata: %w", err)
+ }
+
+ live, err := k.currentSlot(namespace)
+ if err != nil {
+ return err
+ }
+
+ next := otherSlot(live)
+
+ values := map[string]string{
+ fieldAccess: tokens.AccessToken,
+ fieldID: tokens.IDToken,
+ fieldRefresh: tokens.RefreshToken,
+ fieldMeta: string(meta),
+ }
+
+ // Ordering is irrelevant now: nothing written here is reachable until the
+ // pointer moves. tokenFields is used rather than ranging the map so the
+ // write sequence is deterministic, which keeps failures reproducible.
+ for _, field := range tokenFields {
+ if err := k.backend.Set(k.service, fieldAccount(namespace, next, field), values[field]); err != nil {
+ return fmt.Errorf("write %s to the keychain: %w (is the system keychain unlocked?)", field, err)
+ }
+ }
+
+ // The commit.
+ if err := k.backend.Set(k.service, pointerAccount(namespace), next); err != nil {
+ return fmt.Errorf("commit the session to the keychain: %w (is the system keychain unlocked?)", err)
+ }
+
+ // The old slot is now unreachable. Clearing it keeps a superseded token set
+ // out of the keychain, but the session is already committed, so a failure
+ // here is not the caller's problem and must not fail the sign-in.
+ k.clearSlot(namespace, live)
+
+ return nil
+}
+
+// currentSlot reports the live slot, or "" when nothing is stored. An
+// unrecognised value is treated as "nothing live": Save's job is to establish a
+// good session, and it can do that without deciding what the bad value meant.
+// Load is where a corrupt pointer is reported.
+func (k *keyringStore) currentSlot(namespace string) (string, error) {
+ slot, err := k.backend.Get(k.service, pointerAccount(namespace))
+ if errors.Is(err, keyring.ErrNotFound) {
+ return "", nil
+ }
+ if err != nil {
+ return "", fmt.Errorf("read the current token slot: %w (is the system keychain unlocked?)", err)
+ }
+
+ if slot != slotA && slot != slotB {
+ return "", nil
+ }
+
+ return slot, nil
+}
+
+// otherSlot returns the slot to write next. An empty live slot means nothing is
+// stored, so either is free and slotA keeps a first login predictable.
+func otherSlot(live string) string {
+ if live == slotA {
+ return slotB
+ }
+
+ return slotA
+}
+
+// clearSlot removes one slot's field entries, best effort. Callers use it for a
+// slot nothing points at any more.
+func (k *keyringStore) clearSlot(namespace, slot string) {
+ if slot == "" {
+ return
+ }
+
+ for _, field := range tokenFields {
+ _ = k.backend.Delete(k.service, fieldAccount(namespace, slot, field))
+ }
+}
+
+// Delete removes every entry under namespace. Entries already gone are fine:
+// the desired end state is "signed out", so signing out twice must succeed.
+//
+// The pointer goes FIRST, for the same reason Save moves it last: once it is
+// gone the operator is signed out, even if a later delete fails and leaves
+// orphaned field entries behind.
+func (k *keyringStore) Delete(namespace string) error {
+ if err := k.validate(namespace); err != nil {
+ return err
+ }
+
+ err := k.backend.Delete(k.service, pointerAccount(namespace))
+ if err != nil && !errors.Is(err, keyring.ErrNotFound) {
+ return fmt.Errorf("delete the current token slot from the keychain: %w (is the system keychain unlocked?)", err)
+ }
+
+ // Both slots, not just the live one: a torn Save can leave fields in the
+ // inactive slot, and signing out should not leave a token set behind.
+ for _, slot := range []string{slotA, slotB} {
+ for _, field := range tokenFields {
+ err := k.backend.Delete(k.service, fieldAccount(namespace, slot, field))
+ if err != nil && !errors.Is(err, keyring.ErrNotFound) {
+ return fmt.Errorf("delete %s from the keychain: %w (is the system keychain unlocked?)", field, err)
+ }
+ }
+ }
+
+ return nil
+}
+
+// read fetches one field of one slot, treating a missing entry as an error:
+// callers only reach it after the pointer has confirmed a committed session,
+// so anything absent is an inconsistency rather than a logged-out state.
+func (k *keyringStore) read(namespace, slot, field string) (string, error) {
+ value, err := k.backend.Get(k.service, fieldAccount(namespace, slot, field))
+ if err != nil {
+ return "", fmt.Errorf("read %s from the keychain: %w", field, err)
+ }
+ return value, nil
+}
+
+// validate checks the inputs every operation needs.
+func (k *keyringStore) validate(namespace string) error {
+ if wegostrings.IsBlank(k.service) {
+ return errors.New("keychain service name must not be blank")
+ }
+ if wegostrings.IsBlank(namespace) {
+ return errBlankNamespace
+ }
+ return nil
+}
+
+// pointerAccount is the keychain account name of a namespace's commit pointer.
+func pointerAccount(namespace string) string {
+ return namespace + "/" + fieldCurrent
+}
+
+// fieldAccount is the keychain account name for one field of one slot.
+func fieldAccount(namespace, slot, field string) string {
+ return namespace + "/" + slot + "/" + field
+}
diff --git a/cognito/storage/keyring_internal_test.go b/cognito/storage/keyring_internal_test.go
new file mode 100644
index 0000000..4770377
--- /dev/null
+++ b/cognito/storage/keyring_internal_test.go
@@ -0,0 +1,616 @@
+package storage
+
+import (
+ "errors"
+ "sync"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+ wegostrings "github.com/wego/pkg/strings"
+ "github.com/zalando/go-keyring"
+
+ "github.com/wego/pkg/cognito"
+)
+
+const (
+ testService = "pay-admin-test"
+ testNamespace = "pay-admin/staging"
+)
+
+var testExpiry = time.Date(2026, 9, 1, 12, 0, 0, 0, time.UTC)
+
+// TestStoreContract holds every Store implementation to the same behaviour.
+// The keyring case runs against a fake backend, so no test ever pops an OS
+// keychain prompt.
+func TestStoreContract(t *testing.T) {
+ impls := []struct {
+ name string
+ build func() Store
+ }{
+ {
+ name: "memory",
+ build: NewMemory,
+ },
+ {
+ name: "keyring",
+ build: func() Store {
+ return &keyringStore{service: testService, backend: newFakeKeyring()}
+ },
+ },
+ }
+
+ tests := []struct {
+ name string
+ run func(t *testing.T, store Store)
+ }{
+ {
+ name: "a blank namespace is rejected",
+ run: func(t *testing.T, store Store) {
+ _, err := store.Load("")
+ require.Error(t, err)
+ require.Error(t, store.Save("", sampleTokens()))
+ require.Error(t, store.Delete(""))
+ },
+ },
+ {
+ name: "saving nil tokens is rejected",
+ run: func(t *testing.T, store Store) {
+ require.Error(t, store.Save(testNamespace, nil))
+ },
+ },
+ {
+ name: "loading an unknown namespace is not an error",
+ run: func(t *testing.T, store Store) {
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err, "not logged in is a normal state, not a failure")
+ assert.Nil(t, got)
+ },
+ },
+ {
+ name: "deleting an unknown namespace is not an error",
+ run: func(t *testing.T, store Store) {
+ require.NoError(t, store.Delete(testNamespace))
+ },
+ },
+ {
+ name: "delete clears a saved token set",
+ run: func(t *testing.T, store Store) {
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+ require.NoError(t, store.Delete(testNamespace))
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ assert.Nil(t, got)
+ },
+ },
+ {
+ name: "namespaces do not collide",
+ run: func(t *testing.T, store Store) {
+ staging := sampleTokens()
+ staging.AccessToken = "staging-access"
+ production := sampleTokens()
+ production.AccessToken = "production-access"
+
+ require.NoError(t, store.Save("pay-admin/staging", staging))
+ require.NoError(t, store.Save("pay-admin/production", production))
+
+ got, err := store.Load("pay-admin/staging")
+ require.NoError(t, err)
+ require.NotNil(t, got)
+ assert.Equal(t, "staging-access", got.AccessToken)
+
+ require.NoError(t, store.Delete("pay-admin/staging"))
+
+ survivor, err := store.Load("pay-admin/production")
+ require.NoError(t, err)
+ require.NotNil(t, survivor, "deleting one namespace must not touch another")
+ assert.Equal(t, "production-access", survivor.AccessToken)
+ },
+ },
+ {
+ name: "save then load round-trips every field",
+ run: func(t *testing.T, store Store) {
+ want := sampleTokens()
+ require.NoError(t, store.Save(testNamespace, want))
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ require.NotNil(t, got)
+
+ assert.Equal(t, want.AccessToken, got.AccessToken)
+ assert.Equal(t, want.IDToken, got.IDToken)
+ assert.Equal(t, want.RefreshToken, got.RefreshToken)
+ assert.Equal(t, want.ExpiresAt.UTC(), got.ExpiresAt.UTC())
+ },
+ },
+ {
+ name: "save overwrites an earlier token set",
+ run: func(t *testing.T, store Store) {
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+
+ replacement := sampleTokens()
+ replacement.AccessToken = "second-access"
+ require.NoError(t, store.Save(testNamespace, replacement))
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ require.NotNil(t, got)
+ assert.Equal(t, "second-access", got.AccessToken)
+ },
+ },
+ {
+ name: "the loaded token set is a copy",
+ run: func(t *testing.T, store Store) {
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+
+ first, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ require.NotNil(t, first)
+ first.AccessToken = "mutated-by-caller"
+
+ second, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ require.NotNil(t, second)
+ assert.Equal(t, "access-value", second.AccessToken, "a caller must not be able to mutate stored tokens")
+ },
+ },
+ }
+
+ for _, impl := range impls {
+ for _, tt := range tests {
+ t.Run(impl.name+"/"+tt.name, func(t *testing.T) {
+ tt.run(t, impl.build())
+ })
+ }
+ }
+}
+
+func TestKeyringStore_LoadFailures(t *testing.T) {
+ tests := []struct {
+ name string
+ givenSetup func(*fakeKeyring)
+ wantNil bool
+ wantErrContains string
+ }{
+ {
+ name: "a keychain failure on the gate entry is reported",
+ givenSetup: func(f *fakeKeyring) {
+ f.getErr = errors.New("keychain is locked")
+ },
+ wantErrContains: "keychain",
+ },
+ {
+ // The pointer says a session was committed, so anything the slot is
+ // missing is an inconsistency rather than a logged-out state.
+ name: "a missing id_token is an inconsistency, not a logged-out state",
+ givenSetup: func(f *fakeKeyring) {
+ f.putCommitted(map[string]string{fieldAccess: "access-value"})
+ },
+ wantErrContains: "id",
+ },
+ {
+ name: "a missing refresh_token is an inconsistency",
+ givenSetup: func(f *fakeKeyring) {
+ f.putCommitted(map[string]string{
+ fieldAccess: "access-value",
+ fieldID: "id-value",
+ })
+ },
+ wantErrContains: "refresh",
+ },
+ {
+ name: "missing metadata is an inconsistency",
+ givenSetup: func(f *fakeKeyring) {
+ f.putCommitted(map[string]string{
+ fieldAccess: "access-value",
+ fieldID: "id-value",
+ fieldRefresh: "refresh-value",
+ })
+ },
+ wantErrContains: "meta",
+ },
+ {
+ name: "corrupt metadata is reported",
+ givenSetup: func(f *fakeKeyring) {
+ f.putCommitted(map[string]string{
+ fieldAccess: "access-value",
+ fieldID: "id-value",
+ fieldRefresh: "refresh-value",
+ fieldMeta: "{not json",
+ })
+ },
+ wantErrContains: "metadata",
+ },
+ {
+ // Fields present but never committed: the sign-in was torn before
+ // the pointer moved, so the operator is simply not signed in.
+ name: "an uncommitted slot reads as not logged in",
+ givenSetup: func(f *fakeKeyring) {
+ f.putField(slotA, fieldAccess, "access-value")
+ f.putField(slotA, fieldID, "id-value")
+ },
+ wantNil: true,
+ },
+ {
+ name: "nothing stored reads as not logged in",
+ givenSetup: func(_ *fakeKeyring) {},
+ wantNil: true,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ backend := newFakeKeyring()
+ tt.givenSetup(backend)
+ store := &keyringStore{service: testService, backend: backend}
+
+ got, err := store.Load(testNamespace)
+
+ if wegostrings.IsNotEmpty(tt.wantErrContains) {
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), tt.wantErrContains)
+ assert.Nil(t, got)
+ return
+ }
+ require.NoError(t, err)
+ assert.Nil(t, got)
+ })
+ }
+}
+
+// TestKeyringStore_SaveCommitsWithOnePointerWrite pins both halves of the
+// storage contract.
+//
+// Fields stay in SEPARATE entries because zalando/go-keyring shells out to
+// /usr/bin/security on macOS and refuses any command over 4096 bytes
+// (keyring_darwin.go). After its base64 expansion that leaves roughly 3 KB of
+// secret per entry, which a single combined token set would exceed.
+//
+// Atomicity therefore cannot come from writing one entry. It comes from
+// writing the fields into the inactive slot and then moving the pointer, which
+// is ONE small Set and the only write that changes what Load sees.
+func TestKeyringStore_SaveCommitsWithOnePointerWrite(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+
+ require.NotEmpty(t, backend.writes)
+ assert.Equal(t, testNamespace+"/"+fieldCurrent, backend.writes[len(backend.writes)-1],
+ "the pointer must be the last write, so nothing before it is observable")
+ assert.Equal(t, 1, countWrites(backend.writes, testNamespace+"/"+fieldCurrent),
+ "the commit must be a single pointer write")
+
+ for _, field := range []string{fieldAccess, fieldID, fieldRefresh, fieldMeta} {
+ assert.Contains(t, backend.writes, testNamespace+"/"+slotA+"/"+field,
+ "each field keeps its own entry, to stay under the 4 KiB cap")
+ }
+}
+
+// TestKeyringStore_SaveAlternatesSlots covers why there are two slots: a
+// re-login must not overwrite the entries the current session is still being
+// read from, or a torn write would corrupt the live session rather than an
+// unused copy.
+func TestKeyringStore_SaveAlternatesSlots(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+ assert.Equal(t, slotA, backend.value(fieldCurrent))
+
+ second := sampleTokens()
+ second.AccessToken = "second-access"
+ require.NoError(t, store.Save(testNamespace, second))
+ assert.Equal(t, slotB, backend.value(fieldCurrent))
+
+ third := sampleTokens()
+ third.AccessToken = "third-access"
+ require.NoError(t, store.Save(testNamespace, third))
+ assert.Equal(t, slotA, backend.value(fieldCurrent), "slots alternate rather than growing")
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ assert.Equal(t, "third-access", got.AccessToken)
+}
+
+// TestKeyringStore_TornResaveLeavesThePreviousSession is the regression test
+// for the reported defect.
+//
+// With one entry per field and no commit pointer, a re-login that failed
+// partway left the NEW refresh token beside the OLD access token, id token and
+// expiry. Load gated on the access token, which was present, so it handed that
+// mixture back as a live session: a stale expiry paired with a fresh refresh
+// token, or after an account switch one identity's id token paired with
+// another's refresh token.
+func TestKeyringStore_TornResaveLeavesThePreviousSession(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+
+ // Let two writes of the re-login land, then fail.
+ backend.okWrites = 0
+ backend.failSetAfter = 2
+
+ newer := &cognito.TokenSet{
+ AccessToken: "new-access",
+ IDToken: "new-id",
+ RefreshToken: "new-refresh",
+ ExpiresAt: testExpiry.Add(time.Hour),
+ }
+ require.Error(t, store.Save(testNamespace, newer))
+
+ backend.failSetAfter = -1
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err, "the previous session must still be readable")
+ require.NotNil(t, got)
+ assert.Equal(t, sampleTokens(), got,
+ "a torn re-login must leave the previous session exactly, never a mixture of the two")
+}
+
+// TestKeyringStore_LoadRejectsAnUnknownSlot covers a pointer naming a slot that
+// is not one of the two. That cannot arise from this code, so it means the
+// entry was tampered with or written by another version, and guessing which
+// slot was meant would be worse than refusing.
+func TestKeyringStore_LoadRejectsAnUnknownSlot(t *testing.T) {
+ backend := newFakeKeyring()
+ backend.put(fieldCurrent, "somewhere-else")
+ store := &keyringStore{service: testService, backend: backend}
+
+ _, err := store.Load(testNamespace)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "slot")
+}
+
+// countWrites reports how many times account appears in a write log.
+func countWrites(writes []string, account string) int {
+ n := 0
+ for _, w := range writes {
+ if w == account {
+ n++
+ }
+ }
+
+ return n
+}
+
+func TestKeyringStore_BackendFailuresAreReported(t *testing.T) {
+ tests := []struct {
+ name string
+ givenSetup func(*fakeKeyring)
+ givenAction func(Store) error
+ wantErrContains string
+ }{
+ {
+ name: "a save failure is reported",
+ givenSetup: func(f *fakeKeyring) {
+ f.setErr = errors.New("keychain is locked")
+ },
+ givenAction: func(s Store) error { return s.Save(testNamespace, sampleTokens()) },
+ wantErrContains: "keychain",
+ },
+ {
+ name: "a delete failure is reported",
+ givenSetup: func(f *fakeKeyring) {
+ f.deleteErr = errors.New("keychain is locked")
+ },
+ givenAction: func(s Store) error { return s.Delete(testNamespace) },
+ wantErrContains: "keychain",
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ backend := newFakeKeyring()
+ tt.givenSetup(backend)
+ store := &keyringStore{service: testService, backend: backend}
+
+ err := tt.givenAction(store)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), tt.wantErrContains)
+ })
+ }
+}
+
+// TestKeyringStore_DeleteToleratesMissingEntries covers a half-populated
+// namespace: sign-out must succeed rather than stranding the operator, and it
+// must clear the slot a torn Save left behind as well as the live one.
+func TestKeyringStore_DeleteToleratesMissingEntries(t *testing.T) {
+ backend := newFakeKeyring()
+ backend.putField(slotA, fieldAccess, "access-value")
+ backend.putField(slotB, fieldRefresh, "orphaned-refresh")
+ store := &keyringStore{service: testService, backend: backend}
+
+ require.NoError(t, store.Delete(testNamespace))
+ assert.Empty(t, backend.entries, "no token material may survive a sign-out, in either slot")
+}
+
+// TestKeyringStore_FailedCommitKeepsThePreviousSession covers the last write.
+// Every field of the new session is already in the keychain; only the pointer
+// move failed. Load must still return the old session, because the new one was
+// never committed.
+func TestKeyringStore_FailedCommitKeepsThePreviousSession(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+
+ // Let all four field writes land, then fail the pointer write.
+ backend.okWrites = 0
+ backend.failSetAfter = len(tokenFields)
+
+ newer := sampleTokens()
+ newer.AccessToken = "new-access"
+
+ err := store.Save(testNamespace, newer)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "commit", "the failure must name the commit, not a field")
+
+ backend.failSetAfter = -1
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ assert.Equal(t, sampleTokens(), got, "an uncommitted session must not become the live one")
+}
+
+// TestKeyringStore_SaveReportsAnUnreadablePointer covers a keychain that
+// cannot be read at all. Save must not proceed on a guess about which slot is
+// live: writing to the wrong one would overwrite the session it is meant to
+// protect.
+func TestKeyringStore_SaveReportsAnUnreadablePointer(t *testing.T) {
+ backend := newFakeKeyring()
+ backend.getErr = errors.New("keychain is locked")
+ store := &keyringStore{service: testService, backend: backend}
+
+ err := store.Save(testNamespace, sampleTokens())
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "slot")
+ assert.Empty(t, backend.writes, "nothing may be written when the live slot is unknown")
+}
+
+// TestKeyringStore_SaveOverAnUnknownSlotStillSignsIn covers a pointer holding
+// a value this code never writes. Load refuses it, but Save's job is to
+// establish a good session and it can do that without deciding what the bad
+// value meant.
+func TestKeyringStore_SaveOverAnUnknownSlotStillSignsIn(t *testing.T) {
+ backend := newFakeKeyring()
+ backend.put(fieldCurrent, "somewhere-else")
+ store := &keyringStore{service: testService, backend: backend}
+
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+ assert.Equal(t, slotA, backend.value(fieldCurrent))
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ assert.Equal(t, sampleTokens(), got)
+}
+
+// TestKeyringStore_DeleteReportsAPointerFailure covers sign-out when the
+// pointer cannot be removed. That is the entry that decides whether the
+// operator is signed in, so the failure has to be reported rather than
+// swallowed after clearing the fields.
+func TestKeyringStore_DeleteReportsAPointerFailure(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+
+ backend.deleteErr = errors.New("keychain is locked")
+
+ err := store.Delete(testNamespace)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "slot")
+}
+
+func TestKeyringStore_BlankServiceIsRejected(t *testing.T) {
+ store := &keyringStore{service: "", backend: newFakeKeyring()}
+
+ _, err := store.Load(testNamespace)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "service")
+ require.Error(t, store.Save(testNamespace, sampleTokens()))
+ require.Error(t, store.Delete(testNamespace))
+}
+
+func sampleTokens() *cognito.TokenSet {
+ return &cognito.TokenSet{
+ AccessToken: "access-value",
+ IDToken: "id-value",
+ RefreshToken: "refresh-value",
+ ExpiresAt: testExpiry,
+ }
+}
+
+// fakeKeyring is an in-memory stand-in for the OS keychain.
+type fakeKeyring struct {
+ mu sync.Mutex
+ entries map[string]string
+ writes []string
+ setErr error
+ // failSetAfter fails every Set from the (failSetAfter+1)th of this counter
+ // onward, so a test can tear a write sequence in the middle. Negative
+ // disables it. Reset okWrites to re-arm.
+ failSetAfter int
+ okWrites int
+ getErr error
+ deleteErr error
+}
+
+func newFakeKeyring() *fakeKeyring {
+ return &fakeKeyring{entries: make(map[string]string), failSetAfter: -1}
+}
+
+func (f *fakeKeyring) Set(service, user, password string) error {
+ f.mu.Lock()
+ defer f.mu.Unlock()
+ if f.setErr != nil {
+ return f.setErr
+ }
+ if f.failSetAfter >= 0 && f.okWrites >= f.failSetAfter {
+ return errors.New("keychain is locked")
+ }
+ f.okWrites++
+ f.entries[service+"|"+user] = password
+ f.writes = append(f.writes, user)
+ return nil
+}
+
+func (f *fakeKeyring) Get(service, user string) (string, error) {
+ f.mu.Lock()
+ defer f.mu.Unlock()
+ if f.getErr != nil {
+ return "", f.getErr
+ }
+ value, ok := f.entries[service+"|"+user]
+ if !ok {
+ return "", keyring.ErrNotFound
+ }
+ return value, nil
+}
+
+func (f *fakeKeyring) Delete(service, user string) error {
+ f.mu.Lock()
+ defer f.mu.Unlock()
+ if f.deleteErr != nil {
+ return f.deleteErr
+ }
+ key := service + "|" + user
+ if _, ok := f.entries[key]; !ok {
+ return keyring.ErrNotFound
+ }
+ delete(f.entries, key)
+ return nil
+}
+
+// put seeds a namespace-level entry (the pointer) under testNamespace,
+// bypassing the store under test.
+func (f *fakeKeyring) put(field, value string) {
+ f.mu.Lock()
+ defer f.mu.Unlock()
+ f.entries[testService+"|"+testNamespace+"/"+field] = value
+}
+
+// putField seeds one field of one slot, bypassing the store under test.
+func (f *fakeKeyring) putField(slot, field, value string) {
+ f.mu.Lock()
+ defer f.mu.Unlock()
+ f.entries[testService+"|"+testNamespace+"/"+slot+"/"+field] = value
+}
+
+// putCommitted seeds a committed session in slotA, field by field, so a test
+// can then remove or corrupt exactly one part of it.
+func (f *fakeKeyring) putCommitted(fields map[string]string) {
+ for field, value := range fields {
+ f.putField(slotA, field, value)
+ }
+ f.put(fieldCurrent, slotA)
+}
+
+// value reads an entry under testNamespace, bypassing the store under test.
+func (f *fakeKeyring) value(field string) string {
+ f.mu.Lock()
+ defer f.mu.Unlock()
+ return f.entries[testService+"|"+testNamespace+"/"+field]
+}
diff --git a/cognito/storage/store.go b/cognito/storage/store.go
new file mode 100644
index 0000000..abd6bf0
--- /dev/null
+++ b/cognito/storage/store.go
@@ -0,0 +1,99 @@
+// Package storage persists the Cognito token sets obtained by
+// github.com/wego/pkg/cognito.
+//
+// Tokens live under an opaque namespace, e.g. "my-cli/staging". The caller
+// chooses the string and the store never interprets it; namespacing is what
+// keeps a staging login from overwriting a production one.
+//
+// Use NewMemory in tests and anywhere a keychain is unavailable; it satisfies
+// the same interface as NewKeyring, so a caller swaps one for the other without
+// touching its own code.
+package storage
+
+import (
+ "errors"
+ "sync"
+
+ wegostrings "github.com/wego/pkg/strings"
+
+ "github.com/wego/pkg/cognito"
+)
+
+// Store persists a TokenSet under an opaque namespace.
+//
+// A missing entry is NOT an error: Load returns (nil, nil) when nothing is
+// stored, because "not logged in" is a normal state. Callers can therefore
+// tell it apart from a backend that is genuinely broken, which does return an
+// error.
+type Store interface {
+ Load(namespace string) (*cognito.TokenSet, error)
+ Save(namespace string, tokens *cognito.TokenSet) error
+ Delete(namespace string) error
+}
+
+var (
+ errBlankNamespace = errors.New("storage namespace must not be blank")
+ errNoTokens = errors.New("no tokens to save")
+)
+
+// memoryStore keeps tokens in process memory only.
+type memoryStore struct {
+ mu sync.RWMutex
+ tokens map[string]cognito.TokenSet
+}
+
+// NewMemory returns a Store that holds tokens for the lifetime of the process
+// and nothing longer. It serves tests and callers with no usable keychain; a
+// new process always starts signed out.
+func NewMemory() Store {
+ return &memoryStore{tokens: make(map[string]cognito.TokenSet)}
+}
+
+// Load returns the tokens stored under namespace, or (nil, nil) if there are
+// none.
+func (m *memoryStore) Load(namespace string) (*cognito.TokenSet, error) {
+ if wegostrings.IsBlank(namespace) {
+ return nil, errBlankNamespace
+ }
+
+ m.mu.RLock()
+ defer m.mu.RUnlock()
+
+ stored, ok := m.tokens[namespace]
+ if !ok {
+ return nil, nil
+ }
+
+ // stored is already a copy of the map value, so a caller mutating the
+ // result cannot reach back into the store.
+ return &stored, nil
+}
+
+// Save writes tokens under namespace, replacing anything already there.
+func (m *memoryStore) Save(namespace string, tokens *cognito.TokenSet) error {
+ if wegostrings.IsBlank(namespace) {
+ return errBlankNamespace
+ }
+ if tokens == nil {
+ return errNoTokens
+ }
+
+ m.mu.Lock()
+ defer m.mu.Unlock()
+ m.tokens[namespace] = *tokens
+
+ return nil
+}
+
+// Delete removes the tokens under namespace. Deleting nothing is not an error.
+func (m *memoryStore) Delete(namespace string) error {
+ if wegostrings.IsBlank(namespace) {
+ return errBlankNamespace
+ }
+
+ m.mu.Lock()
+ defer m.mu.Unlock()
+ delete(m.tokens, namespace)
+
+ return nil
+}
diff --git a/cognito/storage/store_test.go b/cognito/storage/store_test.go
new file mode 100644
index 0000000..5b33b23
--- /dev/null
+++ b/cognito/storage/store_test.go
@@ -0,0 +1,80 @@
+package storage_test
+
+import (
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+
+ "github.com/wego/pkg/cognito"
+ "github.com/wego/pkg/cognito/storage"
+)
+
+func TestNewMemory_RoundTripsThroughThePublicAPI(t *testing.T) {
+ store := storage.NewMemory()
+ want := &cognito.TokenSet{
+ AccessToken: "access-value",
+ IDToken: "id-value",
+ RefreshToken: "refresh-value",
+ ExpiresAt: time.Date(2026, 9, 1, 12, 0, 0, 0, time.UTC),
+ }
+
+ before, err := store.Load("pay-admin/staging")
+ require.NoError(t, err, "an empty store must report not-logged-in, not an error")
+ require.Nil(t, before)
+
+ require.NoError(t, store.Save("pay-admin/staging", want))
+
+ got, err := store.Load("pay-admin/staging")
+ require.NoError(t, err)
+ require.NotNil(t, got)
+ assert.Equal(t, *want, *got)
+
+ require.NoError(t, store.Delete("pay-admin/staging"))
+
+ after, err := store.Load("pay-admin/staging")
+ require.NoError(t, err)
+ assert.Nil(t, after)
+}
+
+// TestNewKeyring_ValidatesBeforeTouchingTheKeychain exercises the exported
+// constructor without ever reaching the OS keychain: every case here is
+// rejected by validation first, so no test can pop a keychain prompt.
+func TestNewKeyring_ValidatesBeforeTouchingTheKeychain(t *testing.T) {
+ tests := []struct {
+ name string
+ givenService string
+ givenNamespace string
+ }{
+ {
+ name: "a blank service is rejected",
+ givenService: "",
+ givenNamespace: "pay-admin/staging",
+ },
+ {
+ name: "a blank namespace is rejected",
+ givenService: "pay-admin",
+ givenNamespace: "",
+ },
+ {
+ name: "both blank is rejected",
+ givenService: "",
+ givenNamespace: "",
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ store := storage.NewKeyring(tt.givenService)
+ require.NotNil(t, store)
+
+ got, err := store.Load(tt.givenNamespace)
+ require.Error(t, err)
+ assert.Nil(t, got)
+
+ require.Error(t, store.Save(tt.givenNamespace, &cognito.TokenSet{AccessToken: "access-value"}))
+ require.Error(t, store.Delete(tt.givenNamespace))
+ })
+ }
+}
diff --git a/cognito/tokens.go b/cognito/tokens.go
new file mode 100644
index 0000000..b9d64e9
--- /dev/null
+++ b/cognito/tokens.go
@@ -0,0 +1,88 @@
+package cognito
+
+import (
+ "encoding/base64"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "strings"
+ "time"
+
+ wegostrings "github.com/wego/pkg/strings"
+)
+
+// expiryLeeway is shaved off the real expiry so a token cannot lapse midway
+// through a request we have already started.
+const expiryLeeway = 60 * time.Second
+
+// TokenSet is the set of Cognito tokens a signed-in operator holds.
+//
+// Holding tokens is the whole point of the type, so gosec's secret-field
+// warning is expected here rather than a finding: these values are never
+// logged, and the only place they are serialized to is the operator's own
+// keychain (see github.com/wego/pkg/cognito/storage).
+type TokenSet struct {
+ AccessToken string `json:"access_token"` //nolint:gosec // G117: this type exists to carry tokens; stored only in the local keychain, never logged.
+ IDToken string `json:"id_token"`
+ RefreshToken string `json:"refresh_token"` //nolint:gosec // G117: as above - the refresh token is the credential this type is for.
+ ExpiresAt time.Time `json:"expires_at"`
+}
+
+// IsExpired reports whether the access token is spent as of now, treating
+// anything inside expiryLeeway of the deadline as already gone. A nil set, or
+// one with no recorded expiry, counts as expired so callers refresh rather
+// than send a token that may already be dead.
+func (t *TokenSet) IsExpired(now time.Time) bool {
+ if t == nil || t.ExpiresAt.IsZero() {
+ return true
+ }
+ return !now.Before(t.ExpiresAt.Add(-expiryLeeway))
+}
+
+// Email returns the email claim from the id_token.
+//
+// It reads the JWT payload WITHOUT verifying the signature, which is correct
+// here: the token came from our own keychain, and every server that accepts it
+// verifies the signature itself. Treat this as decoding for display and
+// client-side UX, NOT as validation - a caller must never make an
+// authorization decision on the strength of these claims.
+func (t *TokenSet) Email() (string, error) {
+ if t == nil {
+ return "", errors.New("no tokens: sign in first")
+ }
+
+ claims, err := decodeIDTokenClaims(t.IDToken)
+ if err != nil {
+ return "", err
+ }
+
+ email, ok := claims["email"].(string)
+ if !ok || wegostrings.IsBlank(email) {
+ return "", errors.New("id_token has no email claim")
+ }
+ return email, nil
+}
+
+// decodeIDTokenClaims splits a JWT and JSON-decodes its payload. It never
+// panics on a malformed or truncated token: every failure is an error.
+func decodeIDTokenClaims(idToken string) (map[string]any, error) {
+ parts := strings.Split(idToken, ".")
+ if len(parts) != 3 {
+ return nil, errors.New("id_token is not a jwt")
+ }
+
+ payload, err := base64.RawURLEncoding.DecodeString(parts[1])
+ if err != nil {
+ // Some encoders pad their segments; accept those too.
+ payload, err = base64.URLEncoding.DecodeString(parts[1])
+ if err != nil {
+ return nil, fmt.Errorf("decode id_token payload: %w", err)
+ }
+ }
+
+ var claims map[string]any
+ if err := json.Unmarshal(payload, &claims); err != nil {
+ return nil, fmt.Errorf("parse id_token payload: %w", err)
+ }
+ return claims, nil
+}
diff --git a/cognito/tokens_test.go b/cognito/tokens_test.go
new file mode 100644
index 0000000..96ede3a
--- /dev/null
+++ b/cognito/tokens_test.go
@@ -0,0 +1,172 @@
+package cognito_test
+
+import (
+ "encoding/base64"
+ "encoding/json"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+ wegostrings "github.com/wego/pkg/strings"
+
+ "github.com/wego/pkg/cognito"
+)
+
+const testOperatorEmail = "ops@wego.com"
+
+var fixedNow = time.Date(2026, 9, 1, 12, 0, 0, 0, time.UTC)
+
+func TestTokenSet_IsExpired(t *testing.T) {
+ tests := []struct {
+ name string
+ given *cognito.TokenSet
+ want bool
+ }{
+ {
+ name: "nil token set counts as expired",
+ given: nil,
+ want: true,
+ },
+ {
+ name: "zero expiry counts as expired",
+ given: &cognito.TokenSet{},
+ want: true,
+ },
+ {
+ name: "already elapsed",
+ given: &cognito.TokenSet{ExpiresAt: fixedNow.Add(-time.Second)},
+ want: true,
+ },
+ {
+ name: "inside the safety skew",
+ given: &cognito.TokenSet{ExpiresAt: fixedNow.Add(30 * time.Second)},
+ want: true,
+ },
+ {
+ name: "exactly at the safety skew",
+ given: &cognito.TokenSet{ExpiresAt: fixedNow.Add(60 * time.Second)},
+ want: true,
+ },
+ {
+ name: "just beyond the safety skew",
+ given: &cognito.TokenSet{ExpiresAt: fixedNow.Add(61 * time.Second)},
+ want: false,
+ },
+ {
+ name: "an hour of life left",
+ given: &cognito.TokenSet{ExpiresAt: fixedNow.Add(time.Hour)},
+ want: false,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ assert.Equal(t, tt.want, tt.given.IsExpired(fixedNow), "evaluated at %v", fixedNow)
+ })
+ }
+}
+
+func TestTokenSet_Email(t *testing.T) {
+ tests := []struct {
+ name string
+ givenIDToken string
+ want string
+ wantErrContains string
+ }{
+ {
+ name: "blank id_token is rejected",
+ givenIDToken: "",
+ wantErrContains: "not a jwt",
+ },
+ {
+ name: "two-segment token is not a jwt",
+ givenIDToken: "header.payload",
+ wantErrContains: "not a jwt",
+ },
+ {
+ name: "truncated payload does not panic",
+ givenIDToken: "header.!!!not-base64!!!.signature",
+ wantErrContains: "decode",
+ },
+ {
+ name: "payload is not json",
+ givenIDToken: "header." + rawB64("this is not json") + ".signature",
+ wantErrContains: "parse",
+ },
+ {
+ name: "payload is a json array not an object",
+ givenIDToken: "header." + rawB64(`["nope"]`) + ".signature",
+ wantErrContains: "parse",
+ },
+ {
+ name: "no email claim",
+ givenIDToken: "header." + rawB64(`{"sub":"abc"}`) + ".signature",
+ wantErrContains: "no email claim",
+ },
+ {
+ name: "email claim is not a string",
+ givenIDToken: "header." + rawB64(`{"email":42}`) + ".signature",
+ wantErrContains: "no email claim",
+ },
+ {
+ name: "email claim is blank",
+ givenIDToken: "header." + rawB64(`{"email":""}`) + ".signature",
+ wantErrContains: "no email claim",
+ },
+ {
+ name: "padded base64 payload still decodes",
+ givenIDToken: "header." + base64.URLEncoding.EncodeToString([]byte(`{"email":"`+testOperatorEmail+`"}`)) + ".signature",
+ want: testOperatorEmail,
+ },
+ {
+ name: "email is returned",
+ givenIDToken: "header." + rawB64(`{"email":"`+testOperatorEmail+`"}`) + ".signature",
+ want: testOperatorEmail,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ ts := &cognito.TokenSet{IDToken: tt.givenIDToken}
+
+ var (
+ got string
+ err error
+ )
+ require.NotPanics(t, func() { got, err = ts.Email() }, "Email must never panic on a malformed id_token")
+
+ if wegostrings.IsNotEmpty(tt.wantErrContains) {
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), tt.wantErrContains)
+ assert.Empty(t, got)
+ return
+ }
+ require.NoError(t, err)
+ assert.Equal(t, tt.want, got)
+ })
+ }
+}
+
+func TestTokenSet_EmailOnNilReceiver(t *testing.T) {
+ var ts *cognito.TokenSet
+
+ var err error
+ require.NotPanics(t, func() { _, err = ts.Email() })
+ require.Error(t, err)
+}
+
+// rawB64 is unpadded base64url, the encoding JWT segments use.
+func rawB64(s string) string {
+ return base64.RawURLEncoding.EncodeToString([]byte(s))
+}
+
+// mustIDToken builds an unsigned JWT carrying claims. The signature is
+// deliberately junk: Email() must not verify it.
+func mustIDToken(t *testing.T, claims map[string]any) string {
+ t.Helper()
+ payload, err := json.Marshal(claims)
+ require.NoError(t, err)
+ return rawB64(`{"alg":"RS256","typ":"JWT"}`) + "." +
+ base64.RawURLEncoding.EncodeToString(payload) + ".not-a-real-signature"
+}