Make login OAuth state single-use (#193)

Co-authored-by: orchestrator-agent <20+orchestrator-agent@code.cubecraftcreations.com>

internal/auth/auth.go

@@ -79,6 +79,7 @@ type Auth struct {
 	provider *oidc.Provider
 	client   *http.Client
 	cookie   loginCookie
+	pending  *PendingStates
 	sessions Sessions
 	allowed  []string
 	// endSessionURL is the provider's end_session_endpoint; logout ends
@@ -157,6 +158,7 @@ func New(ctx context.Context, cfg Config, sessions Sessions, logger *slog.Logger
 		client:        client,
 		allowed:       slices.Clone(cfg.AllowedGroups),
 		cookie:        newLoginCookie(cfg.SecretKey),
+		pending:       NewPendingStates(loginCookieTTL),
 		sessions:      sessions,
 		endSessionURL: discovered.EndSessionURI,
 		logger:        logger,
@@ -250,7 +252,13 @@ func (a *Auth) login(w http.ResponseWriter, r *http.Request) {
 		Verifier: oauth2.GenerateVerifier(),
 		Return:   returnTo,
 	}
-	if err := a.cookie.set(w, state, time.Now()); err != nil {
+	now := time.Now()
+	if err := a.pending.Add(state.State, now); err != nil {
+		a.logger.WarnContext(r.Context(), "login refused", "error", err)
+		http.Error(w, "too many sign-ins in progress; try again later", http.StatusServiceUnavailable)
+		return
+	}
+	if err := a.cookie.set(w, state, now); err != nil {
 		a.logger.ErrorContext(r.Context(), "set login cookie", "error", err)
 		http.Error(w, "internal error", http.StatusInternalServerError)
 		return
@@ -262,7 +270,7 @@ func (a *Auth) login(w http.ResponseWriter, r *http.Request) {
 }
 
 // callback completes the flow: it checks state against the login cookie,
-// exchanges the code, verifies the ID token and nonce, then starts a session
+// consumes it, exchanges the code, verifies the ID token and nonce, then starts a session
 // if the account is in an allowed group. Membership is checked on every
 // login, so an account removed from the group is refused at its next one.
 func (a *Auth) callback(w http.ResponseWriter, r *http.Request) {
@@ -283,6 +291,13 @@ func (a *Auth) callback(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, "sign-in state mismatch; start again", http.StatusBadRequest)
 		return
 	}
+	// A copy of the cookie verifies until it expires, so the state is
+	// single-use only because it is consumed here.
+	if !a.pending.Consume(state.State) {
+		a.logger.WarnContext(ctx, "login callback state already used")
+		http.Error(w, "sign-in state mismatch; start again", http.StatusBadRequest)
+		return
+	}
 	if q.Get("error") != "" {
 		// error and error_description are provider-supplied text; only the
 		// fixed category is logged.

internal/auth/auth_test.go

@@ -497,6 +497,59 @@ func TestCallbackDefaultReturnPath(t *testing.T) {
 	}
 }
 
+// A saved copy of the login cookie still verifies until it expires, so a
+// replay with a second valid code is refused only because the server
+// consumed the state on first use.
+func TestCallbackRefusesReplayedLoginState(t *testing.T) {
+	t.Parallel()
+	h := newHarness(t)
+	l := h.login(t, LoginPath)
+	idToken := h.provider.signRS256(h.provider.claims(l.query.Get("nonce")))
+
+	if resp := h.callback(l, l.query.Get("state"), idToken); resp.StatusCode != http.StatusSeeOther {
+		t.Fatalf("first callback status = %d, want 303", resp.StatusCode)
+	}
+	resp := h.callback(l, l.query.Get("state"), idToken)
+	if resp.StatusCode != http.StatusBadRequest {
+		t.Errorf("replayed callback status = %d, want 400", resp.StatusCode)
+	}
+	if n := len(h.sessions.started()); n != 1 {
+		t.Errorf("%d sessions started, want 1", n)
+	}
+	assertLoginCookieCleared(t, resp)
+}
+
+// Pending login states live in memory, so a login started before a restart
+// fails closed after it even though its cookie still verifies.
+func TestCallbackRefusesLoginStartedBeforeRestart(t *testing.T) {
+	t.Parallel()
+	h := newHarness(t)
+	l := h.login(t, LoginPath)
+
+	restarted, err := New(context.Background(), Config{
+		Issuer:        h.provider.issuer(),
+		ClientID:      testClientID,
+		ClientSecret:  testClientSecret,
+		PublicURL:     testPublicURL,
+		SecretKey:     []byte("server-secret-key"),
+		AllowedGroups: testAllowedGroups,
+	}, h.sessions, slog.New(slog.DiscardHandler))
+	if err != nil {
+		t.Fatal(err)
+	}
+	r := chi.NewRouter()
+	restarted.Mount(r)
+	h.router = r
+
+	resp := h.callback(l, l.query.Get("state"), h.provider.signRS256(h.provider.claims(l.query.Get("nonce"))))
+	if resp.StatusCode != http.StatusBadRequest {
+		t.Errorf("callback after restart status = %d, want 400", resp.StatusCode)
+	}
+	if n := len(h.sessions.started()); n != 0 {
+		t.Errorf("%d sessions started, want none", n)
+	}
+}
+
 func TestCallbackRejected(t *testing.T) {
 	t.Parallel()
 
@@ -819,6 +872,44 @@ func TestLoginCookieExpires(t *testing.T) {
 	}
 }
 
+func TestPendingStates(t *testing.T) {
+	t.Parallel()
+	p := NewPendingStates(time.Minute)
+	now := time.Now()
+	if err := p.Add("first", now); err != nil {
+		t.Fatal(err)
+	}
+	if !p.Consume("first") {
+		t.Error("Consume of a pending state = false")
+	}
+	if p.Consume("first") {
+		t.Error("second Consume of a state = true")
+	}
+	if p.Consume("never-added") {
+		t.Error("Consume of an unknown state = true")
+	}
+
+	for i := range maxPendingStates - 1 {
+		if err := p.Add(fmt.Sprint("s", i), now); err != nil {
+			t.Fatalf("Add %d: %v", i, err)
+		}
+	}
+	// "first" was consumed but has not expired, so it still counts.
+	if err := p.Add("over", now); !errors.Is(err, ErrTooManyPending) {
+		t.Errorf("Add past the bound = %v, want ErrTooManyPending", err)
+	}
+	// Once the earlier states expire, they are forgotten and a new one fits.
+	if err := p.Add("later", now.Add(time.Minute)); err != nil {
+		t.Errorf("Add after expiry: %v", err)
+	}
+	if p.Consume("s0") {
+		t.Error("an expired state was still pending")
+	}
+	if !p.Consume("later") {
+		t.Error("the state added after expiry is not pending")
+	}
+}
+
 func TestLoginRejectsOffOriginReturn(t *testing.T) {
 	t.Parallel()
 	h := newHarness(t)

internal/auth/cookie.go

@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -153,6 +154,73 @@ func (c FlowCookie) Read(r *http.Request, v flowValue, now time.Time) error {
 	return nil
 }
 
+// maxPendingStates bounds each flow's pending states. The login route is
+// unauthenticated, so without a bound anyone could grow the set without
+// limit; one household never has close to this many flows in progress.
+const maxPendingStates = 10000
+
+// ErrTooManyPending means a flow's pending states are at maxPendingStates,
+// so a new flow is refused until older ones expire.
+var ErrTooManyPending = errors.New("too many sign-in flows in progress")
+
+// PendingStates holds the state of every redirect flow started and not yet
+// completed, until its flow cookie expires. A copy of a flow cookie verifies
+// until it expires, so a flow is single-use only because its callback
+// consumes the state here. The states live in memory — MealQuest runs as one
+// process — so a flow started before a restart fails closed after it.
+type PendingStates struct {
+	ttl    time.Duration
+	mu     sync.Mutex
+	states map[string]time.Time
+	// order is every added state in the order added. Every state is pending
+	// for the same ttl, so it is also expiry order and Add forgets expired
+	// states from its front.
+	order []pendingState
+}
+
+type pendingState struct {
+	state   string
+	expires time.Time
+}
+
+// NewPendingStates returns an empty set whose states are pending for ttl,
+// the flow cookie's lifetime.
+func NewPendingStates(ttl time.Duration) *PendingStates {
+	return &PendingStates{ttl: ttl, states: map[string]time.Time{}}
+}
+
+// Add records state, started at now, as pending and forgets expired ones.
+// It fails with ErrTooManyPending when maxPendingStates are pending.
+func (p *PendingStates) Add(state string, now time.Time) error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	expired := 0
+	for expired < len(p.order) && !now.Before(p.order[expired].expires) {
+		delete(p.states, p.order[expired].state)
+		expired++
+	}
+	p.order = p.order[expired:]
+	// The bound counts every unexpired state, consumed or not: a consumed
+	// state leaves states at once but order only when it expires.
+	if len(p.order) >= maxPendingStates {
+		return ErrTooManyPending
+	}
+	expires := now.Add(p.ttl)
+	p.states[state] = expires
+	p.order = append(p.order, pendingState{state: state, expires: expires})
+	return nil
+}
+
+// Consume removes state from the pending states, reporting whether it was
+// there.
+func (p *PendingStates) Consume(state string) bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	_, ok := p.states[state]
+	delete(p.states, state)
+	return ok
+}
+
 // returnPath validates the optional post-login return path. Empty means "/".
 // Only same-origin absolute paths are allowed: no scheme, host, protocol-
 // relative "//" prefix, backslashes or control characters.

internal/gcal/connect.go

@@ -137,10 +137,9 @@ type Service struct {
 	account sync.Mutex
 
 	// pending holds the state of every connect started and not yet
-	// completed, with when it expires. Complete consumes a state once, so a
-	// replayed cookie finds it gone.
-	pendingMu sync.Mutex
-	pending   map[string]time.Time
+	// completed. Complete consumes a state once, so a replayed cookie finds
+	// it gone.
+	pending *auth.PendingStates
 }
 
 // NewService returns the service over st, sealing refresh tokens with box.
@@ -175,7 +174,7 @@ func NewService(st *store.Store, box *secretbox.Box, cfg Config) (*Service, erro
 		cookie:      auth.NewFlowCookie(connectCookieName, connectCookieKeyLabel, connectCookieTTL, cfg.SecretKey),
 		retryDelay:  firstRetryDelay,
 		now:         time.Now,
-		pending:     map[string]time.Time{},
+		pending:     auth.NewPendingStates(connectCookieTTL),
 	}, nil
 }
 
@@ -210,10 +209,12 @@ func (s *Service) Connect(w http.ResponseWriter, session auth.Session) (string,
 		Session:  sessionBinding(session),
 	}
 	now := s.now()
+	if err := s.pending.Add(state.State, now); err != nil {
+		return "", fmt.Errorf("gcal: %w", err)
+	}
 	if err := s.cookie.Set(w, &state, now); err != nil {
 		return "", fmt.Errorf("gcal: set connect cookie: %w", err)
 	}
-	s.addPending(state.State, now)
 	return s.oauth.AuthCodeURL(state.State,
 		oauth2.AccessTypeOffline,
 		oauth2.ApprovalForce,
@@ -243,7 +244,7 @@ func (s *Service) Complete(w http.ResponseWriter, r *http.Request, session auth.
 	}
 	// A copy of the cookie verifies until it expires, so the state is
 	// single-use only because it is consumed here.
-	if !s.consumePending(state.State) {
+	if !s.pending.Consume(state.State) {
 		return fmt.Errorf("%w: already used", ErrBadState)
 	}
 	if q.Get("error") != "" {
@@ -276,29 +277,6 @@ func (s *Service) Complete(w http.ResponseWriter, r *http.Request, session auth.
 	return s.store(ctx, sealed)
 }
 
-// addPending records state, started at now, as pending until the connect
-// cookie expires, and forgets expired ones.
-func (s *Service) addPending(state string, now time.Time) {
-	s.pendingMu.Lock()
-	defer s.pendingMu.Unlock()
-	for k, expires := range s.pending {
-		if !now.Before(expires) {
-			delete(s.pending, k)
-		}
-	}
-	s.pending[state] = now.Add(connectCookieTTL)
-}
-
-// consumePending removes state from the pending states, reporting whether
-// it was there.
-func (s *Service) consumePending(state string) bool {
-	s.pendingMu.Lock()
-	defer s.pendingMu.Unlock()
-	_, ok := s.pending[state]
-	delete(s.pending, state)
-	return ok
-}
-
 // checkScopes fails with ErrScopes unless token's granted scopes include
 // every scope in Scopes. Google lets the household untick a permission on
 // the consent screen and reports what was granted in "scope".