at-container-registry/pkg/auth/oauth/session_rev_test.go

package oauth

import (
	"context"
	"errors"
	"testing"

	"github.com/bluesky-social/indigo/atproto/auth/oauth"
)

// revStore is a ClientAuthStore that also reports session revisions, standing in
// for the SQLite store.
type revStore struct {
	oauth.ClientAuthStore

	known     map[string]int64
	current   map[string]int64
	exists    map[string]bool
	lookupErr error
}

func newRevStore() *revStore {
	return &revStore{
		known:   map[string]int64{},
		current: map[string]int64{},
		exists:  map[string]bool{},
	}
}

func (s *revStore) KnownRev(did, sessionID string) (int64, bool) {
	rev, ok := s.known[did+":"+sessionID]
	return rev, ok
}

func (s *revStore) GetSessionRev(_ context.Context, did, sessionID string) (int64, bool, error) {
	if s.lookupErr != nil {
		return 0, false, s.lookupErr
	}
	key := did + ":" + sessionID
	return s.current[key], s.exists[key], nil
}

// plainStore implements no revision tracking, standing in for any store that
// predates this mechanism.
type plainStore struct{ oauth.ClientAuthStore }

func refresherWithStore(store oauth.ClientAuthStore) *Refresher {
	return NewRefresher(&oauth.ClientApp{Store: store})
}

// TestSessionRevisionAdvanced covers the decision that gates session deletion.
//
// Getting this wrong in the "true" direction keeps a genuinely dead session
// alive and leaves the user stuck with no way to recover but waiting; getting it
// wrong in the "false" direction costs a re-login. So every case it cannot
// resolve must answer false.
func TestSessionRevisionAdvanced(t *testing.T) {
	const (
		did       = "did:plc:alice"
		sessionID = "session-1"
		key       = did + ":" + sessionID
	)

	t.Run("another instance wrote since we read", func(t *testing.T) {
		s := newRevStore()
		s.known[key] = 3
		s.current[key] = 4
		s.exists[key] = true

		if !refresherWithStore(s).sessionRevisionAdvanced(context.Background(), did, sessionID) {
			t.Error("expected the advanced revision to be detected")
		}
	})

	t.Run("nobody wrote since we read", func(t *testing.T) {
		s := newRevStore()
		s.known[key] = 3
		s.current[key] = 3
		s.exists[key] = true

		if refresherWithStore(s).sessionRevisionAdvanced(context.Background(), did, sessionID) {
			t.Error("an unchanged revision must not look like a concurrent refresh")
		}
	})

	t.Run("session no longer exists", func(t *testing.T) {
		s := newRevStore()
		s.known[key] = 3
		s.exists[key] = false

		if refresherWithStore(s).sessionRevisionAdvanced(context.Background(), did, sessionID) {
			t.Error("a missing session must not be treated as concurrently refreshed")
		}
	})

	t.Run("no revision recorded for this session", func(t *testing.T) {
		s := newRevStore()
		s.current[key] = 9
		s.exists[key] = true

		if refresherWithStore(s).sessionRevisionAdvanced(context.Background(), did, sessionID) {
			t.Error("without a recorded revision there is nothing to compare against")
		}
	})

	t.Run("revision lookup fails", func(t *testing.T) {
		s := newRevStore()
		s.known[key] = 3
		s.lookupErr = errors.New("database unavailable")

		if refresherWithStore(s).sessionRevisionAdvanced(context.Background(), did, sessionID) {
			t.Error("a failed lookup must fall back to the previous behavior")
		}
	})

	t.Run("store does not track revisions", func(t *testing.T) {
		if refresherWithStore(&plainStore{}).sessionRevisionAdvanced(context.Background(), did, sessionID) {
			t.Error("a store without revisions must keep the previous behavior")
		}
	})

	t.Run("canceled request context still allows the check", func(t *testing.T) {
		s := newRevStore()
		s.known[key] = 3
		s.current[key] = 4
		s.exists[key] = true

		// The inbound request being canceled must not push us into deleting a
		// session another instance just refreshed, so the check runs detached.
		ctx, cancel := context.WithCancel(context.Background())
		cancel()

		if !refresherWithStore(s).sessionRevisionAdvanced(ctx, did, sessionID) {
			t.Error("a canceled inbound context must not suppress the check")
		}
	})
}

// TestStaleSessionErrorUnwraps: DoWithSession reports the underlying auth error
// when a retry does not help, so callers must still be able to inspect it.
func TestStaleSessionErrorUnwraps(t *testing.T) {
	cause := errors.New("invalid_grant")
	err := error(&staleSessionError{did: "did:plc:alice", cause: cause})

	if !errors.Is(err, cause) {
		t.Error("staleSessionError does not unwrap to its cause")
	}

	var stale *staleSessionError
	if !errors.As(err, &stale) {
		t.Error("staleSessionError is not detectable with errors.As")
	}
}

// TestDoWithSessionRetriesOnceOnStaleSession pins the retry policy: exactly one
// retry. Looping would hold the per-DID lock while getting the same answer.
func TestDoWithSessionRetriesOnceOnStaleSession(t *testing.T) {
	// resumeSession requires a store implementing GetLatestSessionForDID, which
	// this deliberately does not, so every attempt fails early and we can count
	// attempts without standing up a PDS.
	r := refresherWithStore(&plainStore{})

	attempts := 0
	fn := func(*oauth.ClientSession) error {
		attempts++
		return nil
	}

	// The DID never resolves to a session, so fn is never reached; what matters
	// is that DoWithSession returns rather than looping.
	err := r.DoWithSession(context.Background(), "did:plc:alice", fn)
	if err == nil {
		t.Fatal("expected an error when no session can be resumed")
	}
	if attempts != 0 {
		t.Errorf("fn ran %d times despite the session never resuming", attempts)
	}
}