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..74fdc70
--- /dev/null
+++ b/cognito/callback.go
@@ -0,0 +1,186 @@
+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, expectedState 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(expectedState, 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.
+//
+// expectedState is checked BEFORE anything is delivered, and that ordering is
+// the point. The callback port is predictable and the result channel is
+// single-use, so a handler that published first would let any local process —
+// or a page in the operator's own browser — hit http://127.0.0.1:/?code=x
+// and consume the one delivery, aborting a real sign-in that had not landed
+// yet. Refusing to deliver an uncorrelated request means the flow survives it
+// and the genuine redirect is still accepted.
+//
+// This is a denial-of-sign-in guard, not the CSRF check: Login still compares
+// the state it generated against what came back, which is what stops a code
+// from a different attempt being exchanged. Both run, and the paste-back path
+// relies on the Login-side one alone.
+func callbackHandler(expectedState string, 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")
+
+ // Drop anything that is not this attempt without touching the channel.
+ // Deliberately indistinguishable from any other failure to the caller:
+ // the page says nothing, and the terminal is never told, because a
+ // stray request is not the operator's problem to debug.
+ if !stateMatches(expectedState, query.Get("state")) {
+ w.WriteHeader(http.StatusBadRequest)
+ _, _ = fmt.Fprint(w, failureHTML)
+ return
+ }
+
+ 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..368f6d6
--- /dev/null
+++ b/cognito/callback_internal_test.go
@@ -0,0 +1,292 @@
+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) {
+ const wantedState = "the-state"
+
+ tests := []struct {
+ name string
+ givenQuery string
+ wantStatus int
+ wantCode string
+ wantState string
+ wantErrContains string
+ wantNoDelivery bool
+ }{
+ {
+ name: "authorize error is surfaced with its description",
+ givenQuery: "error=access_denied&error_description=user+said+no&state=" + wantedState,
+ wantStatus: http.StatusBadRequest,
+ wantErrContains: "user said no",
+ },
+ {
+ name: "authorize error without a description still reports",
+ givenQuery: "error=server_error&state=" + wantedState,
+ wantStatus: http.StatusBadRequest,
+ wantErrContains: "server_error",
+ },
+ {
+ name: "missing code is rejected",
+ givenQuery: "state=" + wantedState,
+ wantStatus: http.StatusBadRequest,
+ wantErrContains: "no authorization code",
+ },
+ {
+ name: "code and state are captured",
+ givenQuery: "code=the-code&state=" + wantedState,
+ wantStatus: http.StatusOK,
+ wantCode: "the-code",
+ wantState: wantedState,
+ },
+ // The next three are the denial-of-sign-in guard. Each is a request
+ // that reaches the predictable callback port without belonging to this
+ // attempt, and the assertion that matters is wantNoDelivery: the
+ // single-use channel must still be empty afterwards, so the genuine
+ // redirect can still be accepted.
+ {
+ name: "a code carrying someone else's state is not delivered",
+ givenQuery: "code=blind-code&state=not-this-attempt",
+ wantStatus: http.StatusBadRequest,
+ wantNoDelivery: true,
+ },
+ {
+ name: "a code with no state at all is not delivered",
+ givenQuery: "code=blind-code",
+ wantStatus: http.StatusBadRequest,
+ wantNoDelivery: true,
+ },
+ {
+ name: "an authorize error from another attempt is not delivered",
+ givenQuery: "error=access_denied&state=not-this-attempt",
+ wantStatus: http.StatusBadRequest,
+ wantNoDelivery: true,
+ },
+ }
+
+ 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(wantedState, ch)(rec, req)
+
+ assert.Equal(t, tt.wantStatus, rec.Code)
+ assert.Contains(t, rec.Header().Get("Content-Type"), "text/html")
+
+ if tt.wantNoDelivery {
+ assert.Empty(t, ch, "an uncorrelated request must not consume the single-use delivery")
+ assert.NotContains(t, rec.Body.String(), "blind-code",
+ "the page must not reflect anything from the request")
+ return
+ }
+
+ 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")
+ }
+ })
+ }
+}
+
+// TestCallbackHandler_BlindRequestDoesNotAbortTheRealSignIn is the
+// invalid-then-valid case. Before the state check moved ahead of delivery, the
+// first request here consumed the one delivery the flow gets and the operator's
+// actual redirect arrived to a channel that was already full, so a sign-in
+// could be aborted by anything able to guess the callback port.
+func TestCallbackHandler_BlindRequestDoesNotAbortTheRealSignIn(t *testing.T) {
+ const wantedState = "real-attempt"
+
+ ch := make(chan callbackResult, 1)
+ handler := callbackHandler(wantedState, ch)
+
+ blind := httptest.NewRecorder()
+ handler(blind, httptest.NewRequest(http.MethodGet, "/callback?code=blind&state=guessed", nil))
+
+ assert.Equal(t, http.StatusBadRequest, blind.Code)
+ require.Empty(t, ch, "the blind request must leave the delivery unused")
+
+ real := httptest.NewRecorder()
+ handler(real, httptest.NewRequest(http.MethodGet, "/callback?code=genuine&state="+wantedState, nil))
+
+ assert.Equal(t, http.StatusOK, real.Code)
+ select {
+ case got := <-ch:
+ require.NoError(t, got.err)
+ assert.Equal(t, "genuine", got.code, "the real callback must still be the one delivered")
+ assert.Equal(t, wantedState, got.state)
+ default:
+ t.Fatal("the genuine callback was not delivered after a blind request")
+ }
+}
+
+func TestCallbackHandler_IsSingleUse(t *testing.T) {
+ ch := make(chan callbackResult, 1)
+ handler := callbackHandler("s", 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", "s")
+
+ 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", "xyz")
+ 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..9cbc3b8
--- /dev/null
+++ b/cognito/oauth.go
@@ -0,0 +1,710 @@
+// 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"
+ "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
+
+ schemeHTTP = "http"
+ schemeHTTPS = "https"
+)
+
+// 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 or ReadRedirect is set — a
+ // sign-in whose URL the operator never sees cannot complete — and ignored
+ // otherwise.
+ PromptURL func(url string) error
+
+ // ReadRedirect turns the sign-in into a paste-back exchange: instead of a
+ // loopback listener receiving the redirect, the caller returns the URL the
+ // browser was redirected to and Login reads the code out of it.
+ //
+ // This is the only path that works with NO browser on this machine AND no
+ // way to reach its callback port. Cognito redirects to a loopback URL that
+ // nothing is listening on, the browser shows a connection error, and the
+ // address bar holds ?code=...&state=... for the operator to copy back.
+ //
+ // Cognito has no device authorization grant (RFC 8628): its discovery
+ // document advertises no device_authorization_endpoint, /oauth2/device_
+ // authorization is 404, and the token endpoint answers a device_code grant
+ // with unsupported_grant_type. So this is the substitute, and it needs no
+ // extra Cognito configuration because it is still the authorization code
+ // grant with PKCE, only with the redirect carried by hand.
+ //
+ // Setting it implies NoBrowser: no browser is launched here. PromptURL is
+ // required alongside it; CallbackAddr is not, since nothing binds a port.
+ // The pasted URL carries a single-use authorization code, so it should not
+ // travel through a shared channel; state is still checked, so a code from a
+ // different attempt is rejected.
+ //
+ // LIFETIME CONTRACT. Read this before wiring a reader to it.
+ //
+ // The ctx is the one passed to Login, and honouring it is the caller's
+ // responsibility. Login itself stops waiting when ctx is done, so a
+ // cancelled or timed-out Login always returns — but it CANNOT interrupt a
+ // read already blocked in this function. A reader that ignores ctx (a bare
+ // os.Stdin.Read, or bufio.Reader over it) therefore stays blocked after
+ // Login has returned, holding whatever it reads from.
+ //
+ // Two obligations follow:
+ //
+ // - On cancellation, tear the reader down yourself if you need the read to
+ // stop — close the file or pipe, or use a reader that selects on ctx.
+ // Nothing in this package can do it for you.
+ // - Do NOT start another Login while a previous invocation may still be
+ // blocked in here. Two live invocations reading the same stdin race for
+ // the operator's next line: either may win, so a retry can consume the
+ // paste meant for the other and the losing attempt sees a truncated or
+ // empty read. Wait for the previous reader to yield, or give the retry
+ // its own reader.
+ //
+ // The abandoned invocation is otherwise harmless: Login sends its result
+ // into a buffered channel and exits, so it parks rather than leaks and never
+ // blocks on a receiver that has gone.
+ ReadRedirect func(ctx context.Context) (redirectedURL string, err 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
+ }
+
+ code, callbackState, err := cfg.collectCode(ctx, state, generateChallenge(verifier))
+ 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 requireHTTPS("cognito token url", c.TokenURL)
+}
+
+// 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 err := requireHTTPS("cognito authorize url", c.AuthorizeURL); err != nil {
+ return err
+ }
+ if wegostrings.IsBlank(c.RedirectURI) {
+ return errors.New("cognito redirect uri is not configured")
+ }
+ if err := requireLoopbackRedirect(c.RedirectURI); err != nil {
+ return err
+ }
+ // Only the listener path needs a port; paste-back binds nothing.
+ if c.ReadRedirect == nil && wegostrings.IsBlank(c.CallbackAddr) {
+ return errors.New("local callback address is not configured")
+ }
+ if c.ReadRedirect == nil {
+ if err := requireLoopbackAddr(c.CallbackAddr); err != nil {
+ return err
+ }
+ }
+ if c.ReadRedirect != nil && c.PromptURL == nil {
+ return errors.New("paste-back sign-in needs Config.PromptURL: the operator has to be shown the url they are meant to open")
+ }
+ 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 {
+ // ReadRedirect implies the prompt. A machine that cannot receive the
+ // callback generally cannot open a browser either, and where it could, the
+ // listener path would have worked and been less work for the operator.
+ // Launching a browser and then also asking for a paste is a footgun.
+ if c.NoBrowser || c.ReadRedirect != nil {
+ 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
+}
+
+// collectCode shows the operator the authorize URL and returns the code they
+// came back with, by whichever route this Config selected.
+//
+// The loopback listener is the default. ReadRedirect replaces it with a
+// paste-back exchange for a machine that can neither open a browser nor be
+// reached on its callback port.
+func (c Config) collectCode(ctx context.Context, state, challenge string) (code, gotState string, err error) {
+ if c.ReadRedirect != nil {
+ // No port is bound, so there is nothing to fail early on: show the URL,
+ // then wait on the operator rather than on the network.
+ if err := c.presentAuthorizeURL(c.buildAuthorizeURL(state, challenge)); err != nil {
+ return "", "", err
+ }
+
+ pasted, err := c.readRedirect(ctx)
+ if err != nil {
+ return "", "", err
+ }
+
+ return codeFromRedirect(pasted)
+ }
+
+ // 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(c.CallbackAddr, callbackPath(c.RedirectURI), state)
+ if err != nil {
+ return "", "", err
+ }
+ defer server.shutdown()
+
+ if err := c.presentAuthorizeURL(c.buildAuthorizeURL(state, challenge)); err != nil {
+ return "", "", err
+ }
+
+ return server.wait(ctx, defaultCallbackTimeout)
+}
+
+// codeFromRedirect reads the authorization code and state out of a redirect URL
+// an operator pasted back.
+//
+// The host and scheme are deliberately not checked. The operator is copying
+// from their own address bar, and what binds the code to this attempt is the
+// state check the caller performs, not the shape of the URL. A bare code is
+// refused for the same reason: with no state there is nothing to bind it to.
+func codeFromRedirect(pasted string) (code, state string, err error) {
+ trimmed := strings.TrimSpace(pasted)
+ if wegostrings.IsBlank(trimmed) {
+ return "", "", errors.New("nothing was pasted: copy the whole URL the browser was redirected to, including the ?code=... part")
+ }
+
+ parsed, err := url.Parse(trimmed)
+ if err != nil {
+ return "", "", fmt.Errorf("parse the pasted redirect: %w", err)
+ }
+
+ if wegostrings.IsBlank(parsed.RawQuery) {
+ return "", "", errors.New("the pasted value has no query string: copy the whole URL from the address bar, including the ?code=... part")
+ }
+
+ query := parsed.Query()
+
+ // Cognito reports a refusal in the redirect rather than as a failed
+ // request, so this is where a denied sign-in surfaces.
+ if reported := query.Get("error"); wegostrings.IsNotBlank(reported) {
+ if description := query.Get("error_description"); wegostrings.IsNotBlank(description) {
+ return "", "", fmt.Errorf("the sign-in was refused: %s (%s)", reported, description)
+ }
+
+ return "", "", fmt.Errorf("the sign-in was refused: %s", reported)
+ }
+
+ code = query.Get("code")
+ if wegostrings.IsBlank(code) {
+ return "", "", errors.New("the pasted redirect carries no authorization code")
+ }
+
+ state = query.Get("state")
+ if wegostrings.IsBlank(state) {
+ return "", "", errors.New("the pasted redirect carries no state, so it cannot be tied to this sign-in attempt")
+ }
+
+ return code, state, nil
+}
+
+func (c Config) openBrowserAt(rawURL string) error {
+ if c.OpenBrowser != nil {
+ return c.OpenBrowser(rawURL)
+ }
+ return openBrowser(rawURL)
+}
+
+// ErrInsecureEndpoint is returned for a configured URL or bind address that
+// must not carry OAuth credentials.
+var ErrInsecureEndpoint = errors.New("insecure cognito endpoint")
+
+// requireHTTPS rejects a Cognito endpoint that is not https.
+//
+// Both endpoints carry credentials in the clear if the transport does not
+// protect them: the authorize URL puts the PKCE challenge and state on the
+// wire, and the token endpoint receives the authorization code, the PKCE
+// verifier and the client id on sign-in and the refresh token on renewal. A
+// cleartext http:// value hands all of that to anyone on the path, so it is
+// refused rather than warned about.
+//
+// There is no loopback exemption here, unlike the redirect: these are Cognito's
+// own endpoints on a domain the operator configured, and Cognito serves them
+// over https only. An http:// value is a misconfiguration in every case.
+func requireHTTPS(name, raw string) error {
+ parsed, err := url.Parse(raw)
+ if err != nil {
+ return fmt.Errorf("%w: %s %q cannot be parsed: %w", ErrInsecureEndpoint, name, raw, err)
+ }
+ if parsed.Scheme != schemeHTTPS {
+ return fmt.Errorf(
+ "%w: %s must use https, got %q (the sign-in exchange carries the authorization code, the pkce verifier and the refresh token, so cleartext would expose the whole credential set)",
+ ErrInsecureEndpoint, name, raw)
+ }
+ if wegostrings.IsBlank(parsed.Host) {
+ return fmt.Errorf("%w: %s %q has no host", ErrInsecureEndpoint, name, raw)
+ }
+ if parsed.User != nil {
+ return fmt.Errorf(
+ "%w: %s %q embeds credentials in the url, which would be logged and sent on every request",
+ ErrInsecureEndpoint, name, raw)
+ }
+ return nil
+}
+
+// requireLoopbackRedirect rejects a redirect URI that is not a loopback
+// address.
+//
+// Cleartext http IS allowed here, and only here: RFC 8252 §7.3 has a native app
+// receive the redirect on a loopback interface, where the request never leaves
+// the machine and TLS would buy nothing but a certificate problem. What must
+// hold is that the host is loopback — a redirect to a routable host would send
+// the authorization code across the network to whoever answers it, which is the
+// exact exposure this package's loopback-only design exists to prevent.
+func requireLoopbackRedirect(raw string) error {
+ parsed, err := url.Parse(raw)
+ if err != nil {
+ return fmt.Errorf("%w: cognito redirect uri %q cannot be parsed: %w", ErrInsecureEndpoint, raw, err)
+ }
+ switch parsed.Scheme {
+ case schemeHTTP, schemeHTTPS:
+ default:
+ return fmt.Errorf(
+ "%w: cognito redirect uri %q must be http or https, got scheme %q",
+ ErrInsecureEndpoint, raw, parsed.Scheme)
+ }
+ if !isLoopbackHost(parsed.Hostname()) {
+ return fmt.Errorf(
+ "%w: cognito redirect uri %q must be a loopback address (127.0.0.1, [::1] or localhost) — a routable host would receive the authorization code over the network",
+ ErrInsecureEndpoint, raw)
+ }
+ return nil
+}
+
+// requireLoopbackAddr rejects a callback bind address that is not loopback.
+//
+// Binding 0.0.0.0 or a LAN interface would let anything that can reach the
+// machine deliver a redirect to the single-use callback, so the listener is
+// held to the same loopback rule as the redirect it serves.
+func requireLoopbackAddr(addr string) error {
+ host, port, err := net.SplitHostPort(addr)
+ if err != nil {
+ return fmt.Errorf(
+ "%w: local callback address %q must be host:port, e.g. 127.0.0.1:8110: %w",
+ ErrInsecureEndpoint, addr, err)
+ }
+ if wegostrings.IsBlank(port) {
+ return fmt.Errorf("%w: local callback address %q has no port", ErrInsecureEndpoint, addr)
+ }
+ if !isLoopbackHost(host) {
+ return fmt.Errorf(
+ "%w: local callback address %q must bind a loopback interface (127.0.0.1 or [::1]) — binding 0.0.0.0 or a routable interface would let any host that can reach this machine deliver the sign-in callback",
+ ErrInsecureEndpoint, addr)
+ }
+ return nil
+}
+
+// isLoopbackHost reports whether host is LITERALLY loopback.
+//
+// Only a literal loopback IP counts, plus the name "localhost". A name that
+// merely resolves to 127.0.0.1 does not: resolution can change between this
+// check and the request, so trusting it would make the check decorative. Go
+// itself treats "localhost" as loopback (net/http.isLocalhost, and the URL
+// spec's special-casing), and a poisoned "localhost" costs the attacker nothing
+// they do not already have on a machine they can edit /etc/hosts on.
+func isLoopbackHost(host string) bool {
+ if strings.EqualFold(host, "localhost") {
+ return true
+ }
+ ip := net.ParseIP(strings.Trim(host, "[]"))
+ return ip != nil && ip.IsLoopback()
+}
+
+// readRedirect runs ReadRedirect without letting it pin Login to a blocked
+// reader.
+//
+// ReadRedirect is synchronous and typically reads a line from stdin, which
+// ignores context: os.Stdin.Read stays blocked until the operator types
+// something or the descriptor closes. Called directly, that made a cancelled
+// or timed-out Login unable to return — the caller asked to stop and the
+// process sat there anyway, which is the whole defect.
+//
+// ctx is handed to the callback so a reader that can honour it will, and Login
+// additionally selects on ctx.Done() so it returns even when the reader cannot.
+// Those are different guarantees and both are needed: the select bounds Login,
+// the ctx argument is what lets the read itself be abandoned. Neither can
+// interrupt a read already blocked inside an uncooperative reader — the
+// goroutine survives until it yields, then sends into a buffered channel and
+// exits, so it parks rather than leaks and never blocks on a receiver that has
+// gone. Config.ReadRedirect documents the teardown and no-concurrent-retry
+// obligations that fall to the caller as a result.
+func (c Config) readRedirect(ctx context.Context) (string, error) {
+ type pasteResult struct {
+ url string
+ err error
+ }
+ // Buffered: the goroutine must be able to finish and exit after Login has
+ // already returned on ctx.Done() and stopped receiving.
+ done := make(chan pasteResult, 1)
+
+ go func() {
+ url, err := c.ReadRedirect(ctx)
+ done <- pasteResult{url: url, err: err}
+ }()
+
+ select {
+ case <-ctx.Done():
+ return "", fmt.Errorf("waiting for the pasted redirect: %w", ctx.Err())
+ case result := <-done:
+ if result.err != nil {
+ return "", fmt.Errorf("read the pasted redirect: %w", result.err)
+ }
+ return result.url, nil
+ }
+}
diff --git a/cognito/oauth_internal_test.go b/cognito/oauth_internal_test.go
new file mode 100644
index 0000000..b5815d1
--- /dev/null
+++ b/cognito/oauth_internal_test.go
@@ -0,0 +1,368 @@
+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")
+ })
+ }
+}
+
+// TestCodeFromRedirect covers pulling the authorization code out of the URL an
+// operator pasted back.
+//
+// This is the sign-in path for a machine with no browser AND no way to forward
+// the callback port: Cognito redirects to a loopback URL nothing is listening
+// on, the browser shows a connection error, and the address bar holds the code.
+func TestCodeFromRedirect(t *testing.T) {
+ tests := []struct {
+ name string
+ givenInput string
+ wantCode string
+ wantState string
+ wantErr string
+ }{
+ {
+ name: "a pasted redirect yields its code and state",
+ givenInput: "http://localhost:8100/callback?code=abc123&state=xyz789",
+ wantCode: "abc123",
+ wantState: "xyz789",
+ },
+ {
+ name: "surrounding whitespace from a copy is tolerated",
+ givenInput: " http://localhost:8100/callback?code=abc123&state=xyz789\n",
+ wantCode: "abc123",
+ wantState: "xyz789",
+ },
+ {
+ name: "https and an unexpected host are accepted; the state check is the guard",
+ givenInput: "https://example.test/cb?state=xyz789&code=abc123",
+ wantCode: "abc123",
+ wantState: "xyz789",
+ },
+ {
+ name: "percent-encoded values are decoded",
+ givenInput: "http://localhost:8100/callback?code=a%2Fb%2Bc&state=s%3D1",
+ wantCode: "a/b+c",
+ wantState: "s=1",
+ },
+ {name: "nothing pasted", givenInput: " ", wantErr: "nothing was pasted"},
+ {
+ // The commonest mistake: copying the path but not the query.
+ name: "a url with no query says what is missing",
+ givenInput: "http://localhost:8100/callback",
+ wantErr: "no query",
+ },
+ {
+ name: "a bare code is refused because it carries no state",
+ givenInput: "abc123",
+ wantErr: "no query",
+ },
+ {
+ name: "a denied sign-in reports the provider's error",
+ givenInput: "http://localhost:8100/callback?error=access_denied&error_description=User+denied",
+ wantErr: "access_denied",
+ },
+ {
+ name: "a denied sign-in includes the description when there is one",
+ givenInput: "http://localhost:8100/callback?error=access_denied&error_description=User+denied",
+ wantErr: "User denied",
+ },
+ {
+ name: "a query with state but no code is refused",
+ givenInput: "http://localhost:8100/callback?state=xyz789",
+ wantErr: "no authorization code",
+ },
+ {
+ // Without state there is nothing to bind the code to this attempt.
+ name: "a query with code but no state is refused",
+ givenInput: "http://localhost:8100/callback?code=abc123",
+ wantErr: "no state",
+ },
+ {
+ name: "an unparsable url is reported",
+ givenInput: "http://[::1",
+ wantErr: "parse",
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ code, state, err := codeFromRedirect(tt.givenInput)
+
+ if tt.wantErr != "" {
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), tt.wantErr)
+ assert.Empty(t, code)
+
+ return
+ }
+
+ require.NoError(t, err)
+ assert.Equal(t, tt.wantCode, code)
+ assert.Equal(t, tt.wantState, state)
+ })
+ }
+}
+
+// TestValidateForLogin_PasteBack covers the config rules the paste-back path
+// adds: it needs somewhere to show the URL, and it does NOT need a callback
+// port, since nothing binds one.
+func TestValidateForLogin_PasteBack(t *testing.T) {
+ base := func() Config {
+ return Config{
+ AuthorizeURL: "https://cognito.test/oauth2/authorize",
+ TokenURL: "https://cognito.test/oauth2/token",
+ ClientID: "client",
+ RedirectURI: "http://localhost:8100/callback",
+ }
+ }
+
+ t.Run("paste-back without PromptURL is refused", func(t *testing.T) {
+ cfg := base()
+ cfg.ReadRedirect = func(context.Context) (string, error) { return "", nil }
+
+ require.ErrorContains(t, cfg.validateForLogin(), "PromptURL")
+ })
+
+ t.Run("paste-back needs no callback address", func(t *testing.T) {
+ cfg := base()
+ cfg.ReadRedirect = func(context.Context) (string, error) { return "", nil }
+ cfg.PromptURL = func(string) error { return nil }
+
+ require.NoError(t, cfg.validateForLogin(), "nothing binds a port on this path")
+ })
+
+ t.Run("the listener path still needs a callback address", func(t *testing.T) {
+ require.ErrorContains(t, base().validateForLogin(), "callback address")
+ })
+}
+
+// TestPresentAuthorizeURL_PasteBackImpliesThePrompt pins the coupling: setting
+// ReadRedirect must not leave a browser being launched, or the operator gets a
+// browser AND a paste prompt for the same sign-in.
+func TestPresentAuthorizeURL_PasteBackImpliesThePrompt(t *testing.T) {
+ var prompted, opened string
+
+ cfg := Config{
+ PromptURL: func(u string) error { prompted = u; return nil },
+ OpenBrowser: func(u string) error { opened = u; return nil },
+ ReadRedirect: func(context.Context) (string, error) { return "", nil },
+ }
+
+ require.NoError(t, cfg.presentAuthorizeURL("https://cognito.test/authorize?state=s"))
+
+ assert.Equal(t, "https://cognito.test/authorize?state=s", prompted)
+ assert.Empty(t, opened, "no browser may be launched on the paste-back path")
+}
diff --git a/cognito/oauth_test.go b/cognito/oauth_test.go
new file mode 100644
index 0000000..a0a6002
--- /dev/null
+++ b/cognito/oauth_test.go
@@ -0,0 +1,1085 @@
+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
+ givenTimeout time.Duration
+ 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",
+ },
+ // On the listener path a callback whose state is not this attempt's is
+ // now DROPPED BY THE HANDLER rather than delivered and then rejected by
+ // Login, so the observable failure is the wait timing out instead of a
+ // "state mismatch". That is the point: delivering it would consume the
+ // single-use result and let a blind request to the predictable callback
+ // port abort a real sign-in. See
+ // TestCallbackHandler_BlindRequestDoesNotAbortTheRealSignIn for the
+ // invalid-then-valid proof, and TestLogin_PasteBackRejectsAForeignState
+ // for the paste-back path, where Login's own state check is still the
+ // only guard and still reports a mismatch.
+ {
+ name: "a tampered state never completes the sign-in",
+ givenTimeout: 500 * time.Millisecond,
+ givenBrowser: func(f *fakeBrowser) {
+ f.tamper = func(q url.Values) { q.Set("state", "tampered-state") }
+ },
+ wantErrContains: "waiting for the sign-in callback",
+ },
+ {
+ name: "a callback with no state never completes the sign-in",
+ givenTimeout: 500 * time.Millisecond,
+ givenBrowser: func(f *fakeBrowser) {
+ f.tamper = func(q url.Values) { q.Del("state") }
+ },
+ wantErrContains: "waiting for the sign-in callback",
+ },
+ {
+ 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 := tlsConfig(t, ts, mustFreeAddr(t))
+ cfg.OpenBrowser = browser.open
+ if tt.givenConfig != nil {
+ tt.givenConfig(&cfg)
+ }
+
+ timeout := 10 * time.Second
+ if tt.givenTimeout > 0 {
+ timeout = tt.givenTimeout
+ }
+ ctx, cancel := context.WithTimeout(context.Background(), timeout)
+ 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 := tlsConfig(t, ts, 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 := tlsConfig(t, ts, 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 := tlsConfig(t, ts, 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 := tlsConfig(t, ts, 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 := tlsConfig(t, ts, 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 := tlsConfig(t, ts, "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 := tlsConfig(t, ts, "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 := tlsConfig(t, ts, "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 := tlsConfig(t, ts, "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 },
+ }
+}
+
+// tlsConfig is baseConfig wired to trust ts's self-signed certificate. Config
+// copies the client and forces CheckRedirect, but keeps the Transport, so the
+// test root travels with it.
+func tlsConfig(t *testing.T, ts *tokenServer, callbackAddr string) cognito.Config {
+ t.Helper()
+ cfg := baseConfig(t, ts.URL, callbackAddr)
+ cfg.HTTPClient = ts.Client()
+ return cfg
+}
+
+// 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))
+ }
+
+ // NewTLSServer, not NewServer: Config.validateClient requires https for the
+ // token endpoint, so a plaintext test server would exercise the rejection
+ // path instead of the flow. ts.Client() below trusts this server's cert.
+ ts.Server = httptest.NewTLSServer(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 := tlsConfig(t, ts, 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 := tlsConfig(t, ts, 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)
+ })
+ }
+}
+
+// TestLogin_PasteBack drives a whole sign-in through the paste-back path: no
+// browser is launched, no callback port is bound, and the code arrives as the
+// redirect URL an operator copied out of their address bar.
+func TestLogin_PasteBack(t *testing.T) {
+ ts := newTokenServer(t, nil)
+
+ var shown string
+
+ cfg := cognito.Config{
+ AuthorizeURL: "https://cognito.test/oauth2/authorize",
+ TokenURL: ts.URL,
+ HTTPClient: ts.Client(),
+ ClientID: "test-client-id",
+ RedirectURI: "http://localhost:8100/callback",
+ Scopes: "openid email",
+ AllowedDomain: "@wego.com",
+ Now: func() time.Time { return fixedNow },
+ PromptURL: func(u string) error { shown = u; return nil },
+ }
+ // CallbackAddr is deliberately left empty: nothing binds a port here, and
+ // that is the whole reason this path exists.
+ cfg.ReadRedirect = func(context.Context) (string, error) {
+ return "http://localhost:8100/callback?code=pasted-code&state=" + stateOf(t, shown), nil
+ }
+
+ tokens, err := cognito.Login(context.Background(), cfg)
+ require.NoError(t, err)
+ require.NotNil(t, tokens)
+
+ assert.Contains(t, shown, "code_challenge_method=S256", "the operator must be shown a PKCE authorize url")
+
+ form := ts.lastForm(t)
+ assert.Equal(t, "pasted-code", form.Get("code"), "the pasted code must be what is redeemed")
+ assert.NotEmpty(t, form.Get("code_verifier"), "PKCE must still bind the exchange to this process")
+ assert.Equal(t, "authorization_code", form.Get("grant_type"))
+}
+
+// TestLogin_PasteBackRejectsAForeignState covers the guard that replaces the
+// listener's origin check. Nothing about a pasted URL proves where it came
+// from except the state, so a code from another attempt must be refused.
+func TestLogin_PasteBackRejectsAForeignState(t *testing.T) {
+ ts := newTokenServer(t, nil)
+
+ cfg := cognito.Config{
+ AuthorizeURL: "https://cognito.test/oauth2/authorize",
+ TokenURL: ts.URL,
+ HTTPClient: ts.Client(),
+ ClientID: "test-client-id",
+ RedirectURI: "http://localhost:8100/callback",
+ AllowedDomain: "@wego.com",
+ Now: func() time.Time { return fixedNow },
+ PromptURL: func(string) error { return nil },
+ ReadRedirect: func(context.Context) (string, error) {
+ return "http://localhost:8100/callback?code=pasted-code&state=someone-elses-state", nil
+ },
+ }
+
+ _, err := cognito.Login(context.Background(), cfg)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "state mismatch")
+ assert.Empty(t, ts.forms, "a code with a foreign state must never reach the token endpoint")
+}
+
+// TestLogin_PasteBackReportsAReadFailure covers the operator abandoning the
+// prompt, e.g. ctrl-D at the paste.
+func TestLogin_PasteBackReportsAReadFailure(t *testing.T) {
+ cfg := cognito.Config{
+ AuthorizeURL: "https://cognito.test/oauth2/authorize",
+ TokenURL: "https://unused.test/oauth2/token",
+ ClientID: "test-client-id",
+ RedirectURI: "http://localhost:8100/callback",
+ AllowedDomain: "@wego.com",
+ PromptURL: func(string) error { return nil },
+ ReadRedirect: func(context.Context) (string, error) { return "", errors.New("stdin closed") },
+ }
+
+ _, err := cognito.Login(context.Background(), cfg)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "stdin closed")
+}
+
+// stateOf pulls the state parameter out of an authorize URL.
+func stateOf(t *testing.T, authorizeURL string) string {
+ t.Helper()
+
+ parsed, err := url.Parse(authorizeURL)
+ require.NoError(t, err)
+
+ state := parsed.Query().Get("state")
+ require.NotEmpty(t, state, "the authorize url must carry a state")
+
+ return state
+}
+
+// TestLogin_RejectsInsecureEndpoints covers the transport rules. The sign-in
+// exchange carries the authorization code, the PKCE verifier and the refresh
+// token, so the two Cognito endpoints must be https and the two local ones must
+// be loopback — a routable redirect or bind would put a single-use code on the
+// network for whoever answers.
+//
+// Every case must fail BEFORE any request is made, which is why none of these
+// configs point at a live server.
+func TestLogin_RejectsInsecureEndpoints(t *testing.T) {
+ tests := []struct {
+ name string
+ givenConfig func(*cognito.Config)
+ wantErrContains string
+ }{
+ {
+ name: "cleartext token url",
+ givenConfig: func(c *cognito.Config) { c.TokenURL = "http://cognito.test/oauth2/token" },
+ wantErrContains: "token url must use https",
+ },
+ {
+ name: "cleartext token url on loopback is refused too",
+ givenConfig: func(c *cognito.Config) { c.TokenURL = "http://127.0.0.1:9000/oauth2/token" },
+ wantErrContains: "token url must use https",
+ },
+ {
+ name: "token url embedding credentials",
+ givenConfig: func(c *cognito.Config) { c.TokenURL = "https://user:pw@cognito.test/oauth2/token" },
+ wantErrContains: "embeds credentials in the url",
+ },
+ {
+ name: "cleartext authorize url",
+ givenConfig: func(c *cognito.Config) { c.AuthorizeURL = "http://cognito.test/oauth2/authorize" },
+ wantErrContains: "authorize url must use https",
+ },
+ {
+ name: "routable redirect uri",
+ givenConfig: func(c *cognito.Config) { c.RedirectURI = "http://attacker.test/callback" },
+ wantErrContains: "redirect uri",
+ },
+ {
+ name: "non-http redirect scheme",
+ givenConfig: func(c *cognito.Config) { c.RedirectURI = "ftp://127.0.0.1/callback" },
+ wantErrContains: "must be http or https",
+ },
+ {
+ name: "wildcard callback bind",
+ givenConfig: func(c *cognito.Config) { c.CallbackAddr = "0.0.0.0:8110" },
+ wantErrContains: "must bind a loopback interface",
+ },
+ {
+ name: "routable callback bind",
+ givenConfig: func(c *cognito.Config) { c.CallbackAddr = "192.168.1.20:8110" },
+ wantErrContains: "must bind a loopback interface",
+ },
+ {
+ name: "callback bind without a port",
+ givenConfig: func(c *cognito.Config) { c.CallbackAddr = "127.0.0.1" },
+ wantErrContains: "must be host:port",
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ cfg := baseConfig(t, "https://cognito.test/oauth2/token", "127.0.0.1:8110")
+ cfg.OpenBrowser = func(string) error {
+ t.Fatal("an insecure endpoint must be rejected before a browser is opened")
+ return nil
+ }
+ tt.givenConfig(&cfg)
+
+ got, err := cognito.Login(context.Background(), cfg)
+
+ require.Error(t, err)
+ assert.ErrorIs(t, err, cognito.ErrInsecureEndpoint,
+ "callers need to distinguish a refused endpoint from a failed sign-in")
+ assert.Contains(t, err.Error(), tt.wantErrContains)
+ assert.Nil(t, got)
+ })
+ }
+}
+
+// TestLogin_AcceptsLoopbackHTTPRedirect pins the one cleartext exemption:
+// RFC 8252 has a native app receive the redirect on loopback, where the request
+// never leaves the machine, so requiring TLS there would buy nothing but a
+// certificate problem.
+func TestLogin_AcceptsLoopbackHTTPRedirect(t *testing.T) {
+ for _, host := range []string{"127.0.0.1", "localhost", "[::1]"} {
+ t.Run(host, func(t *testing.T) {
+ cfg := baseConfig(t, "https://cognito.test/oauth2/token", "127.0.0.1:8110")
+ cfg.RedirectURI = "http://" + host + ":8110/callback"
+ cfg.OpenBrowser = func(string) error { return errors.New("stop here") }
+
+ _, err := cognito.Login(context.Background(), cfg)
+
+ require.Error(t, err, "the fake browser stops the flow")
+ assert.NotErrorIs(t, err, cognito.ErrInsecureEndpoint,
+ "a loopback http redirect is the documented native-app callback and must be accepted")
+ })
+ }
+}
+
+// TestRefresh_RejectsInsecureTokenURL covers the renewal path, which sends the
+// refresh token rather than the code and so needs the same transport rule.
+func TestRefresh_RejectsInsecureTokenURL(t *testing.T) {
+ cfg := baseConfig(t, "http://cognito.test/oauth2/token", "127.0.0.1:8110")
+
+ got, err := cognito.Refresh(context.Background(), cfg, testRefreshValue)
+
+ require.Error(t, err)
+ assert.ErrorIs(t, err, cognito.ErrInsecureEndpoint)
+ assert.Nil(t, got)
+}
+
+// TestLogin_PasteBackIsCancellable pins that cancelling Login returns even
+// though the paste-back reader is still blocked.
+//
+// ReadRedirect is synchronous and normally reads stdin, which ignores context:
+// called directly, a cancelled Login could not return until the operator typed
+// something, so the caller asked to stop and the process sat there. Login now
+// waits on ctx alongside the read.
+func TestLogin_PasteBackIsCancellable(t *testing.T) {
+ reading := make(chan struct{})
+ release := make(chan struct{})
+ t.Cleanup(func() { close(release) })
+
+ // Paste-back binds no port, but RedirectURI still has to be a valid loopback
+ // URL — it is what Cognito redirects the operator's browser to.
+ cfg := baseConfig(t, "https://cognito.test/oauth2/token", "127.0.0.1:8110")
+ cfg.CallbackAddr = ""
+ cfg.PromptURL = func(string) error { return nil }
+ cfg.ReadRedirect = func(context.Context) (string, error) {
+ close(reading)
+ <-release // a reader that cannot be interrupted, like os.Stdin.Read
+ return "", errors.New("never reached in this test")
+ }
+
+ ctx, cancel := context.WithCancel(context.Background())
+
+ done := make(chan error, 1)
+ go func() {
+ _, err := cognito.Login(ctx, cfg)
+ done <- err
+ }()
+
+ <-reading // the read is now blocked
+ cancel()
+
+ select {
+ case err := <-done:
+ require.Error(t, err)
+ assert.ErrorIs(t, err, context.Canceled)
+ assert.Contains(t, err.Error(), "waiting for the pasted redirect")
+ case <-time.After(5 * time.Second):
+ t.Fatal("Login did not return after cancellation while the paste-back reader was blocked")
+ }
+}
+
+// TestLogin_PasteBackReaderSeesCancellation pins the other half of the
+// cancellation contract. TestLogin_PasteBackIsCancellable proves Login returns
+// when the reader CANNOT honour ctx; this proves a reader that CAN honour it is
+// given what it needs to — the ctx reaching the callback is the only way a
+// caller can abandon the read itself rather than leaving it parked on stdin.
+func TestLogin_PasteBackReaderSeesCancellation(t *testing.T) {
+ cfg := baseConfig(t, "https://cognito.test/oauth2/token", "127.0.0.1:8110")
+ cfg.CallbackAddr = ""
+ cfg.PromptURL = func(string) error { return nil }
+
+ readerReturned := make(chan error, 1)
+ cfg.ReadRedirect = func(ctx context.Context) (string, error) {
+ // A cooperative reader: selects rather than blocking on the fd.
+ <-ctx.Done()
+ err := ctx.Err()
+ readerReturned <- err
+ return "", err
+ }
+
+ ctx, cancel := context.WithCancel(context.Background())
+ go func() {
+ time.Sleep(50 * time.Millisecond)
+ cancel()
+ }()
+
+ _, err := cognito.Login(ctx, cfg)
+
+ require.Error(t, err)
+ assert.ErrorIs(t, err, context.Canceled)
+
+ select {
+ case readerErr := <-readerReturned:
+ assert.ErrorIs(t, readerErr, context.Canceled,
+ "the callback must receive Login's ctx, not a detached one")
+ case <-time.After(2 * time.Second):
+ t.Fatal("the reader never observed cancellation, so ctx did not reach the callback")
+ }
+}
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..aeb0317
--- /dev/null
+++ b/cognito/storage/keyring.go
@@ -0,0 +1,410 @@
+package storage
+
+import (
+ "crypto/rand"
+ "encoding/hex"
+ "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 under a fresh GENERATION
+// name, and a separate pointer entry names the generation that counts. Save
+// writes its fields 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 generation untouched, so a torn write costs
+// the new session, never the old one.
+//
+// Generations are UNIQUE AND NEVER REUSED, and that is load-bearing twice
+// over. An earlier version alternated between two fixed slots, which broke in
+// two ways review found:
+//
+// - Two concurrent Save calls both read the same live slot, both computed
+// the same "other" slot, and interleaved their field writes into it. Both
+// committed, leaving one slot holding one session's access token beside
+// another's id and refresh tokens. Fresh names make the writers disjoint,
+// so each generation is whole and the later commit simply wins.
+// - The pointer could cycle A->B->A, so a generation that HAD changed under
+// an in-flight Load looked unchanged and Load returned fields it never
+// selected. A name that never repeats makes any change detectable.
+//
+// A hybrid is not merely untidy: the id token carries the operator identity
+// that admin writes are audited against, so pairing it with another session's
+// access or refresh token can attribute a production change to the wrong
+// person.
+//
+// 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 generation. Committing
+ // a token set is exactly one write of this entry.
+ fieldCurrent = "current"
+)
+
+// generationBytes is the entropy in a generation name. A generation is never
+// reused, so this only has to make an accidental collision between two
+// concurrent writers impossible in practice.
+const generationBytes = 12
+
+// loadAttempts is how many times Load will re-read after a concurrent commit
+// moves the pointer out from under it. Each retry needs another process to
+// commit a whole session in the gap, so three is far past what a real overlap
+// produces and still terminates.
+const loadAttempts = 3
+
+// 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"`
+}
+
+// keyringPointer is the commit record: which generation is live, and the one
+// it replaced. Previous is kept only so Delete can reap it, since a keychain
+// cannot be enumerated.
+type keyringPointer struct {
+ Current string `json:"current"`
+ Previous string `json:"previous,omitempty"`
+}
+
+// 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.
+//
+// A token set is read across five keychain entries, so a concurrent commit can
+// move the pointer and clear the slot midway through. That is not an error and
+// must not surface as one: Load re-reads the pointer and starts again on the
+// slot that is now live, which is the ordinary case of a second pay-admin
+// process finishing a login or a refresh while this one reads.
+func (k *keyringStore) Load(namespace string) (*cognito.TokenSet, error) {
+ if err := k.validate(namespace); err != nil {
+ return nil, err
+ }
+
+ for range loadAttempts {
+ pointer, err := k.readPointer(namespace)
+ if err != nil || wegostrings.IsBlank(pointer.Current) {
+ // No pointer means not signed in, which is (nil, nil).
+ return nil, err
+ }
+
+ tokens, err := k.readGeneration(namespace, pointer.Current)
+ if err == nil {
+ return tokens, nil
+ }
+
+ // The read failed. If the pointer has moved since it was chosen, this
+ // generation was superseded and cleared under us, so try the new one.
+ // If it has not, the namespace really is inconsistent and the original
+ // error is the honest one to report.
+ moved, checkErr := k.generationMoved(namespace, pointer.Current)
+ if checkErr != nil {
+ return nil, checkErr
+ }
+
+ if !moved {
+ return nil, err
+ }
+ }
+
+ // Every attempt lost the same race, which needs a sustained run of commits
+ // against one namespace. Report it rather than looping.
+ return nil, fmt.Errorf("the keychain session for %q was replaced %d times while being read", namespace, loadAttempts)
+}
+
+// readPointer returns the commit record, or a zero record when nothing is
+// stored.
+func (k *keyringStore) readPointer(namespace string) (keyringPointer, error) {
+ raw, err := k.backend.Get(k.service, pointerAccount(namespace))
+ if errors.Is(err, keyring.ErrNotFound) {
+ return keyringPointer{}, nil
+ }
+ if err != nil {
+ return keyringPointer{}, fmt.Errorf("read the keychain: %w (is the system keychain unlocked?)", err)
+ }
+
+ var pointer keyringPointer
+ if err := json.Unmarshal([]byte(raw), &pointer); err != nil {
+ // Nothing here writes anything but this record, so the entry was
+ // tampered with, truncated, restored from a backup taken mid-write, or
+ // written by a build that stored a different shape.
+ //
+ // The recovery is always the same and the operator cannot guess it from
+ // a json error, so name it: Delete rewrites the whole namespace, so
+ // signing out and back in fixes any unreadable pointer. Reported rather
+ // than done here, because silently discarding a session store is not
+ // this function's call to make.
+ return keyringPointer{}, fmt.Errorf(
+ "parse the token pointer for %q: %w (the stored session is unreadable - sign out of this environment and sign in again to rewrite it)",
+ namespace, err)
+ }
+
+ if wegostrings.IsBlank(pointer.Current) {
+ return keyringPointer{}, fmt.Errorf(
+ "the token pointer for %q names no generation (the stored session is unreadable - sign out of this environment and sign in again to rewrite it)",
+ namespace)
+ }
+
+ return pointer, nil
+}
+
+// generationMoved reports whether the pointer now names a generation other
+// than the one a read was using.
+//
+// Because a generation name is never reused, this is exact: an unchanged value
+// really means nothing committed in between. Under the old two-slot layout the
+// pointer could cycle back to the value a reader started on, so a changed
+// generation looked unchanged.
+func (k *keyringStore) generationMoved(namespace, generation string) (bool, error) {
+ pointer, err := k.readPointer(namespace)
+ if err != nil {
+ return false, err
+ }
+
+ return pointer.Current != generation, nil
+}
+
+// newGeneration mints a name no other writer will pick.
+func newGeneration() (string, error) {
+ buf := make([]byte, generationBytes)
+ if _, err := rand.Read(buf); err != nil {
+ return "", fmt.Errorf("generate a token generation name: %w", err)
+ }
+
+ return hex.EncodeToString(buf), nil
+}
+
+// readGeneration assembles the token set held under one generation.
+func (k *keyringStore) readGeneration(namespace, generation string) (*cognito.TokenSet, error) {
+ access, err := k.read(namespace, generation, fieldAccess)
+ if err != nil {
+ return nil, err
+ }
+ idToken, err := k.read(namespace, generation, fieldID)
+ if err != nil {
+ return nil, err
+ }
+ refresh, err := k.read(namespace, generation, fieldRefresh)
+ if err != nil {
+ return nil, err
+ }
+ rawMeta, err := k.read(namespace, generation, 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 and safe against a concurrent Save.
+// The fields go under a fresh generation no other writer will pick, 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; if another Save
+// overlaps, the two write disjoint generations and the later commit wins with
+// its own session intact.
+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)
+ }
+
+ // A pointer that cannot be read is not a reason to refuse a sign-in: the
+ // new generation does not collide with anything either way. What is lost is
+ // only the chance to reap the generation being replaced.
+ previous, _ := k.readPointer(namespace)
+
+ generation, err := newGeneration()
+ if err != nil {
+ return err
+ }
+
+ values := map[string]string{
+ fieldAccess: tokens.AccessToken,
+ fieldID: tokens.IDToken,
+ fieldRefresh: tokens.RefreshToken,
+ fieldMeta: string(meta),
+ }
+
+ // Ordering is irrelevant: nothing written here is reachable until the
+ // pointer moves, and no other writer shares this generation. tokenFields is
+ // used rather than ranging the map so the sequence is deterministic, which
+ // keeps failures reproducible.
+ for _, field := range tokenFields {
+ if err := k.backend.Set(k.service, fieldAccount(namespace, generation, field), values[field]); err != nil {
+ return fmt.Errorf("write %s to the keychain: %w (is the system keychain unlocked?)", field, err)
+ }
+ }
+
+ committed, err := json.Marshal(keyringPointer{Current: generation, Previous: previous.Current})
+ if err != nil {
+ return fmt.Errorf("encode the token pointer: %w", err)
+ }
+
+ // The commit.
+ if err := k.backend.Set(k.service, pointerAccount(namespace), string(committed)); err != nil {
+ return fmt.Errorf("commit the session to the keychain: %w (is the system keychain unlocked?)", err)
+ }
+
+ // The generation before the one just replaced is now unreachable by any
+ // reader that started before this commit, so it is safe to clear. The one
+ // directly replaced is deliberately KEPT until the next Save, because a
+ // Load may still be part way through reading it; Load's retry handles the
+ // rest. Failure here is not the caller's problem: the session is committed.
+ k.clearGeneration(namespace, previous.Previous)
+
+ return nil
+}
+
+// clearGeneration removes one generation's field entries, best effort. Callers
+// use it for a generation nothing can still be reading.
+func (k *keyringStore) clearGeneration(namespace, generation string) {
+ if wegostrings.IsBlank(generation) {
+ return
+ }
+
+ for _, field := range tokenFields {
+ _ = k.backend.Delete(k.service, fieldAccount(namespace, generation, 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 is read first so the generations it names can be cleared, then
+// removed BEFORE them: once it is gone the operator is signed out even if a
+// later delete fails.
+//
+// A keychain cannot be enumerated, so only the generations the pointer knows
+// about can be reaped. A generation orphaned by a crash between its field
+// writes and its commit, or by two Saves overlapping, is therefore not
+// reachable here. It holds a superseded token set that no code path will ever
+// return, and Cognito refresh tokens expire, so it decays rather than
+// accumulating indefinitely. Eliminating that window needs process-shared
+// locking, which is a larger change than this seam warrants.
+func (k *keyringStore) Delete(namespace string) error {
+ if err := k.validate(namespace); err != nil {
+ return err
+ }
+
+ // A pointer that cannot be parsed must not block signing out; the entry is
+ // removed below regardless.
+ pointer, _ := k.readPointer(namespace)
+
+ err := k.backend.Delete(k.service, pointerAccount(namespace))
+ if err != nil && !errors.Is(err, keyring.ErrNotFound) {
+ return fmt.Errorf("delete the token pointer from the keychain: %w (is the system keychain unlocked?)", err)
+ }
+
+ for _, generation := range []string{pointer.Current, pointer.Previous} {
+ k.clearGeneration(namespace, generation)
+ }
+
+ 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, generation, field string) (string, error) {
+ value, err := k.backend.Get(k.service, fieldAccount(namespace, generation, 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 generation.
+func fieldAccount(namespace, generation, field string) string {
+ return namespace + "/" + generation + "/" + field
+}
diff --git a/cognito/storage/keyring_internal_test.go b/cognito/storage/keyring_internal_test.go
new file mode 100644
index 0000000..55b8bb6
--- /dev/null
+++ b/cognito/storage/keyring_internal_test.go
@@ -0,0 +1,1009 @@
+package storage
+
+import (
+ "encoding/json"
+ "errors"
+ "fmt"
+ "strings"
+ "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 generation reads as not logged in",
+ givenSetup: func(f *fakeKeyring) {
+ f.putField(testGeneration, fieldAccess, "access-value")
+ f.putField(testGeneration, 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")
+
+ generation := backend.liveGeneration(t)
+ for _, field := range tokenFields {
+ assert.Contains(t, backend.writes, testNamespace+"/"+generation+"/"+field,
+ "each field keeps its own entry, to stay under the 4 KiB cap")
+ }
+}
+
+// TestKeyringStore_SaveNeverReusesAGeneration covers the property both
+// round-4 findings turn on. A name that can come round again lets two writers
+// collide and lets a pointer cycle back to a value a reader started on.
+func TestKeyringStore_SaveNeverReusesAGeneration(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+
+ seen := map[string]bool{}
+
+ for i := range 6 {
+ tokens := sampleTokens()
+ tokens.AccessToken = fmt.Sprintf("access-%d", i)
+ require.NoError(t, store.Save(testNamespace, tokens))
+
+ generation := backend.liveGeneration(t)
+ assert.False(t, seen[generation], "generation %q was reused", generation)
+ seen[generation] = true
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ assert.Equal(t, tokens.AccessToken, got.AccessToken, "the newest session must be the live one")
+ }
+
+ assert.Len(t, seen, 6)
+}
+
+// TestKeyringStore_SaveReapsTheGenerationBeforeLast covers cleanup. The
+// generation directly replaced is kept, because a Load may still be reading
+// it; the one before that cannot have a reader who started after its
+// replacement committed, so it goes.
+func TestKeyringStore_SaveReapsTheGenerationBeforeLast(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+
+ require.NoError(t, store.Save(testNamespace, session("first")))
+ first := backend.liveGeneration(t)
+
+ require.NoError(t, store.Save(testNamespace, session("second")))
+ second := backend.liveGeneration(t)
+
+ require.NoError(t, store.Save(testNamespace, session("third")))
+
+ assert.Empty(t, backend.entries[testService+"|"+testNamespace+"/"+first+"/"+fieldAccess],
+ "the generation before last must be cleared")
+ assert.NotEmpty(t, backend.entries[testService+"|"+testNamespace+"/"+second+"/"+fieldAccess],
+ "the generation just replaced is kept for an in-flight Load")
+}
+
+// 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_LoadRejectsAnUnreadablePointer covers a pointer entry this
+// code did not write. Guessing what it meant would be worse than refusing.
+func TestKeyringStore_LoadRejectsAnUnreadablePointer(t *testing.T) {
+ tests := []struct {
+ name string
+ givenValue string
+ wantErr string
+ }{
+ {name: "not json", givenValue: "somewhere-else", wantErr: "parse the token pointer"},
+ {name: "json naming no generation", givenValue: `{"current":""}`, wantErr: "names no generation"},
+ // A bare slot name is what a pre-release build of this package wrote,
+ // before generations replaced two reusable slots. That shape never
+ // shipped, so no released version can produce it and there is nothing
+ // to migrate — but a keychain carrying one still has to fail readably
+ // rather than with a raw json error.
+ {name: "a bare slot name from a pre-release layout", givenValue: "a", wantErr: "parse the token pointer"},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ backend := newFakeKeyring()
+ backend.put(fieldCurrent, tt.givenValue)
+ store := &keyringStore{service: testService, backend: backend}
+
+ _, err := store.Load(testNamespace)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), tt.wantErr)
+ assert.Contains(t, err.Error(), "sign out of this environment and sign in again",
+ "an unreadable pointer must name the recovery, which is not guessable from a json error")
+ })
+ }
+}
+
+// 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("gen-current", fieldAccess, "access-value")
+ backend.putField("gen-previous", fieldRefresh, "superseded-refresh")
+ backend.putPointer("gen-current", "gen-previous")
+ 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 generation the pointer knows")
+}
+
+// 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_SaveStillSignsInWhenThePointerCannotBeRead covers a
+// keychain that cannot be read at all.
+//
+// Under the old two-slot layout Save HAD to read the pointer, because it wrote
+// to whichever slot the live one was not using; an unreadable pointer
+// therefore had to fail rather than risk overwriting the live session. A fresh
+// generation collides with nothing, so the read is now only an optimisation
+// for reaping, and a sign-in can succeed without it.
+func TestKeyringStore_SaveStillSignsInWhenThePointerCannotBeRead(t *testing.T) {
+ backend := newFakeKeyring()
+ backend.getErr = errors.New("keychain is locked")
+ store := &keyringStore{service: testService, backend: backend}
+
+ require.NoError(t, store.Save(testNamespace, sampleTokens()),
+ "an unreadable pointer must not block establishing a session")
+
+ backend.getErr = nil
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ assert.Equal(t, sampleTokens(), got)
+}
+
+// TestKeyringStore_SaveOverAnUnreadablePointerStillSignsIn covers a pointer
+// holding a value this code never wrote. 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_SaveOverAnUnreadablePointerStillSignsIn(t *testing.T) {
+ backend := newFakeKeyring()
+ backend.put(fieldCurrent, "somewhere-else")
+ store := &keyringStore{service: testService, backend: backend}
+
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+
+ 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(), "token pointer")
+}
+
+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
+ // onGet fires once, on the next Get, before the read happens.
+ onGet func(account string)
+ // onSet fires on every Set, before the write happens, so a test can drive
+ // a second writer into the middle of one Save.
+ onSet func(account string)
+}
+
+// clearOnSet disarms the write hook. A hook that drives a nested Save MUST
+// call this before doing so: the nested writes re-enter the hook, and guarding
+// with sync.Once instead deadlocks, because Once.Do cannot be re-entered.
+func (f *fakeKeyring) clearOnSet() {
+ f.mu.Lock()
+ defer f.mu.Unlock()
+
+ f.onSet = nil
+}
+
+// clearOnGet disarms the read hook, for the same reason as clearOnSet.
+func (f *fakeKeyring) clearOnGet() {
+ f.mu.Lock()
+ defer f.mu.Unlock()
+
+ f.onGet = nil
+}
+
+// takeOnSet returns the write hook, if any.
+func (f *fakeKeyring) takeOnSet() func(string) {
+ f.mu.Lock()
+ defer f.mu.Unlock()
+
+ return f.onSet
+}
+
+// takeOnGet returns the hook, if any. It does NOT clear it: the hook decides
+// which account it wants to fire on, so consuming it on the first unrelated
+// read would silently disarm the test.
+func (f *fakeKeyring) takeOnGet() func(string) {
+ f.mu.Lock()
+ defer f.mu.Unlock()
+
+ return f.onGet
+}
+
+func newFakeKeyring() *fakeKeyring {
+ return &fakeKeyring{entries: make(map[string]string), failSetAfter: -1}
+}
+
+func (f *fakeKeyring) Set(service, user, password string) error {
+ // Outside the lock, like onGet, so a hook can run a whole Save.
+ if hook := f.takeOnSet(); hook != nil {
+ hook(user)
+ }
+
+ 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) {
+ // onGet runs OUTSIDE the lock and before the read, so a hook can drive a
+ // whole Save through this same backend without deadlocking.
+ if hook := f.takeOnGet(); hook != nil {
+ hook(user)
+ }
+
+ 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
+}
+
+// testGeneration is a fixed generation name for seeding, standing in for the
+// random one Save would mint.
+const testGeneration = "gen0"
+
+// putCommitted seeds a committed session under testGeneration, 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(testGeneration, field, value)
+ }
+ f.putPointer(testGeneration, "")
+}
+
+// putPointer seeds the commit record directly.
+func (f *fakeKeyring) putPointer(current, previous string) {
+ raw, err := json.Marshal(keyringPointer{Current: current, Previous: previous})
+ if err != nil {
+ panic(err)
+ }
+ f.put(fieldCurrent, string(raw))
+}
+
+// liveGeneration reads the committed generation name straight out of the fake.
+func (f *fakeKeyring) liveGeneration(t *testing.T) string {
+ t.Helper()
+
+ raw := f.value(fieldCurrent)
+ require.NotEmpty(t, raw, "no pointer entry was written")
+
+ var pointer keyringPointer
+ require.NoError(t, json.Unmarshal([]byte(raw), &pointer))
+
+ return pointer.Current
+}
+
+// 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]
+}
+
+// TestKeyringStore_LoadSurvivesAConcurrentCommit is the regression for a race
+// the two-slot commit introduced.
+//
+// Load reads the pointer, then reads that slot's four fields one at a time.
+// Between those reads another pay-admin process can finish a login or a
+// refresh, commit the replacement slot and clear the one this Load already
+// chose, so the read failed with "secret not found in keyring" even though a
+// perfectly good session existed the whole time.
+//
+// Two overlapping processes is ordinary: a shell running a command while a
+// refresh fires, or two terminals against the same namespace.
+func TestKeyringStore_LoadSurvivesAConcurrentCommit(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+
+ replacement := &cognito.TokenSet{
+ AccessToken: "refreshed-access",
+ IDToken: "refreshed-id",
+ RefreshToken: "refreshed-refresh",
+ ExpiresAt: testExpiry.Add(time.Hour),
+ }
+
+ // Commit the replacement the moment this Load starts reading fields, which
+ // is exactly the window the race lives in.
+ backend.onGet = func(account string) {
+ if account == testNamespace+"/"+fieldCurrent {
+ return // the pointer read itself; let it through
+ }
+
+ backend.clearOnGet()
+ require.NoError(t, store.Save(testNamespace, replacement))
+ }
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err, "a concurrent commit must not fail a Load that was already in flight")
+ require.NotNil(t, got)
+
+ // Either session is a correct answer; a mixture is not.
+ assert.Contains(t, []string{sampleTokens().AccessToken, replacement.AccessToken}, got.AccessToken)
+ if got.AccessToken == replacement.AccessToken {
+ assert.Equal(t, replacement.RefreshToken, got.RefreshToken, "the two sessions must not mix")
+ assert.Equal(t, replacement.IDToken, got.IDToken)
+ } else {
+ assert.Equal(t, sampleTokens().RefreshToken, got.RefreshToken, "the two sessions must not mix")
+ }
+}
+
+// TestKeyringStore_LoadGivesUpOnEndlessCommits covers the retry bound. A
+// namespace being rewritten faster than it can be read is not something a
+// retry can fix, so Load must report it rather than spin.
+func TestKeyringStore_LoadGivesUpOnEndlessCommits(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+
+ // Commit a fresh session on every field read, so the pointer has always
+ // moved by the time the retry check looks.
+ backend.onGet = func(account string) {
+ if account == testNamespace+"/"+fieldCurrent {
+ return
+ }
+
+ next := sampleTokens()
+ next.AccessToken = "rewritten"
+ _ = store.Save(testNamespace, next)
+ }
+
+ _, err := store.Load(testNamespace)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "replaced")
+}
+
+// TestKeyringStore_LoadReportsARealInconsistency pins the other side of the
+// retry: when the pointer has NOT moved, a missing field is a genuinely broken
+// namespace and must be reported, not retried into a generic timeout.
+func TestKeyringStore_LoadReportsARealInconsistency(t *testing.T) {
+ backend := newFakeKeyring()
+ backend.putCommitted(map[string]string{fieldAccess: "access-value"})
+ store := &keyringStore{service: testService, backend: backend}
+
+ _, err := store.Load(testNamespace)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "id", "the missing field must be named")
+ assert.NotContains(t, err.Error(), "replaced")
+}
+
+// sessionIsWhole asserts a loaded token set came entirely from one session.
+// A set pairing one login's access token with another's id and refresh tokens
+// is the specific corruption these tests exist to prevent: the id token
+// supplies the operator identity that admin writes are audited against, so a
+// mismatched pair can attribute a production change to the wrong person.
+func sessionIsWhole(t *testing.T, got *cognito.TokenSet, sessions map[string]*cognito.TokenSet) {
+ t.Helper()
+ require.NotNil(t, got)
+
+ for name, want := range sessions {
+ if got.AccessToken != want.AccessToken {
+ continue
+ }
+
+ assert.Equal(t, want.IDToken, got.IDToken, "id token must come from the same session as the access token (%s)", name)
+ assert.Equal(t, want.RefreshToken, got.RefreshToken, "refresh token must come from the same session as the access token (%s)", name)
+ assert.Equal(t, want.ExpiresAt, got.ExpiresAt, "expiry must come from the same session as the access token (%s)", name)
+
+ return
+ }
+
+ t.Fatalf("loaded access token %q belongs to no known session", got.AccessToken)
+}
+
+func session(tag string) *cognito.TokenSet {
+ return &cognito.TokenSet{
+ AccessToken: tag + "-access",
+ IDToken: tag + "-id",
+ RefreshToken: tag + "-refresh",
+ ExpiresAt: testExpiry,
+ }
+}
+
+// TestKeyringStore_ConcurrentSavesNeverCommitAHybrid is the first regression
+// named in review round 4.
+//
+// Two pay-admin processes overlapping a login and a refresh both used to read
+// the same live slot, both compute the same "other" slot, and interleave their
+// field writes into it. Both then committed, leaving one slot holding one
+// session's access token beside another's id and refresh tokens.
+func TestKeyringStore_ConcurrentSavesNeverCommitAHybrid(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+
+ seed, first, second := session("seed"), session("first"), session("second")
+ require.NoError(t, store.Save(testNamespace, seed))
+
+ // The second writer runs to completion inside the first writer's very first
+ // field write, which is the schedule that produced the hybrid.
+ // Fire once the first writer has already stored its access token but
+ // before its id token, so the second writer's whole session lands in
+ // between. Firing earlier lets the first writer simply overwrite
+ // everything, which is consistent and proves nothing.
+ backend.onSet = func(account string) {
+ if !strings.HasSuffix(account, "/"+fieldID) {
+ return
+ }
+
+ backend.clearOnSet()
+ require.NoError(t, store.Save(testNamespace, second))
+ }
+
+ require.NoError(t, store.Save(testNamespace, first))
+
+ backend.onSet = nil
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ sessionIsWhole(t, got, map[string]*cognito.TokenSet{"seed": seed, "first": first, "second": second})
+}
+
+// TestKeyringStore_LoadIsNotFooledByPointerReuse is the second regression named
+// in review round 4, the A->B->A schedule.
+//
+// Load re-reads the pointer after a failed field read and retried only when it
+// had changed. With two reusable slots the pointer could cycle back to the one
+// Load started on, so a changed generation looked unchanged and Load returned
+// fields belonging to a session it never selected.
+func TestKeyringStore_LoadIsNotFooledByPointerReuse(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+
+ seed, first, second := session("seed"), session("first"), session("second")
+ require.NoError(t, store.Save(testNamespace, seed))
+
+ // Two commits land after Load has chosen its generation and read part of
+ // it. Under slot reuse the second commit lands back on the first slot.
+ // Fire after Load has read the access token but before the rest, so the
+ // commits land inside one read of one generation.
+ backend.onGet = func(account string) {
+ if !strings.HasSuffix(account, "/"+fieldID) {
+ return
+ }
+
+ backend.clearOnGet()
+ require.NoError(t, store.Save(testNamespace, first))
+ require.NoError(t, store.Save(testNamespace, second))
+ }
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ sessionIsWhole(t, got, map[string]*cognito.TokenSet{"seed": seed, "first": first, "second": second})
+}
+
+// TestKeyringStore_LoadReportsAPointerFailureDuringRetry covers the branch
+// where the pointer read that decides retry-or-report itself fails. Reporting
+// that beats retrying blindly or claiming the namespace is inconsistent.
+func TestKeyringStore_LoadReportsAPointerFailureDuringRetry(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+ require.NoError(t, store.Save(testNamespace, sampleTokens()))
+
+ // Break the field read, then break the keychain outright so the retry
+ // check cannot establish whether the generation moved.
+ generation := backend.liveGeneration(t)
+ delete(backend.entries, testService+"|"+testNamespace+"/"+generation+"/"+fieldID)
+
+ backend.onGet = func(account string) {
+ if !strings.HasSuffix(account, "/"+fieldID) {
+ return
+ }
+
+ backend.clearOnGet()
+ backend.mu.Lock()
+ backend.getErr = errors.New("keychain is locked")
+ backend.mu.Unlock()
+ }
+
+ _, err := store.Load(testNamespace)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "keychain")
+}
+
+// TestKeyringStore_SaveReportsACommitEncodingFailure is not reachable through
+// the public API -- keyringPointer always marshals -- so the pointer encoder is
+// exercised directly to prove the record round-trips exactly.
+func TestKeyringStore_PointerRoundTrips(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+
+ require.NoError(t, store.Save(testNamespace, session("one")))
+ first := backend.liveGeneration(t)
+
+ require.NoError(t, store.Save(testNamespace, session("two")))
+
+ pointer, err := store.readPointer(testNamespace)
+ require.NoError(t, err)
+ assert.Equal(t, backend.liveGeneration(t), pointer.Current)
+ assert.Equal(t, first, pointer.Previous, "the record must remember what it replaced, so Delete can reap it")
+}
+
+// TestKeyringStore_SaveReportsAFieldWriteFailure covers the write loop's error
+// path with the generation layout, and that a failure leaves the live session
+// untouched.
+func TestKeyringStore_SaveReportsAFieldWriteFailure(t *testing.T) {
+ backend := newFakeKeyring()
+ store := &keyringStore{service: testService, backend: backend}
+ require.NoError(t, store.Save(testNamespace, session("live")))
+
+ backend.okWrites = 0
+ backend.failSetAfter = 0
+
+ err := store.Save(testNamespace, session("doomed"))
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "keychain")
+
+ backend.failSetAfter = -1
+
+ got, err := store.Load(testNamespace)
+ require.NoError(t, err)
+ assert.Equal(t, session("live"), got, "a failed write must not disturb the live session")
+}
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"
+}