at-container-registry/pkg/auth/oauth/client_testmode_test.go

//go:build testmode

// This test drives the real OAuth client against a fake PDS on loopback, which only a testmode build can reach.
package oauth

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

	"github.com/bluesky-social/indigo/atproto/atcrypto"
	"github.com/bluesky-social/indigo/atproto/auth/oauth"
	"github.com/bluesky-social/indigo/atproto/syntax"
)

// seedSession stores a session for did pointing at the given resource server
// and token endpoint, with a freshly generated DPoP key.
func seedSession(t *testing.T, store *fakeAuthStore, did, hostURL, tokenEndpoint string) {
	t.Helper()
	key, err := atcrypto.GeneratePrivateKeyP256()
	if err != nil {
		t.Fatal(err)
	}
	parsedDID, err := syntax.ParseDID(did)
	if err != nil {
		t.Fatal(err)
	}
	store.sessions[did] = oauth.ClientSessionData{
		AccountDID:              parsedDID,
		SessionID:               "test-session",
		HostURL:                 hostURL,
		AuthServerURL:           hostURL,
		AuthServerTokenEndpoint: tokenEndpoint,
		Scopes:                  []string{"atproto"},
		AccessToken:             "old-access",
		RefreshToken:            "old-refresh",
		DPoPPrivateKeyMultibase: key.Multibase(),
	}
}

func newTestRefresher(t *testing.T, store oauth.ClientAuthStore) *Refresher {
	t.Helper()
	clientApp, err := NewClientApp("http://localhost:5000", store, []string{"atproto"}, "", "test")
	if err != nil {
		t.Fatal(err)
	}
	return NewRefresher(clientApp)
}

// TestDoWithSession_RefreshSurvivesRequestCancellation reproduces the
// production incident: the caller's context is canceled while the token
// refresh is in flight (after the auth server has already rotated the
// refresh token). The rotated token MUST be persisted and the session MUST
// NOT be deleted, or the next refresh fails with invalid_grant "Refresh
// token replayed" and the user is signed out.
func TestDoWithSession_RefreshSurvivesRequestCancellation(t *testing.T) {
	const did = "did:web:refresh-cancel.example.com"

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

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

	// Resource endpoint: reject the stale access token so DoWithAuth triggers
	// a refresh; accept the rotated one.
	mux.HandleFunc("/xrpc/test.resource", func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") == "DPoP new-access" {
			w.WriteHeader(http.StatusOK)
			return
		}
		w.Header().Set("WWW-Authenticate", `Bearer error="invalid_token", error_description="expired"`)
		w.WriteHeader(http.StatusUnauthorized)
	})

	// Token endpoint: simulate the Docker client hanging up mid-refresh by
	// canceling the caller's context BEFORE responding, then return rotated
	// tokens (the auth server has already committed the rotation by then).
	mux.HandleFunc("/oauth/token", func(w http.ResponseWriter, r *http.Request) {
		cancel()
		fmt.Fprintf(w, `{"sub":%q,"access_token":"new-access","refresh_token":"new-refresh"}`, did)
	})

	store := newFakeAuthStore()
	seedSession(t, store, did, srv.URL, srv.URL+"/oauth/token")
	refresher := newTestRefresher(t, store)
	uiStore := &spyUISessionStore{}
	refresher.SetUISessionStore(uiStore)

	err := refresher.DoWithSession(ctx, did, func(session *oauth.ClientSession) error {
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/xrpc/test.resource", nil)
		if err != nil {
			return err
		}
		resp, err := session.DoWithAuth(session.Client, req, syntax.NSID("com.atproto.server.getServiceAuth"))
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		return nil
	})

	// The overall operation may fail (the post-refresh retry of the resource
	// request runs on the canceled inbound context) — that is fine, Docker
	// retries. What must hold is that the rotated refresh token was saved and
	// the session survived.
	if got := store.refreshToken(did); got != "new-refresh" {
		t.Errorf("rotated refresh token not persisted: got %q, want %q (op err: %v)", got, "new-refresh", err)
	}
	if deleted := store.deletedDIDs(); len(deleted) != 0 {
		t.Errorf("session was deleted: %v", deleted)
	}
	if len(uiStore.deleted) != 0 {
		t.Errorf("UI session was invalidated: %v", uiStore.deleted)
	}
}

// fakeAuthStore is a ClientAuthStore whose SaveSession honors context
// cancellation, so a test fails if session persistence runs on a canceled
// context. It also implements GetLatestSessionForDID (the sessionGetter
// extension the Refresher requires).
type fakeAuthStore struct {
	mu       sync.Mutex
	sessions map[string]oauth.ClientSessionData // keyed by DID
	deleted  []string
}

func newFakeAuthStore() *fakeAuthStore {
	return &fakeAuthStore{sessions: make(map[string]oauth.ClientSessionData)}
}

func (s *fakeAuthStore) GetSession(ctx context.Context, did syntax.DID, sessionID string) (*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 *fakeAuthStore) SaveSession(ctx context.Context, sess 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 *fakeAuthStore) 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 *fakeAuthStore) GetAuthRequestInfo(ctx context.Context, state string) (*oauth.AuthRequestData, error) {
	return nil, fmt.Errorf("not implemented")
}
func (s *fakeAuthStore) SaveAuthRequestInfo(ctx context.Context, info oauth.AuthRequestData) error {
	return nil
}
func (s *fakeAuthStore) DeleteAuthRequestInfo(ctx context.Context, state string) error {
	return nil
}

func (s *fakeAuthStore) GetLatestSessionForDID(ctx context.Context, did string) (*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 *fakeAuthStore) refreshToken(did string) string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.sessions[did].RefreshToken
}

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

// spyUISessionStore records DeleteByDID calls.
type spyUISessionStore struct {
	mu      sync.Mutex
	deleted []string
}

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

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