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