at-container-registry/internal/testharness/harness.go

//go:build testmode

package testharness

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/bluesky-social/indigo/atproto/identity"
	"github.com/bluesky-social/indigo/atproto/syntax"
	"github.com/distribution/distribution/v3/configuration"
	"github.com/google/go-containerregistry/pkg/authn"
	"github.com/johannesboyne/gofakes3"
	"github.com/johannesboyne/gofakes3/backend/s3mem"

	"atcr.io/pkg/appview"
	"atcr.io/pkg/appview/registryauth"
	"atcr.io/pkg/atproto"
	atprotodid "atcr.io/pkg/atproto/did"
	"atcr.io/pkg/billing"
	"atcr.io/pkg/hold"
	"atcr.io/pkg/hold/quota"
	"atcr.io/pkg/testpds"
)

// Option configures Harness construction. Use New(t, WithX(...)) to set them.
type Option func(*options)

type options struct {
	tap         func(backend string, next http.Handler) http.Handler
	quota       *quota.Config
	billing     *billing.Config
	privateHold bool
}

// WithBackendTap wraps every backend the stack talks to: the fake PDS handler
// ("pds"), the gofakes3 handler ("s3"), and a reverse proxy placed in front of
// the hold ("hold") so the appview's XRPC calls pass through it. Benchmarks
// use it to count round trips per operation and to inject latency.
func WithBackendTap(tap func(backend string, next http.Handler) http.Handler) Option {
	return func(o *options) { o.tap = tap }
}

// WithPrivateHold builds the hold with captain.Public = false. Reads then
// require the owner or a crew member, so anonymous pulls are refused and a
// PDS-known stranger is refused too — the mirror of the default public hold,
// where anyone may pull and only crew may push.
func WithPrivateHold() Option {
	return func(o *options) {
		o.privateHold = true
	}
}

// WithQuotaTiers configures the hold's quota manager with the given tier
// definitions. NewCrewTier names the tier applied to new crew members. Pass
// a tiny limit (e.g. "1KB") plus NewCrewTier="tiny" to make a single push
// exhaust the quota for a non-captain user, so the next /auth/token gate
// denies with "quota exceeded".
func WithQuotaTiers(tiers []quota.TierConfig, newCrewTier string) Option {
	return func(o *options) {
		o.quota = &quota.Config{
			Tiers:    tiers,
			Defaults: quota.DefaultsConfig{NewCrewTier: newCrewTier},
		}
	}
}

// WithBilling wires a billing.Config into the appview built by the harness.
// Only meaningful under `-tags billing`: without the tag the billing package
// compiles to no-op stubs and Manager.Enabled() stays false regardless of
// what's set here.
func WithBilling(cfg billing.Config) Option {
	return func(o *options) {
		o.billing = &cfg
	}
}

// Harness owns all in-process servers and tears them down on test cleanup.
type Harness struct {
	t          *testing.T
	PDS        *testpds.Server
	S3URL      string
	HoldDID    string
	HoldURL    string
	AppViewURL string // 127.0.0.1:PORT — host used for /v2/ registry requests
	UIBaseURL  string // localhost:PORT — host used for UI/API routes (e.g. /api/stripe/webhook)
	AppView    *appview.AppViewServer
	Hold       *hold.HoldServer
	Captain    *Sailor // hold owner; set before appview boots
}

// Sailor combines the fake-PDS identity with the synthetic OAuth bits AppView
// needs to mint a registry JWT. Today the harness drives /auth/token via the
// app-password Basic auth path, so callers don't have to think about OAuth.
type Sailor struct {
	Identity *testpds.Identity
}

// DID returns the sailor's DID.
func (s *Sailor) DID() string { return s.Identity.DID.String() }

// Handle returns the sailor's handle.
func (s *Sailor) Handle() string { return s.Identity.Handle.String() }

// New brings up the full stack on random localhost ports and registers
// t.Cleanup to tear everything down. The bucket name used inside gofakes3 is
// fixed; the test never cares about its value.
func New(t *testing.T, opts ...Option) *Harness {
	t.Helper()
	var o options
	for _, opt := range opts {
		opt(&o)
	}
	h := &Harness{t: t}

	// 1. Fake PDS.
	var pdsOpts []testpds.Option
	if o.tap != nil {
		pdsOpts = append(pdsOpts, testpds.WithMiddleware(func(next http.Handler) http.Handler {
			return o.tap("pds", next)
		}))
	}
	h.PDS = testpds.New(t, pdsOpts...)
	atproto.SetDirectory(h.PDS.Directory())
	t.Cleanup(func() {
		// Reset to a fresh default so a later non-test process won't see our
		// fake. SetDirectory(nil) re-arms lazy init in GetDirectory.
		atproto.SetDirectory(nil)
	})

	// 2. gofakes3 (S3-compatible in-memory).
	backend := s3mem.New()
	if err := backend.CreateBucket("atcr-test"); err != nil {
		t.Fatalf("create test bucket: %v", err)
	}
	faker := gofakes3.New(backend)
	var s3Handler = faker.Server()
	if o.tap != nil {
		s3Handler = o.tap("s3", s3Handler)
	}
	s3ts := httptest.NewServer(s3Handler)
	t.Cleanup(s3ts.Close)
	h.S3URL = s3ts.URL

	// 3. Hold.
	holdListener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("hold listen: %v", err)
	}
	holdAddr := holdListener.Addr().String()
	holdPublicURL := "http://" + holdAddr
	if o.tap != nil {
		// Everything that addresses the hold by its public URL (the appview,
		// the DID document, the did:web itself) goes through a counting
		// reverse proxy; the hold keeps serving on its own listener behind it.
		proxyListener, err := net.Listen("tcp", "127.0.0.1:0")
		if err != nil {
			t.Fatalf("hold proxy listen: %v", err)
		}
		target, _ := url.Parse("http://" + holdAddr)
		proxy := httputil.NewSingleHostReverseProxy(target)
		proxySrv := &http.Server{Handler: o.tap("hold", proxy)}
		go func() { _ = proxySrv.Serve(proxyListener) }()
		t.Cleanup(func() { _ = proxySrv.Close() })
		holdPublicURL = "http://" + proxyListener.Addr().String()
	}
	h.HoldURL = holdPublicURL
	h.HoldDID = atprotodid.GenerateDIDFromURL(holdPublicURL)

	// Owner DID for the hold. The captain identity lives on the fake PDS; we
	// register it BEFORE constructing the hold so any synchronous bootstrap
	// step that needs to resolve the owner finds it. The captain is exposed
	// as h.Captain so tests that push to their own hold can use it directly.
	captainIdent, err := h.PDS.AddIdentity("captain.test")
	if err != nil {
		t.Fatalf("add captain: %v", err)
	}
	h.Captain = &Sailor{Identity: captainIdent}

	// Pre-register hold's own did:web in the directory so AppView can resolve
	// it without hitting the network (we still serve /.well-known/did.json
	// from the hold for parity, but the in-memory directory wins).
	h.PDS.Directory().Register(&identity.Identity{
		DID:    syntax.DID(h.HoldDID),
		Handle: syntax.HandleInvalid,
		Services: map[string]identity.ServiceEndpoint{
			"atproto_pds": {Type: "AtprotoPersonalDataServer", URL: holdPublicURL},
			"atcr_hold":   {Type: "AtcrHoldService", URL: holdPublicURL},
		},
		Keys: map[string]identity.VerificationMethod{},
	})

	holdTmp := t.TempDir()
	holdCfg := &hold.Config{
		LogLevel: "warn",
		Storage: hold.StorageConfig{
			AccessKey: "test", SecretKey: "test", Region: "us-east-1",
			Bucket: "atcr-test", Endpoint: h.S3URL,
		},
		Server: hold.ServerConfig{
			Addr:         holdAddr,
			PublicURL:    holdPublicURL,
			Public:       !o.privateHold, // public: anyone may pull, crew may push. private: crew only, both ways.
			ReadTimeout:  60 * time.Second,
			WriteTimeout: 5 * time.Minute,
		},
		Registration: hold.RegistrationConfig{
			OwnerDID:           captainIdent.DID.String(),
			AllowAllCrew:       true, // lets the appview auto-crew flow add sailors on first push
			ProfileDisplayName: "Test Captain",
			ProfileDescription: "harness-owned hold",
		},
		Database: hold.DatabaseConfig{
			// Real directory rather than :memory: — libsql's connection pool
			// opens a fresh in-memory DB per connection, so schemas created
			// on one don't appear on others. A tempdir sidesteps the issue
			// and gets cleaned up automatically by t.TempDir().
			Path:      holdTmp,
			KeyPath:   filepath.Join(holdTmp, "signing.key"),
			DIDMethod: "web",
		},
	}

	if o.quota != nil {
		holdCfg.Quota = *o.quota
	}

	holdSrv, err := hold.NewHoldServer(holdCfg)
	if err != nil {
		holdListener.Close()
		t.Fatalf("new hold: %v", err)
	}
	h.Hold = holdSrv

	holdDone := make(chan struct{})
	go func() {
		_ = holdSrv.ServeWithListener(holdListener)
		close(holdDone)
	}()
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = holdSrv.Shutdown(ctx)
		<-holdDone
	})

	if err := waitForHTTP(holdPublicURL+"/.well-known/did.json", 5*time.Second); err != nil {
		t.Fatalf("hold did doc not reachable: %v", err)
	}

	// 4. AppView. The DomainRoutingMiddleware compares the request's host
	// against RegistryDomains (port-stripped) to decide whether to allow /v2/
	// or redirect. RegistryDomains must therefore be a bare hostname (no
	// port), AND must differ from the UI hostname extracted from BaseURL.
	// We bind to 127.0.0.1, route /v2/ via that, and use "localhost" as the
	// UI hostname (same socket, different name) so the routing branches
	// don't collide.
	avListener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("appview listen: %v", err)
	}
	avAddr := avListener.Addr().String()
	_, avPort, err := net.SplitHostPort(avAddr)
	if err != nil {
		t.Fatalf("split appview addr: %v", err)
	}
	avBaseURL := "http://localhost:" + avPort
	h.AppViewURL = "http://" + avAddr
	h.UIBaseURL = avBaseURL

	avDBPath := filepath.Join(t.TempDir(), "appview.db")
	avCfg := buildAppViewConfig(avAddr, avBaseURL, h.HoldDID, avDBPath)
	if o.billing != nil {
		avCfg.Billing = *o.billing
	}
	avSrv, err := appview.NewAppViewServer(avCfg, nil)
	if err != nil {
		avListener.Close()
		t.Fatalf("new appview: %v", err)
	}
	h.AppView = avSrv

	avDone := make(chan struct{})
	go func() {
		_ = avSrv.ServeWithListener(avListener)
		close(avDone)
	}()
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = avSrv.Shutdown(ctx)
		<-avDone
	})

	if err := waitForHTTP(avBaseURL+"/healthz", 5*time.Second); err != nil {
		// Not fatal: some appview versions don't expose /healthz at this path.
		// Fall back to a TCP probe just to be sure the listener is up.
		if err2 := waitForTCP(avAddr, 5*time.Second); err2 != nil {
			t.Fatalf("appview not reachable: %v / %v", err, err2)
		}
	}

	// Seed the appview's local tables that the /auth/token authorizer reads.
	// In production these are populated by the Jetstream worker consuming
	// hold firehose events; tests disable Jetstream and seed directly so
	// authorization decisions resolve against known state.
	h.seedCaptainRecord(captainIdent, !o.privateHold)
	h.seedUserRow(captainIdent)
	h.seedCrewMember(captainIdent, []string{"blob:read", "blob:write", "crew:admin"})

	return h
}

// AddSailor creates a new identity on the fake PDS and registers it as a
// crew member of the harness's hold with blob:write permission. The returned
// Sailor is immediately usable for push and pull: AppView mints registry
// JWTs via the app-password Basic auth path on /auth/token, and the auth
// gate finds the seeded crew row in its local table.
func (h *Harness) AddSailor(handle string) *Sailor {
	return h.AddSailorWithPermissions(handle, []string{"blob:read", "blob:write"})
}

// AddSailorWithPermissions is like AddSailor but lets the caller choose the
// permissions written to the appview's hold_crew_members table. Use
// []string{"blob:read"} for a read-only crew member who can pull but not push.
func (h *Harness) AddSailorWithPermissions(handle string, permissions []string) *Sailor {
	h.t.Helper()
	ident, err := h.PDS.AddIdentity(handle)
	if err != nil {
		h.t.Fatalf("add sailor %q: %v", handle, err)
	}
	h.seedUserRow(ident)
	h.seedCrewMember(ident, permissions)
	return &Sailor{Identity: ident}
}

// AddStranger creates an identity on the fake PDS and seeds the appview's
// users row (so PDS resolution works for the auth-token service-auth pre-mint)
// but does NOT add a crew_members row. The returned sailor authenticates fine
// but is not a member of the hold — push token requests should be denied with
// "crew membership required", while pull token requests succeed because the
// gate's membership requirement is push-only.
func (h *Harness) AddStranger(handle string) *Sailor {
	h.t.Helper()
	ident, err := h.PDS.AddIdentity(handle)
	if err != nil {
		h.t.Fatalf("add stranger %q: %v", handle, err)
	}
	h.seedUserRow(ident)
	return &Sailor{Identity: ident}
}

// AnonAuth returns the anonymous authenticator. crane uses it when no
// credentials are configured; the token endpoint requires Basic auth, so the
// resulting /v2/* requests fail with 401 unauthorized.
func (h *Harness) AnonAuth() authn.Authenticator {
	return authn.Anonymous
}

// seedCaptainRecord writes a row to hold_captain_records so the auth gate's
// isCaptain check returns true for the captain.
//
// public must match the hold's own captain.Public. This row is what the
// appview's hold authorizer reads, so seeding it true against a private hold
// would let the appview admit reads the hold then refuses, which is the
// disagreement these tests exist to catch.
func (h *Harness) seedCaptainRecord(captain *testpds.Identity, public bool) {
	h.t.Helper()
	_, err := h.AppView.Database.Exec(
		`INSERT INTO hold_captain_records (hold_did, owner_did, public, allow_all_crew) VALUES (?, ?, ?, ?)`,
		h.HoldDID, captain.DID.String(), public, true,
	)
	if err != nil {
		h.t.Fatalf("seed captain record: %v", err)
	}
}

// seedUserRow writes a row to users so the appview's lookups (PDS endpoint
// for the user, default hold, etc.) resolve locally. The hold_resolver in
// authgate reads default_hold_did from this table.
func (h *Harness) seedUserRow(ident *testpds.Identity) {
	h.t.Helper()
	_, err := h.AppView.Database.Exec(
		`INSERT OR REPLACE INTO users (did, handle, pds_endpoint, default_hold_did, last_seen)
		 VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)`,
		ident.DID.String(), ident.Handle.String(), h.PDS.URL(), h.HoldDID,
	)
	if err != nil {
		h.t.Fatalf("seed user row: %v", err)
	}
}

// seedCrewMember writes a row to hold_crew_members so checkCrewBlobWrite
// finds the member without waiting on Jetstream-fed updates.
func (h *Harness) seedCrewMember(ident *testpds.Identity, permissions []string) {
	h.t.Helper()
	permsJSON, err := json.Marshal(permissions)
	if err != nil {
		h.t.Fatalf("marshal permissions: %v", err)
	}
	_, err = h.AppView.Database.Exec(
		`INSERT OR REPLACE INTO hold_crew_members
		 (hold_did, member_did, rkey, role, permissions, added_at)
		 VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)`,
		h.HoldDID, ident.DID.String(), "test-"+ident.Handle.String(), "member", string(permsJSON),
	)
	if err != nil {
		h.t.Fatalf("seed crew member: %v", err)
	}
}

// AppViewHostPort returns the host:port the registry was bound to, for use
// as the registry component of an OCI image reference.
func (h *Harness) AppViewHostPort() string {
	return strings.TrimPrefix(h.AppViewURL, "http://")
}

// RegistryAuth returns an authn.Authenticator that performs the Docker token
// dance against this harness. Handle is used as the username because
// production's parseBasicAuthDID() reconstructs DIDs assuming the canonical
// 2-segment did:web shape; our synthesized DIDs are longer
// (did:web:host:user:alice). Handles resolve through the fake directory just
// as well as DIDs do.
func (h *Harness) RegistryAuth(s *Sailor) authn.Authenticator {
	return &authn.Basic{Username: s.Handle(), Password: s.Identity.Password}
}

// Auth carries neutral basic-auth credentials for the test OCI client
// abstraction. Both Username and Password empty means anonymous.
type Auth struct {
	Username, Password string
}

// RegistryCreds returns the sailor's basic-auth credentials as a neutral Auth
// value, for client-library-agnostic test code.
func (h *Harness) RegistryCreds(s *Sailor) Auth {
	return Auth{Username: s.Handle(), Password: s.Identity.Password}
}

// AnonCreds returns the anonymous (empty) Auth.
func (h *Harness) AnonCreds() Auth { return Auth{} }

// --- helpers ---------------------------------------------------------------

func buildAppViewConfig(addr, baseURL, holdDID, dbPath string) *appview.Config {
	cfg := appview.DefaultConfig()
	cfg.LogLevel = "warn"
	cfg.Server.Addr = addr
	cfg.Server.BaseURL = baseURL
	cfg.Server.ManagedHolds = []string{holdDID}
	// Registry domain is a bare hostname (no port). DomainRoutingMiddleware
	// strips ports before matching, so "127.0.0.1" is what /v2/ requests
	// will hit (since the listener binds to 127.0.0.1). BaseURL uses
	// "localhost" for the UI hostname so the two branches differ.
	cfg.Server.RegistryDomains = []string{"127.0.0.1"}
	// Real file path under t.TempDir(). `:memory:` is rejected because
	// libsql's connection pool opens multiple distinct in-memory databases,
	// so the schema applied by InitDB doesn't reach subsequent connections.
	cfg.UI.DatabasePath = dbPath
	cfg.UI.LibsqlSyncURL = ""
	// Disable jetstream so the test doesn't open WebSockets to the public network.
	cfg.Jetstream.URLs = []string{}
	cfg.Jetstream.BackfillEnabled = false
	cfg.Jetstream.RelayEndpoints = []string{}

	cfg.Auth.TokenExpiration = 5 * time.Minute
	cfg.Auth.Services = cfg.Server.RegistryDomains
	cfg.Auth.CertPath = filepath.Join(os.TempDir(), fmt.Sprintf("atcr-test-cert-%d.pem", time.Now().UnixNano()))
	cfg.Distribution = buildDistributionConfig(addr, baseURL, holdDID, cfg.Auth.Services, cfg.Auth.CertPath)
	return cfg
}

func buildDistributionConfig(addr, baseURL, holdDID string, services []string, certPath string) *configuration.Configuration {
	serviceName := services[0]
	distConfig := &configuration.Configuration{}
	distConfig.Version = configuration.MajorMinorVersion(0, 1)
	distConfig.Log = configuration.Log{
		Level:     configuration.Loglevel("warn"),
		Formatter: "text",
		Fields:    map[string]any{"service": "atcr-appview"},
	}
	distConfig.HTTP = configuration.HTTP{
		Addr:   addr,
		Secret: "test-http-secret-do-not-use-in-prod",
		Headers: map[string][]string{
			"X-Content-Type-Options": {"nosniff"},
		},
	}
	distConfig.Storage = configuration.Storage{
		"inmemory": configuration.Parameters{},
		"maintenance": configuration.Parameters{
			"uploadpurging": map[any]any{
				"enabled":  false,
				"age":      7 * 24 * time.Hour,
				"interval": 24 * time.Hour,
				"dryrun":   false,
			},
		},
		// Mirror buildStorageConfig: distribution v3.1.1's DeleteManifest
		// handler returns UNSUPPORTED unless delete is enabled here.
		"delete": configuration.Parameters{"enabled": true},
	}
	distConfig.Middleware = map[string][]configuration.Middleware{
		"registry": {{
			Name: "atproto-resolver",
			Options: configuration.Parameters{
				"default_hold_did": holdDID,
				"base_url":         baseURL,
			},
		}},
	}
	distConfig.Auth = configuration.Auth{
		registryauth.AuthType: configuration.Parameters{
			"realm":          baseURL + "/auth/token",
			"services":       services,
			"issuer":         serviceName,
			"rootcertbundle": certPath,
			"expiration":     int((5 * time.Minute).Seconds()),
		},
	}
	distConfig.Health = configuration.Health{
		StorageDriver: configuration.StorageDriver{
			Enabled:   false,
			Interval:  10 * time.Second,
			Threshold: 3,
		},
	}
	return distConfig
}

// waitForHTTP polls a URL until it returns any response or timeout elapses.
func waitForHTTP(u string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		resp, err := http.Get(u)
		if err == nil {
			resp.Body.Close()
			return nil
		}
		time.Sleep(25 * time.Millisecond)
	}
	return fmt.Errorf("timed out waiting for %s", u)
}

func waitForTCP(addr string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		c, err := net.DialTimeout("tcp", addr, 250*time.Millisecond)
		if err == nil {
			c.Close()
			return nil
		}
		time.Sleep(25 * time.Millisecond)
	}
	return fmt.Errorf("timed out waiting for tcp %s", addr)
}