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

package db

import (
	"context"
	"testing"
	"time"

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

func TestInvalidateSessionsWithMismatchedScopes(t *testing.T) {
	// Create in-memory test database
	db, err := InitDB(":memory:")
	if err != nil {
		t.Fatalf("Failed to init database: %v", err)
	}
	defer db.Close()

	store := NewOAuthStore(db)
	ctx := context.Background()

	// Test 1: Empty database - should return 0
	count, err := store.InvalidateSessionsWithMismatchedScopes(ctx, []string{"atproto", "blob:image/png"})
	if err != nil {
		t.Fatalf("Expected no error with empty DB, got: %v", err)
	}
	if count != 0 {
		t.Errorf("Expected 0 invalidated sessions in empty DB, got %d", count)
	}

	// Helper to create session data
	createSession := func(did, sessionID string, scopes []string) oauth.ClientSessionData {
		parsedDID, _ := syntax.ParseDID(did)
		return oauth.ClientSessionData{
			AccountDID:              parsedDID,
			SessionID:               sessionID,
			HostURL:                 "https://bsky.social",
			AuthServerURL:           "https://bsky.social",
			AuthServerTokenEndpoint: "https://bsky.social/oauth/token",
			Scopes:                  scopes,
			AccessToken:             "test_access_token",
			RefreshToken:            "test_refresh_token",
			DPoPAuthServerNonce:     "test_nonce",
			DPoPHostNonce:           "test_host_nonce",
			DPoPPrivateKeyMultibase: "test_key",
		}
	}

	// Test 2: Session with matching scopes - should not be invalidated
	matchingSession := createSession("did:plc:test1", "session1", []string{"atproto", "blob:image/png"})
	if err := store.SaveSession(ctx, matchingSession); err != nil {
		t.Fatalf("Failed to save matching session: %v", err)
	}

	count, err = store.InvalidateSessionsWithMismatchedScopes(ctx, []string{"atproto", "blob:image/png"})
	if err != nil {
		t.Fatalf("Expected no error, got: %v", err)
	}
	if count != 0 {
		t.Errorf("Expected 0 invalidated sessions (all match), got %d", count)
	}

	// Verify session still exists
	retrieved, err := store.GetSession(ctx, matchingSession.AccountDID, matchingSession.SessionID)
	if err != nil {
		t.Errorf("Expected session to still exist, got error: %v", err)
	}
	if retrieved == nil {
		t.Error("Expected session to still exist, got nil")
	}

	// Test 3: Session with mismatched scopes (missing scope) - should be invalidated
	mismatchedSession := createSession("did:plc:test2", "session2", []string{"atproto"}) // Missing blob scope
	if err := store.SaveSession(ctx, mismatchedSession); err != nil {
		t.Fatalf("Failed to save mismatched session: %v", err)
	}

	count, err = store.InvalidateSessionsWithMismatchedScopes(ctx, []string{"atproto", "blob:image/png"})
	if err != nil {
		t.Fatalf("Expected no error, got: %v", err)
	}
	if count != 1 {
		t.Errorf("Expected 1 invalidated session, got %d", count)
	}

	// Verify mismatched session was deleted
	_, err = store.GetSession(ctx, mismatchedSession.AccountDID, mismatchedSession.SessionID)
	if err == nil {
		t.Error("Expected session to be deleted (should error), but got no error")
	}

	// Test 4: Session with extra scopes - should be invalidated
	extraScopeSession := createSession("did:plc:test3", "session3", []string{"atproto", "blob:image/png", "extra:scope"})
	if err := store.SaveSession(ctx, extraScopeSession); err != nil {
		t.Fatalf("Failed to save extra scope session: %v", err)
	}

	count, err = store.InvalidateSessionsWithMismatchedScopes(ctx, []string{"atproto", "blob:image/png"})
	if err != nil {
		t.Fatalf("Expected no error, got: %v", err)
	}
	if count != 1 {
		t.Errorf("Expected 1 invalidated session (extra scope), got %d", count)
	}

	// Test 5: Multiple sessions with mixed matches - only mismatch should be invalidated
	matching1 := createSession("did:plc:test4", "session4", []string{"atproto", "blob:image/png"})
	matching2 := createSession("did:plc:test5", "session5", []string{"blob:image/png", "atproto"}) // Different order
	mismatched1 := createSession("did:plc:test6", "session6", []string{"atproto"})
	mismatched2 := createSession("did:plc:test7", "session7", []string{"wrong", "scopes"})

	for _, sess := range []oauth.ClientSessionData{matching1, matching2, mismatched1, mismatched2} {
		if err := store.SaveSession(ctx, sess); err != nil {
			t.Fatalf("Failed to save session: %v", err)
		}
	}

	count, err = store.InvalidateSessionsWithMismatchedScopes(ctx, []string{"atproto", "blob:image/png"})
	if err != nil {
		t.Fatalf("Expected no error, got: %v", err)
	}
	if count != 2 {
		t.Errorf("Expected 2 invalidated sessions, got %d", count)
	}

	// Verify matching sessions still exist
	for _, sess := range []oauth.ClientSessionData{matching1, matching2} {
		retrieved, err := store.GetSession(ctx, sess.AccountDID, sess.SessionID)
		if err != nil {
			t.Errorf("Expected matching session %s to exist, got error: %v", sess.SessionID, err)
		}
		if retrieved == nil {
			t.Errorf("Expected matching session %s to exist, got nil", sess.SessionID)
		}
	}

	// Test 6: Malformed session data - should be deleted
	parsedDID, _ := syntax.ParseDID("did:plc:test8")
	_, err = db.ExecContext(ctx, `
		INSERT INTO oauth_sessions (session_key, account_did, session_id, session_data, created_at, updated_at)
		VALUES (?, ?, ?, ?, datetime('now'), datetime('now'))
	`, makeSessionKey("did:plc:test8", "malformed"), "did:plc:test8", "malformed", "invalid json data")
	if err != nil {
		t.Fatalf("Failed to insert malformed session: %v", err)
	}

	count, err = store.InvalidateSessionsWithMismatchedScopes(ctx, []string{"atproto", "blob:image/png"})
	if err != nil {
		t.Fatalf("Expected no error handling malformed data, got: %v", err)
	}
	if count != 1 {
		t.Errorf("Expected 1 invalidated session (malformed), got %d", count)
	}

	// Verify malformed session was deleted
	_, err = store.GetSession(ctx, parsedDID, "malformed")
	if err == nil {
		t.Error("Expected malformed session to be deleted, but got no error")
	}
}

func TestScopesMatch(t *testing.T) {
	// Test oauth.ScopesMatch function including include: scope expansion
	tests := []struct {
		name     string
		stored   []string
		desired  []string
		expected bool
	}{
		{
			name:     "exact match",
			stored:   []string{"atproto", "blob:image/png"},
			desired:  []string{"atproto", "blob:image/png"},
			expected: true,
		},
		{
			name:     "different order",
			stored:   []string{"blob:image/png", "atproto"},
			desired:  []string{"atproto", "blob:image/png"},
			expected: true,
		},
		{
			name:     "missing scope",
			stored:   []string{"atproto"},
			desired:  []string{"atproto", "blob:image/png"},
			expected: false,
		},
		{
			name:     "extra scope",
			stored:   []string{"atproto", "blob:image/png", "extra"},
			desired:  []string{"atproto", "blob:image/png"},
			expected: false,
		},
		{
			name:     "both empty",
			stored:   []string{},
			desired:  []string{},
			expected: true,
		},
		{
			name:     "nil vs empty",
			stored:   nil,
			desired:  []string{},
			expected: true,
		},
		{
			name: "include scope expansion",
			stored: []string{
				"atproto",
				"repo?collection=io.atcr.manifest&collection=io.atcr.repo.page&collection=io.atcr.sailor.profile&collection=io.atcr.sailor.star&collection=io.atcr.tag",
			},
			desired: []string{
				"atproto",
				"include:io.atcr.authFullApp",
			},
			expected: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := atcroauth.ScopesMatch(tt.stored, tt.desired)
			if result != tt.expected {
				t.Errorf("ScopesMatch(%v, %v) = %v, want %v",
					tt.stored, tt.desired, result, tt.expected)
			}
		})
	}
}

func TestOAuthStoreSessionLifecycle(t *testing.T) {
	// Basic test to ensure SaveSession, GetSession, DeleteSession work correctly
	db, err := InitDB(":memory:")
	if err != nil {
		t.Fatalf("Failed to init database: %v", err)
	}
	defer db.Close()

	store := NewOAuthStore(db)
	ctx := context.Background()

	// Create test session
	did, _ := syntax.ParseDID("did:plc:testuser")
	sessionData := oauth.ClientSessionData{
		AccountDID:              did,
		SessionID:               "test_session_id",
		HostURL:                 "https://bsky.social",
		AuthServerURL:           "https://bsky.social",
		AuthServerTokenEndpoint: "https://bsky.social/oauth/token",
		Scopes:                  []string{"atproto", "blob:image/png"},
		AccessToken:             "test_access_token",
		RefreshToken:            "test_refresh_token",
		DPoPAuthServerNonce:     "test_nonce",
		DPoPHostNonce:           "test_host_nonce",
		DPoPPrivateKeyMultibase: "test_key",
	}

	// Test SaveSession
	if err := store.SaveSession(ctx, sessionData); err != nil {
		t.Fatalf("Failed to save session: %v", err)
	}

	// Test GetSession
	retrieved, err := store.GetSession(ctx, did, "test_session_id")
	if err != nil {
		t.Fatalf("Failed to get session: %v", err)
	}
	if retrieved == nil {
		t.Fatal("Retrieved session is nil")
	}
	if retrieved.SessionID != sessionData.SessionID {
		t.Errorf("Expected session ID %s, got %s", sessionData.SessionID, retrieved.SessionID)
	}
	if len(retrieved.Scopes) != len(sessionData.Scopes) {
		t.Errorf("Expected %d scopes, got %d", len(sessionData.Scopes), len(retrieved.Scopes))
	}

	// Test UpdateSession (upsert)
	sessionData.AccessToken = "new_access_token"
	if err := store.SaveSession(ctx, sessionData); err != nil {
		t.Fatalf("Failed to update session: %v", err)
	}

	retrieved, err = store.GetSession(ctx, did, "test_session_id")
	if err != nil {
		t.Fatalf("Failed to get updated session: %v", err)
	}
	if retrieved.AccessToken != "new_access_token" {
		t.Errorf("Expected updated access token, got %s", retrieved.AccessToken)
	}

	// Test DeleteSession
	if err := store.DeleteSession(ctx, did, "test_session_id"); err != nil {
		t.Fatalf("Failed to delete session: %v", err)
	}

	// Verify deletion
	_, err = store.GetSession(ctx, did, "test_session_id")
	if err == nil {
		t.Error("Expected error after deletion, got nil")
	}
}

func TestCleanupOldSessions(t *testing.T) {
	db, err := InitDB(":memory:")
	if err != nil {
		t.Fatalf("Failed to init database: %v", err)
	}
	defer db.Close()

	store := NewOAuthStore(db)
	ctx := context.Background()

	// Insert old session (31 days ago)
	did1, _ := syntax.ParseDID("did:plc:old")
	oldSessionData := oauth.ClientSessionData{
		AccountDID:              did1,
		SessionID:               "old_session",
		HostURL:                 "https://bsky.social",
		AuthServerURL:           "https://bsky.social",
		AuthServerTokenEndpoint: "https://bsky.social/oauth/token",
		Scopes:                  []string{"atproto"},
		AccessToken:             "old_token",
		RefreshToken:            "old_refresh",
		DPoPAuthServerNonce:     "old_nonce",
		DPoPHostNonce:           "old_host_nonce",
		DPoPPrivateKeyMultibase: "old_key",
	}

	// Save and manually update timestamp to be old
	if err := store.SaveSession(ctx, oldSessionData); err != nil {
		t.Fatalf("Failed to save old session: %v", err)
	}

	// Update timestamp to 31 days ago
	oldTime := time.Now().Add(-31 * 24 * time.Hour)
	_, err = db.ExecContext(ctx, `
		UPDATE oauth_sessions
		SET updated_at = ?
		WHERE session_key = ?
	`, oldTime, makeSessionKey(did1.String(), "old_session"))
	if err != nil {
		t.Fatalf("Failed to update session timestamp: %v", err)
	}

	// Insert recent session (1 day ago)
	did2, _ := syntax.ParseDID("did:plc:recent")
	recentSessionData := oauth.ClientSessionData{
		AccountDID:              did2,
		SessionID:               "recent_session",
		HostURL:                 "https://bsky.social",
		AuthServerURL:           "https://bsky.social",
		AuthServerTokenEndpoint: "https://bsky.social/oauth/token",
		Scopes:                  []string{"atproto"},
		AccessToken:             "recent_token",
		RefreshToken:            "recent_refresh",
		DPoPAuthServerNonce:     "recent_nonce",
		DPoPHostNonce:           "recent_host_nonce",
		DPoPPrivateKeyMultibase: "recent_key",
	}

	if err := store.SaveSession(ctx, recentSessionData); err != nil {
		t.Fatalf("Failed to save recent session: %v", err)
	}

	// Run cleanup (remove sessions older than 30 days)
	store.CleanupOldSessions(ctx, 30*24*time.Hour)

	// Verify old session was deleted
	_, err = store.GetSession(ctx, did1, "old_session")
	if err == nil {
		t.Error("Expected old session to be deleted")
	}

	// Verify recent session still exists
	_, err = store.GetSession(ctx, did2, "recent_session")
	if err != nil {
		t.Errorf("Expected recent session to exist, got error: %v", err)
	}
}

// TestMakeSessionKey tests the session key generation function
func TestMakeSessionKey(t *testing.T) {
	tests := []struct {
		name      string
		did       string
		sessionID string
		expected  string
	}{
		{
			name:      "normal case",
			did:       "did:plc:abc123",
			sessionID: "session_xyz789",
			expected:  "did:plc:abc123:session_xyz789",
		},
		{
			name:      "empty did",
			did:       "",
			sessionID: "session123",
			expected:  ":session123",
		},
		{
			name:      "empty session",
			did:       "did:plc:test",
			sessionID: "",
			expected:  "did:plc:test:",
		},
		{
			name:      "both empty",
			did:       "",
			sessionID: "",
			expected:  ":",
		},
		{
			name:      "with colon in did",
			did:       "did:web:example.com",
			sessionID: "session123",
			expected:  "did:web:example.com:session123",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := makeSessionKey(tt.did, tt.sessionID)
			if result != tt.expected {
				t.Errorf("makeSessionKey(%q, %q) = %q, want %q", tt.did, tt.sessionID, result, tt.expected)
			}
		})
	}
}