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

package auth

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"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"
)

func TestGetOrFetchServiceToken_NilRefresher(t *testing.T) {
	ctx := context.Background()
	did := "did:plc:test123"
	holdDID := "did:web:hold.example.com"
	pdsEndpoint := "https://pds.example.com"

	// Test with nil refresher - should return error
	_, err := GetOrFetchServiceToken(ctx, nil, did, holdDID, pdsEndpoint)
	if err == nil {
		t.Error("Expected error when refresher is nil")
	}

	expectedErrMsg := "refresher is nil"
	if err.Error() != "refresher is nil (OAuth session required for service tokens)" {
		t.Errorf("Expected error message to contain %q, got %q", expectedErrMsg, err.Error())
	}
}

// Note: Full tests with mocked OAuth refresher and HTTP client will be added
// in the comprehensive test implementation phase

// ----------------------------------------------------------------------------
// Session-deletion gating tests (refresh-cancellation incident regression)
// ----------------------------------------------------------------------------

// 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)
}

// 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)
	}
}

func TestIsStaleBearerToken(t *testing.T) {
	stale := []string{"ExpiredToken", "InvalidToken", "ExpiredSession", "InvalidSession"}
	for _, name := range stale {
		if !isStaleBearerToken(name) {
			t.Errorf("isStaleBearerToken(%q) = false, want true", name)
		}
	}
	// Anything else must not discard the credential — an unrelated server-side
	// failure should not force the user to re-authenticate.
	for _, name := range []string{"", "InvalidRequest", "RateLimitExceeded", "InternalServerError"} {
		if isStaleBearerToken(name) {
			t.Errorf("isStaleBearerToken(%q) = true, want false", name)
		}
	}
}

// TestAppPasswordServiceToken_EvictsOnStaleToken is the regression test for a
// wedged-account loop: a PDS reporting an expired bearer token as 400 with an
// atproto error name (rather than 401) left the dead token in the cache, so
// every subsequent request replayed it. One account produced 16,110 such errors
// and 4,254 retryable 503s over 33 hours, recovering only on process restart.
func TestAppPasswordServiceToken_EvictsOnStaleToken(t *testing.T) {
	tests := []struct {
		name        string
		status      int
		body        string
		wantEvicted bool
	}{
		{
			name:        "400 with ExpiredToken evicts",
			status:      http.StatusBadRequest,
			body:        `{"error":"ExpiredToken","message":"Token has expired"}`,
			wantEvicted: true,
		},
		{
			name:        "401 evicts (pre-existing path)",
			status:      http.StatusUnauthorized,
			body:        `{"error":"AuthMissing"}`,
			wantEvicted: true,
		},
		{
			name:        "400 with an unrelated error keeps the token",
			status:      http.StatusBadRequest,
			body:        `{"error":"InvalidRequest","message":"bad aud"}`,
			wantEvicted: false,
		},
		{
			name:        "500 keeps the token",
			status:      http.StatusInternalServerError,
			body:        `{"error":"InternalServerError"}`,
			wantEvicted: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			pds := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "application/json")
				w.WriteHeader(tt.status)
				_, _ = w.Write([]byte(tt.body))
			}))
			defer pds.Close()

			did := "did:plc:" + strings.ToLower(strings.ReplaceAll(t.Name(), "/", ""))
			holdDID := "did:web:hold.example.com"

			GetGlobalTokenCache().Set(did, "stale-access-token", time.Hour)
			t.Cleanup(func() { GetGlobalTokenCache().Delete(did) })

			if _, err := GetOrFetchServiceTokenWithAppPassword(
				context.Background(), did, holdDID, pds.URL,
			); err == nil {
				t.Fatal("expected an error from a non-200 PDS response")
			}

			_, stillCached := GetGlobalTokenCache().Get(did)
			if tt.wantEvicted && stillCached {
				t.Error("expected the stale app-password token to be evicted, but it is still cached")
			}
			if !tt.wantEvicted && !stillCached {
				t.Error("token was evicted on an unrelated failure; the user is forced to re-authenticate needlessly")
			}
		})
	}
}