at-container-registry/pkg/auth/servicetoken_testmode_test.go

//go:build testmode

// These tests drive the real OAuth client against a fake PDS on loopback, which only a testmode build can reach.
package auth

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

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

// setupServiceTokenScenario builds a fake PDS whose getServiceAuth endpoint
// always 401s with invalid_token (forcing a token refresh) and whose token
// endpoint behavior is supplied by the caller, plus a Refresher seeded with a
// session for did.
func setupServiceTokenScenario(t *testing.T, did string, tokenHandler http.HandlerFunc) (*oauth.Refresher, *stubAuthStore, *recordingUIStore, string) {
	t.Helper()

	mux := http.NewServeMux()
	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)

	mux.HandleFunc(atproto.ServerGetServiceAuth, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("WWW-Authenticate", `Bearer error="invalid_token", error_description="expired"`)
		w.WriteHeader(http.StatusUnauthorized)
	})
	mux.HandleFunc("/oauth/token", tokenHandler)

	key, err := atcrypto.GeneratePrivateKeyP256()
	if err != nil {
		t.Fatal(err)
	}
	parsedDID, err := syntax.ParseDID(did)
	if err != nil {
		t.Fatal(err)
	}

	store := newStubAuthStore()
	store.sessions[did] = indigo_oauth.ClientSessionData{
		AccountDID:              parsedDID,
		SessionID:               "test-session",
		HostURL:                 srv.URL,
		AuthServerURL:           srv.URL,
		AuthServerTokenEndpoint: srv.URL + "/oauth/token",
		Scopes:                  []string{"atproto"},
		AccessToken:             "old-access",
		RefreshToken:            "old-refresh",
		DPoPPrivateKeyMultibase: key.Multibase(),
	}

	clientApp, err := oauth.NewClientApp("http://localhost:5000", store, []string{"atproto"}, "", "test")
	if err != nil {
		t.Fatal(err)
	}
	refresher := oauth.NewRefresher(clientApp)
	uiStore := &recordingUIStore{}
	refresher.SetUISessionStore(uiStore)

	return refresher, store, uiStore, srv.URL
}

// TestGetOrFetchServiceToken_TransientErrorKeepsSession asserts that a
// transient refresh failure (5xx from the auth server) does NOT delete the
// OAuth session or invalidate UI sessions.
func TestGetOrFetchServiceToken_TransientErrorKeepsSession(t *testing.T) {
	const did = "did:web:transient.example.com"
	refresher, store, uiStore, pdsURL := setupServiceTokenScenario(t, did,
		func(w http.ResponseWriter, r *http.Request) {
			http.Error(w, "boom", http.StatusInternalServerError)
		})

	_, err := GetOrFetchServiceToken(context.Background(), refresher, did, "did:web:hold1.example.com", pdsURL)
	if err == nil {
		t.Fatal("expected error from failed service token fetch")
	}
	if deleted := store.deletedDIDs(); len(deleted) != 0 {
		t.Errorf("session deleted on transient error: %v", deleted)
	}
	if len(uiStore.deleted) != 0 {
		t.Errorf("UI sessions invalidated on transient error: %v", uiStore.deleted)
	}
}

// TestGetOrFetchServiceToken_InvalidGrantDeletesSession asserts that a
// genuine invalid_grant (refresh token replayed/revoked) still deletes the
// session and invalidates UI sessions.
func TestGetOrFetchServiceToken_InvalidGrantDeletesSession(t *testing.T) {
	const did = "did:web:invalidgrant.example.com"
	refresher, store, uiStore, pdsURL := setupServiceTokenScenario(t, did,
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusBadRequest)
			_, _ = w.Write([]byte(`{"error":"invalid_grant","error_description":"Refresh token replayed"}`))
		})

	_, err := GetOrFetchServiceToken(context.Background(), refresher, did, "did:web:hold2.example.com", pdsURL)
	if err == nil {
		t.Fatal("expected error from invalid_grant")
	}
	if deleted := store.deletedDIDs(); len(deleted) != 1 || deleted[0] != did {
		t.Errorf("expected session deletion for %s, got: %v", did, deleted)
	}
	if len(uiStore.deleted) != 1 || uiStore.deleted[0] != did {
		t.Errorf("expected UI session invalidation for %s, got: %v", did, uiStore.deleted)
	}
}

// stubAuthStore is a minimal indigo ClientAuthStore + GetLatestSessionForDID
// implementation for driving a real Refresher against httptest servers.
type stubAuthStore struct {
	mu       sync.Mutex
	sessions map[string]indigo_oauth.ClientSessionData
	deleted  []string
}

func newStubAuthStore() *stubAuthStore {
	return &stubAuthStore{sessions: make(map[string]indigo_oauth.ClientSessionData)}
}

func (s *stubAuthStore) GetSession(ctx context.Context, did syntax.DID, sessionID string) (*indigo_oauth.ClientSessionData, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	sess, ok := s.sessions[did.String()]
	if !ok {
		return nil, fmt.Errorf("session not found")
	}
	return &sess, nil
}

func (s *stubAuthStore) SaveSession(ctx context.Context, sess indigo_oauth.ClientSessionData) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.sessions[sess.AccountDID.String()] = sess
	return nil
}

func (s *stubAuthStore) DeleteSession(ctx context.Context, did syntax.DID, sessionID string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.sessions, did.String())
	s.deleted = append(s.deleted, did.String())
	return nil
}

func (s *stubAuthStore) GetAuthRequestInfo(ctx context.Context, state string) (*indigo_oauth.AuthRequestData, error) {
	return nil, fmt.Errorf("not implemented")
}
func (s *stubAuthStore) SaveAuthRequestInfo(ctx context.Context, info indigo_oauth.AuthRequestData) error {
	return nil
}
func (s *stubAuthStore) DeleteAuthRequestInfo(ctx context.Context, state string) error {
	return nil
}

func (s *stubAuthStore) GetLatestSessionForDID(ctx context.Context, did string) (*indigo_oauth.ClientSessionData, string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	sess, ok := s.sessions[did]
	if !ok {
		return nil, "", fmt.Errorf("no session for DID")
	}
	return &sess, sess.SessionID, nil
}

func (s *stubAuthStore) deletedDIDs() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]string(nil), s.deleted...)
}

type recordingUIStore struct {
	mu      sync.Mutex
	deleted []string
}

func (s *recordingUIStore) Create(did, handle, pdsEndpoint string, duration time.Duration) (string, error) {
	return "", fmt.Errorf("not implemented")
}

func (s *recordingUIStore) DeleteByDID(did string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.deleted = append(s.deleted, did)
}