at-container-registry/pkg/appview/middleware/hold_prefs_test.go

package middleware

import (
	"context"
	"database/sql"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"atcr.io/pkg/appview/db"
	"atcr.io/pkg/atproto"
)

// The registry used to read the owner's sailor profile from their PDS on every
// single /v2/ request, which made a ten-layer push forty-odd round trips to a
// third-party server before a byte of image data moved. These tests count the
// PDS requests rather than only checking the values that come back, because the
// whole point of the change is the request that does not happen.

const prefsOwnerDID = "did:plc:prefsowner"
const prefsOwnerHandle = "prefs.example.com"

// profilePDS starts a PDS whose getRecord handler is supplied by the caller,
// and returns the server plus a counter of profile reads it served.
func profilePDS(t *testing.T, handler http.HandlerFunc) (*httptest.Server, *atomic.Int64) {
	t.Helper()
	var calls atomic.Int64
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/xrpc/com.atproto.repo.getRecord" {
			calls.Add(1)
			handler(w, r)
			return
		}
		w.WriteHeader(http.StatusNotFound)
	}))
	t.Cleanup(srv.Close)
	return srv, &calls
}

// servesProfile answers every profile read with the given record.
func servesProfile(profile *atproto.SailorProfileRecord) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{"value": profile})
	}
}

func prefsTestDB(t *testing.T) *sql.DB {
	t.Helper()
	database, err := db.InitDB(":memory:", db.LibsqlConfig{})
	require.NoError(t, err, "open test database")
	t.Cleanup(func() { database.Close() })
	return database
}

// seedUser inserts a users row. autoRemove nil leaves auto_remove_untagged
// NULL, which is the "never learned" state.
func seedUser(t *testing.T, database *sql.DB, holdDID string, autoRemove *bool) {
	t.Helper()
	var val any
	if autoRemove != nil {
		val = *autoRemove
	}
	_, err := database.Exec(`
		INSERT INTO users (did, handle, pds_endpoint, default_hold_did, auto_remove_untagged, last_seen)
		VALUES (?, ?, ?, ?, ?, ?)`,
		prefsOwnerDID, prefsOwnerHandle, "https://pds.example.com", holdDID, val, time.Now())
	require.NoError(t, err, "seed user")
}

// readCachedPrefs returns the raw row state, with auto_remove_untagged as a
// tri-state so a written 0 is distinguishable from a still-NULL column.
func readCachedPrefs(t *testing.T, database *sql.DB) (found bool, holdDID sql.NullString, autoRemove sql.NullBool) {
	t.Helper()
	err := database.QueryRow(
		`SELECT default_hold_did, auto_remove_untagged FROM users WHERE did = ?`,
		prefsOwnerDID).Scan(&holdDID, &autoRemove)
	if err == sql.ErrNoRows {
		return false, holdDID, autoRemove
	}
	require.NoError(t, err, "read cached prefs")
	return true, holdDID, autoRemove
}

// TestFindHoldDIDAndPrefs_CachedRowMakesNoPDSRequest is the whole point of the
// change: when the Jetstream-fed row has both answers, the hot path must not
// touch the owner's PDS at all.
func TestFindHoldDIDAndPrefs_CachedRowMakesNoPDSRequest(t *testing.T) {
	database := prefsTestDB(t)
	known := true
	seedUser(t, database, "did:web:user.hold.io", &known)

	pds, calls := profilePDS(t, servesProfile(atproto.NewSailorProfileRecord("did:web:should.not.be.read")))

	resolver := &NamespaceResolver{
		defaultHoldDID: "did:web:default.atcr.io",
		userPrefs:      db.NewHoldDIDDB(database),
	}

	holdDID, prefs := resolver.findHoldDIDAndPrefs(context.Background(), prefsOwnerDID, prefsOwnerHandle, pds.URL)

	assert.Equal(t, "did:web:user.hold.io", holdDID, "hold DID should come from the cached row")
	assert.True(t, prefs.AutoRemoveUntagged, "auto-remove should come from the cached row")
	assert.Equal(t, int64(0), calls.Load(), "a fully cached row must not read the PDS")
}

// TestFindHoldDIDAndPrefs_CachedEmptyHoldUsesDefault pins the meaning of an
// empty default_hold_did: not "unknown", but "use the operator's hold", the
// same reading the auth gate's resolveHoldDID gives it.
func TestFindHoldDIDAndPrefs_CachedEmptyHoldUsesDefault(t *testing.T) {
	database := prefsTestDB(t)
	known := false
	seedUser(t, database, "", &known)

	pds, calls := profilePDS(t, servesProfile(atproto.NewSailorProfileRecord("did:web:should.not.be.read")))

	resolver := &NamespaceResolver{
		defaultHoldDID: "did:web:default.atcr.io",
		userPrefs:      db.NewHoldDIDDB(database),
	}

	holdDID, prefs := resolver.findHoldDIDAndPrefs(context.Background(), prefsOwnerDID, prefsOwnerHandle, pds.URL)

	assert.Equal(t, "did:web:default.atcr.io", holdDID, "an empty cached hold means the appview default")
	assert.False(t, prefs.AutoRemoveUntagged)
	assert.Equal(t, int64(0), calls.Load(), "a known-false auto-remove is still a known value")
}

// TestFindHoldDIDAndPrefs_NullAutoRemoveFetchesOnce covers the window after a
// deploy, while the backfill is still filling the new column: one live fetch,
// written back, and never again.
func TestFindHoldDIDAndPrefs_NullAutoRemoveFetchesOnce(t *testing.T) {
	database := prefsTestDB(t)
	seedUser(t, database, "did:web:stale.hold.io", nil) // auto_remove_untagged NULL

	profile := atproto.NewSailorProfileRecord("did:web:user.hold.io")
	profile.AutoRemoveUntagged = true
	pds, calls := profilePDS(t, servesProfile(profile))

	resolver := &NamespaceResolver{
		defaultHoldDID: "did:web:default.atcr.io",
		userPrefs:      db.NewHoldDIDDB(database),
	}
	ctx := context.Background()

	holdDID, prefs := resolver.findHoldDIDAndPrefs(ctx, prefsOwnerDID, prefsOwnerHandle, pds.URL)
	assert.Equal(t, "did:web:user.hold.io", holdDID, "the live profile should serve this request")
	assert.True(t, prefs.AutoRemoveUntagged)
	assert.Equal(t, int64(1), calls.Load(), "a NULL auto-remove should cost exactly one fetch")

	found, cachedHold, cachedAuto := readCachedPrefs(t, database)
	require.True(t, found)
	assert.Equal(t, "did:web:user.hold.io", cachedHold.String, "the fetch should be written back")
	require.True(t, cachedAuto.Valid, "auto_remove_untagged must no longer be NULL")
	assert.True(t, cachedAuto.Bool)

	holdDID, prefs = resolver.findHoldDIDAndPrefs(ctx, prefsOwnerDID, prefsOwnerHandle, pds.URL)
	assert.Equal(t, "did:web:user.hold.io", holdDID)
	assert.True(t, prefs.AutoRemoveUntagged)
	assert.Equal(t, int64(1), calls.Load(), "the second request must be served from the row")
}

// TestFindHoldDIDAndPrefs_MissingRowFetchesOnce covers a user the backfill has
// never reached: no row at all, and the fallback still has to bound itself.
func TestFindHoldDIDAndPrefs_MissingRowFetchesOnce(t *testing.T) {
	database := prefsTestDB(t)

	profile := atproto.NewSailorProfileRecord("did:web:user.hold.io")
	profile.AutoRemoveUntagged = true
	pds, calls := profilePDS(t, servesProfile(profile))

	resolver := &NamespaceResolver{
		defaultHoldDID: "did:web:default.atcr.io",
		userPrefs:      db.NewHoldDIDDB(database),
	}
	ctx := context.Background()

	holdDID, prefs := resolver.findHoldDIDAndPrefs(ctx, prefsOwnerDID, prefsOwnerHandle, pds.URL)
	assert.Equal(t, "did:web:user.hold.io", holdDID)
	assert.True(t, prefs.AutoRemoveUntagged)
	assert.Equal(t, int64(1), calls.Load())

	found, cachedHold, cachedAuto := readCachedPrefs(t, database)
	require.True(t, found, "the fallback should create the row it was missing")
	assert.Equal(t, "did:web:user.hold.io", cachedHold.String)
	require.True(t, cachedAuto.Valid)
	assert.True(t, cachedAuto.Bool)

	_, _ = resolver.findHoldDIDAndPrefs(ctx, prefsOwnerDID, prefsOwnerHandle, pds.URL)
	assert.Equal(t, int64(1), calls.Load(), "a created row must stop the fallback repeating")
}

// TestFindHoldDIDAndPrefs_MissingProfileCachedAsKnown covers the user who has
// never written a sailor profile at all. A 404 is an answer, so it has to be
// written down as 0 rather than left NULL, or that user pays for a PDS round
// trip on every request forever.
func TestFindHoldDIDAndPrefs_MissingProfileCachedAsKnown(t *testing.T) {
	database := prefsTestDB(t)

	pds, calls := profilePDS(t, func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotFound)
	})

	resolver := &NamespaceResolver{
		defaultHoldDID: "did:web:default.atcr.io",
		userPrefs:      db.NewHoldDIDDB(database),
	}
	ctx := context.Background()

	holdDID, prefs := resolver.findHoldDIDAndPrefs(ctx, prefsOwnerDID, prefsOwnerHandle, pds.URL)
	assert.Equal(t, "did:web:default.atcr.io", holdDID, "no profile means the appview default")
	assert.False(t, prefs.AutoRemoveUntagged)
	assert.Equal(t, int64(1), calls.Load())

	found, cachedHold, cachedAuto := readCachedPrefs(t, database)
	require.True(t, found, "a 404 profile should still leave a row behind")
	assert.Equal(t, "", cachedHold.String, "no custom hold: empty means the operator default")
	require.True(t, cachedAuto.Valid, "a 404 is a known answer, not an unknown one")
	assert.False(t, cachedAuto.Bool)

	_, _ = resolver.findHoldDIDAndPrefs(ctx, prefsOwnerDID, prefsOwnerHandle, pds.URL)
	assert.Equal(t, int64(1), calls.Load(), "a user with no profile must not refetch every request")
}

// TestFindHoldDIDAndPrefs_FailedFetchIsNotCached is the other side of the same
// coin: a network failure is not an answer. Caching it would freeze a transient
// blip into the user's routing for good.
func TestFindHoldDIDAndPrefs_FailedFetchIsNotCached(t *testing.T) {
	database := prefsTestDB(t)

	pds, calls := profilePDS(t, func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	})

	resolver := &NamespaceResolver{
		defaultHoldDID: "did:web:default.atcr.io",
		userPrefs:      db.NewHoldDIDDB(database),
	}
	ctx := context.Background()

	holdDID, prefs := resolver.findHoldDIDAndPrefs(ctx, prefsOwnerDID, prefsOwnerHandle, pds.URL)
	assert.Equal(t, "did:web:default.atcr.io", holdDID, "a failed fetch should degrade to the default hold")
	assert.False(t, prefs.AutoRemoveUntagged)
	assert.Equal(t, int64(1), calls.Load())

	found, _, _ := readCachedPrefs(t, database)
	assert.False(t, found, "a transient error must not be written down as the user's preference")

	_, _ = resolver.findHoldDIDAndPrefs(ctx, prefsOwnerDID, prefsOwnerHandle, pds.URL)
	assert.Equal(t, int64(2), calls.Load(), "an uncached failure should be retried on the next request")
}

// TestFindHoldDIDAndPrefs_FailedFetchLeavesExistingRowAlone is the same rule
// applied to a user who already has a row: the blip must not clear the hold
// they are actually using, nor stamp a guessed auto-remove over the unknown.
func TestFindHoldDIDAndPrefs_FailedFetchLeavesExistingRowAlone(t *testing.T) {
	database := prefsTestDB(t)
	seedUser(t, database, "did:web:user.hold.io", nil)

	pds, _ := profilePDS(t, func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	})

	resolver := &NamespaceResolver{
		defaultHoldDID: "did:web:default.atcr.io",
		userPrefs:      db.NewHoldDIDDB(database),
	}

	_, _ = resolver.findHoldDIDAndPrefs(context.Background(), prefsOwnerDID, prefsOwnerHandle, pds.URL)

	found, cachedHold, cachedAuto := readCachedPrefs(t, database)
	require.True(t, found)
	assert.Equal(t, "did:web:user.hold.io", cachedHold.String, "a failed fetch must not clear the cached hold")
	assert.False(t, cachedAuto.Valid, "a failed fetch must leave the unknown unknown")
}