at-container-registry/pkg/appview/db/oauth_session_rev_test.go

package db

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"path/filepath"
	"testing"

	atcroauth "atcr.io/pkg/auth/oauth"
	"github.com/bluesky-social/indigo/atproto/auth/oauth"
	"github.com/bluesky-social/indigo/atproto/syntax"
)

// revTestDB returns a file-backed database.
//
// Not ":memory:" — go-libsql gives each connection to an in-memory DSN its own
// private database, and these tests deliberately use two OAuthStore values to
// stand in for two AppView instances sharing one database.
func revTestDB(t *testing.T) *sql.DB {
	t.Helper()
	database, err := InitDB(filepath.Join(t.TempDir(), "oauth.db"), LibsqlConfig{})
	if err != nil {
		t.Fatalf("InitDB: %v", err)
	}
	t.Cleanup(func() { database.Close() })
	return database
}

func sessionData(t *testing.T, did, sessionID, accessToken, refreshToken string) oauth.ClientSessionData {
	t.Helper()
	parsed, err := syntax.ParseDID(did)
	if err != nil {
		t.Fatalf("ParseDID(%q): %v", did, err)
	}
	return oauth.ClientSessionData{
		AccountDID:              parsed,
		SessionID:               sessionID,
		HostURL:                 "https://bsky.social",
		AuthServerURL:           "https://bsky.social",
		AuthServerTokenEndpoint: "https://bsky.social/oauth/token",
		Scopes:                  []string{"atproto"},
		AccessToken:             accessToken,
		RefreshToken:            refreshToken,
		DPoPPrivateKeyMultibase: "test_key",
	}
}

func storedTokens(t *testing.T, database *sql.DB, did, sessionID string) (access, refresh string) {
	t.Helper()
	var raw string
	err := database.QueryRow(
		`SELECT session_data FROM oauth_sessions WHERE session_key = ?`,
		makeSessionKey(did, sessionID),
	).Scan(&raw)
	if err != nil {
		t.Fatalf("read session: %v", err)
	}
	var data oauth.ClientSessionData
	if err := unmarshalSession(raw, &data); err != nil {
		t.Fatalf("parse session: %v", err)
	}
	return data.AccessToken, data.RefreshToken
}

// TestConcurrentRefreshDoesNotClobberRotatedTokens is the scenario this whole
// mechanism exists for.
//
// Refresh tokens rotate on use. Two AppView instances refreshing the same
// account both hold a session they believe is current. Without a revision check
// the one that writes second replaces the freshly rotated tokens with ones the
// auth server has already invalidated, and the next refresh fails with
// invalid_grant for everyone.
func TestConcurrentRefreshDoesNotClobberRotatedTokens(t *testing.T) {
	database := revTestDB(t)
	ctx := context.Background()

	const (
		did       = "did:plc:alice"
		sessionID = "session-1"
	)

	// One instance creates the session.
	instanceA := NewOAuthStore(database)
	if err := instanceA.SaveSession(ctx, sessionData(t, did, sessionID, "access-v1", "refresh-v1")); err != nil {
		t.Fatalf("initial save: %v", err)
	}

	// A second instance, with its own in-process revision map, reads it.
	instanceB := NewOAuthStore(database)
	parsed, _ := syntax.ParseDID(did)
	if _, err := instanceB.GetSession(ctx, parsed, sessionID); err != nil {
		t.Fatalf("instance B read: %v", err)
	}
	if _, err := instanceA.GetSession(ctx, parsed, sessionID); err != nil {
		t.Fatalf("instance A read: %v", err)
	}

	// B refreshes first and persists the rotated tokens.
	if err := instanceB.SaveSession(ctx, sessionData(t, did, sessionID, "access-v2", "refresh-v2")); err != nil {
		t.Fatalf("instance B save: %v", err)
	}

	// A now tries to persist what it read before B's write.
	err := instanceA.SaveSession(ctx, sessionData(t, did, sessionID, "access-stale", "refresh-stale"))
	if !errors.Is(err, ErrSessionRevConflict) {
		t.Fatalf("expected ErrSessionRevConflict from the losing writer, got %v", err)
	}

	access, refresh := storedTokens(t, database, did, sessionID)
	if access != "access-v2" || refresh != "refresh-v2" {
		t.Errorf("the stale writer clobbered the rotated tokens: stored access=%q refresh=%q, want the v2 pair",
			access, refresh)
	}
}

// TestSaveSessionAdvancesRevision: a successful write must bump the revision, or
// nothing else can detect that it happened.
func TestSaveSessionAdvancesRevision(t *testing.T) {
	database := revTestDB(t)
	ctx := context.Background()
	store := NewOAuthStore(database)

	const (
		did       = "did:plc:alice"
		sessionID = "session-1"
	)

	if err := store.SaveSession(ctx, sessionData(t, did, sessionID, "a1", "r1")); err != nil {
		t.Fatalf("save: %v", err)
	}
	first, exists, err := store.GetSessionRev(ctx, did, sessionID)
	if err != nil || !exists {
		t.Fatalf("GetSessionRev: exists=%v err=%v", exists, err)
	}

	if err := store.SaveSession(ctx, sessionData(t, did, sessionID, "a2", "r2")); err != nil {
		t.Fatalf("second save: %v", err)
	}
	second, _, err := store.GetSessionRev(ctx, did, sessionID)
	if err != nil {
		t.Fatalf("GetSessionRev: %v", err)
	}

	if second <= first {
		t.Errorf("revision did not advance: %d then %d", first, second)
	}
}

// TestSameStoreRepeatedSavesSucceed: the CAS must not get in the way of the
// ordinary case, one instance writing a session it owns over and over as DPoP
// nonces change.
func TestSameStoreRepeatedSavesSucceed(t *testing.T) {
	database := revTestDB(t)
	ctx := context.Background()
	store := NewOAuthStore(database)

	const (
		did       = "did:plc:alice"
		sessionID = "session-1"
	)

	for i := range 10 {
		if err := store.SaveSession(ctx, sessionData(t, did, sessionID, "a", "r")); err != nil {
			t.Fatalf("save %d: %v", i, err)
		}
	}
}

// TestGetSessionRecordsRevisionForLaterCAS: reading is what establishes the
// expectation a later write compares against.
func TestGetSessionRecordsRevisionForLaterCAS(t *testing.T) {
	database := revTestDB(t)
	ctx := context.Background()

	const (
		did       = "did:plc:alice"
		sessionID = "session-1"
	)

	writer := NewOAuthStore(database)
	if err := writer.SaveSession(ctx, sessionData(t, did, sessionID, "a1", "r1")); err != nil {
		t.Fatalf("save: %v", err)
	}

	reader := NewOAuthStore(database)
	if _, ok := reader.KnownRev(did, sessionID); ok {
		t.Fatal("a store that has never read the session should know no revision")
	}

	parsed, _ := syntax.ParseDID(did)
	if _, err := reader.GetSession(ctx, parsed, sessionID); err != nil {
		t.Fatalf("GetSession: %v", err)
	}
	if _, ok := reader.KnownRev(did, sessionID); !ok {
		t.Error("GetSession did not record a revision")
	}
}

// TestGetLatestSessionForDIDRecordsRevision covers the path the refresher
// actually uses: resumeSession loads through GetLatestSessionForDID, so if that
// did not record a revision the CAS would never engage where it matters most.
func TestGetLatestSessionForDIDRecordsRevision(t *testing.T) {
	database := revTestDB(t)
	ctx := context.Background()

	const (
		did       = "did:plc:alice"
		sessionID = "session-1"
	)

	writer := NewOAuthStore(database)
	if err := writer.SaveSession(ctx, sessionData(t, did, sessionID, "a1", "r1")); err != nil {
		t.Fatalf("save: %v", err)
	}

	reader := NewOAuthStore(database)
	if _, _, err := reader.GetLatestSessionForDID(ctx, did); err != nil {
		t.Fatalf("GetLatestSessionForDID: %v", err)
	}
	if _, ok := reader.KnownRev(did, sessionID); !ok {
		t.Error("GetLatestSessionForDID did not record a revision")
	}
}

// TestDeleteSessionForgetsRevision: a remembered revision must not survive the
// session it describes, or a re-created session would be compared against a
// revision from the dead one.
func TestDeleteSessionForgetsRevision(t *testing.T) {
	database := revTestDB(t)
	ctx := context.Background()
	store := NewOAuthStore(database)

	const (
		did       = "did:plc:alice"
		sessionID = "session-1"
	)

	if err := store.SaveSession(ctx, sessionData(t, did, sessionID, "a1", "r1")); err != nil {
		t.Fatalf("save: %v", err)
	}
	parsed, _ := syntax.ParseDID(did)
	if err := store.DeleteSession(ctx, parsed, sessionID); err != nil {
		t.Fatalf("DeleteSession: %v", err)
	}
	if _, ok := store.KnownRev(did, sessionID); ok {
		t.Error("revision survived the session being deleted")
	}

	// A re-created session must save cleanly rather than hitting a phantom
	// conflict.
	if err := store.SaveSession(ctx, sessionData(t, did, sessionID, "a2", "r2")); err != nil {
		t.Errorf("re-created session failed to save: %v", err)
	}
}

func TestDeleteSessionsForDIDForgetsRevisions(t *testing.T) {
	database := revTestDB(t)
	ctx := context.Background()
	store := NewOAuthStore(database)

	const did = "did:plc:alice"
	for _, sid := range []string{"s1", "s2"} {
		if err := store.SaveSession(ctx, sessionData(t, did, sid, "a", "r")); err != nil {
			t.Fatalf("save %s: %v", sid, err)
		}
	}

	if err := store.DeleteSessionsForDID(ctx, did); err != nil {
		t.Fatalf("DeleteSessionsForDID: %v", err)
	}
	for _, sid := range []string{"s1", "s2"} {
		if _, ok := store.KnownRev(did, sid); ok {
			t.Errorf("revision for %s survived the bulk delete", sid)
		}
	}
}

// TestErrSessionRevConflictIsSharedAcrossPackages: pkg/auth/oauth compares
// against its own sentinel, so the two must be the same value or the refresher
// would treat a conflict as an ordinary failure and log it as an error.
func TestErrSessionRevConflictIsSharedAcrossPackages(t *testing.T) {
	if !errors.Is(ErrSessionRevConflict, atcroauth.ErrSessionRevConflict) {
		t.Error("db.ErrSessionRevConflict and oauth.ErrSessionRevConflict are different errors")
	}
}

// unmarshalSession is a tiny helper so the test does not need to import
// encoding/json separately from the store's own representation.
func unmarshalSession(raw string, out *oauth.ClientSessionData) error {
	return json.Unmarshal([]byte(raw), out)
}