internal/mcpgateway/oauth_test.go
Rama main · réplica de lectura del repositorio original.
package mcpgateway
import (
"context"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/modelcontextprotocol/go-sdk/oauthex"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/enterpilot/gomodel/internal/encryption/encryptiontest"
"github.com/enterpilot/gomodel/internal/storage/sqlx/sqlxtest"
)
const testRedirectURI = "http://127.0.0.1:8080/admin/mcp-oauth/callback"
// fakeAuthServer is an OAuth 2.1 authorization server that checks PKCE and
// the resource indicator, issues a short-lived token, and rotates refresh
// tokens.
type fakeAuthServer struct {
*httptest.Server
registration bool
resource string
mu sync.Mutex
challenge string
refreshToken string
accessToken string
refreshes int
}
func newFakeAuthServer(t *testing.T, registration bool) *fakeAuthServer {
t.Helper()
as := &fakeAuthServer{registration: registration}
as.Server = httptest.NewServer(http.HandlerFunc(as.serve))
t.Cleanup(as.Close)
return as
}
func writeTestJSON(w http.ResponseWriter, status int, body any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(body)
}
func (as *fakeAuthServer) serve(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/.well-known/oauth-authorization-server":
meta := map[string]any{
"issuer": as.URL,
"authorization_endpoint": as.URL + "/authorize",
"token_endpoint": as.URL + "/token",
"code_challenge_methods_supported": []string{"S256"},
"scopes_supported": []string{"mcp:read", "offline_access"},
"response_types_supported": []string{"code"},
}
if as.registration {
meta["registration_endpoint"] = as.URL + "/register"
}
writeTestJSON(w, http.StatusOK, meta)
case "/register":
var meta map[string]any
_ = json.NewDecoder(r.Body).Decode(&meta)
writeTestJSON(w, http.StatusCreated, map[string]any{
"client_id": "dcr-client",
"redirect_uris": meta["redirect_uris"],
"token_endpoint_auth_method": "none",
})
case "/token":
as.token(w, r)
default:
http.NotFound(w, r)
}
}
func (as *fakeAuthServer) token(w http.ResponseWriter, r *http.Request) {
if err := r.ParseForm(); err != nil {
writeTestJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid_request"})
return
}
as.mu.Lock()
defer as.mu.Unlock()
if r.PostForm.Get("resource") != as.resource {
writeTestJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid_target"})
return
}
switch r.PostForm.Get("grant_type") {
case "authorization_code":
sum := sha256.Sum256([]byte(r.PostForm.Get("code_verifier")))
if r.PostForm.Get("code") != "code-1" || base64.RawURLEncoding.EncodeToString(sum[:]) != as.challenge {
writeTestJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid_grant"})
return
}
// Expires inside oauth2's expiry delta, so the first use refreshes.
as.accessToken, as.refreshToken = "access-1", "refresh-1"
writeTestJSON(w, http.StatusOK, map[string]any{
"access_token": "access-1", "refresh_token": "refresh-1", "token_type": "Bearer", "expires_in": 1, "scope": "mcp:read offline_access",
})
case "refresh_token":
if r.PostForm.Get("refresh_token") != as.refreshToken {
writeTestJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid_grant", "error_description": "refresh token was rotated"})
return
}
as.refreshes++
as.accessToken = "access-2"
as.refreshToken = "refresh-2"
writeTestJSON(w, http.StatusOK, map[string]any{
"access_token": "access-2", "refresh_token": "refresh-2", "token_type": "Bearer", "expires_in": 3600,
})
default:
writeTestJSON(w, http.StatusBadRequest, map[string]string{"error": "unsupported_grant_type"})
}
}
func (as *fakeAuthServer) validAccessToken(token string) bool {
as.mu.Lock()
defer as.mu.Unlock()
return as.accessToken != "" && token == as.accessToken
}
// newProtectedUpstream serves an MCP server at /mcp that answers 401 with a
// resource_metadata challenge unless the request carries a token as issued.
func newProtectedUpstream(t *testing.T, as *fakeAuthServer) string {
t.Helper()
server := mcp.NewServer(&mcp.Implementation{Name: "protected", Version: "test"}, nil)
addEchoTool("echo")(server)
handler := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { return server }, nil)
var ts *httptest.Server
ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/.well-known/oauth-protected-resource/mcp" {
writeTestJSON(w, http.StatusOK, map[string]any{
"resource": ts.URL + "/mcp",
"authorization_servers": []string{as.URL},
"scopes_supported": []string{"mcp:read"},
})
return
}
token, _ := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
if !as.validAccessToken(token) {
w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+ts.URL+`/.well-known/oauth-protected-resource/mcp", scope="mcp:read"`)
http.Error(w, "unauthorized", http.StatusUnauthorized)
return
}
handler.ServeHTTP(w, r)
}))
t.Cleanup(ts.Close)
as.resource = ts.URL + "/mcp"
return ts.URL + "/mcp"
}
func newOAuthTestService(t *testing.T, serverURL string) (*Service, *oauthStore) {
t.Helper()
store, err := NewSQLStore(context.Background(), sqlxtest.NewSQLite(t))
require.NoError(t, err)
box, _ := encryptiontest.NewBox(t)
oauth := newOAuthStore(store, box)
service, _ := newTestServiceWithOptions(t, Options{
ConfigServers: map[string]ServerSpec{"notion": testSpec("notion", serverURL, nil)},
Store: store,
oauthStore: oauth,
})
return service, oauth
}
func serverView(t *testing.T, service *Service, name string) ServerView {
t.Helper()
for _, view := range service.Views() {
if view.Spec.Name == name {
return view
}
}
require.FailNow(t, "server not found", name)
return ServerView{}
}
func waitForStatus(t *testing.T, service *Service, name string, status ServerStatus) ServerView {
t.Helper()
var view ServerView
require.Eventually(t, func() bool {
view = serverView(t, service, name)
return view.Status == status
}, 10*time.Second, 20*time.Millisecond, "server %s did not become %s", name, status)
return view
}
// authorize plays the browser and authorization server: it records the PKCE
// challenge and returns the callback the server would redirect to.
func authorize(t *testing.T, as *fakeAuthServer, authURL string, binding string) OAuthCallback {
t.Helper()
parsed, err := url.Parse(authURL)
require.NoError(t, err)
query := parsed.Query()
as.mu.Lock()
as.challenge = query.Get("code_challenge")
as.mu.Unlock()
return OAuthCallback{State: query.Get("state"), Code: "code-1", Binding: binding}
}
func TestOAuthSignInFlow(t *testing.T) {
as := newFakeAuthServer(t, true)
serverURL := newProtectedUpstream(t, as)
service, store := newOAuthTestService(t, serverURL)
ctx := context.Background()
var notified atomic.Int32
service.SetChangeNotifier(func(context.Context) { notified.Add(1) })
before := serverView(t, service, "notion")
assert.Equal(t, StatusDegraded, before.Status)
assert.True(t, before.AuthRequired, "a 401 handshake marks the server as needing sign-in")
assert.Nil(t, before.OAuth)
start, err := service.StartOAuth(ctx, "notion", OAuthStartRequest{RedirectURI: testRedirectURI})
require.NoError(t, err)
require.NotEmpty(t, start.Binding)
authURL, err := url.Parse(start.AuthorizationURL)
require.NoError(t, err)
query := authURL.Query()
assert.Equal(t, as.URL+"/authorize", authURL.Scheme+"://"+authURL.Host+authURL.Path)
assert.Equal(t, "dcr-client", query.Get("client_id"), "the client came from dynamic registration")
assert.Equal(t, "S256", query.Get("code_challenge_method"))
assert.Equal(t, serverURL, query.Get("resource"))
assert.Equal(t, testRedirectURI, query.Get("redirect_uri"))
assert.Equal(t, "mcp:read offline_access", query.Get("scope"), "challenge scope plus offline_access")
assert.NotEmpty(t, query.Get("state"))
name, err := service.CompleteOAuth(ctx, authorize(t, as, start.AuthorizationURL, start.Binding))
require.NoError(t, err)
assert.Equal(t, "notion", name)
assert.Equal(t, int32(1), notified.Load(), "a sign-in is announced to other gateway instances")
view := waitForStatus(t, service, "notion", StatusConnected)
assert.False(t, view.AuthRequired)
assert.Equal(t, 1, view.ToolCount)
require.NotNil(t, view.OAuth)
assert.Equal(t, OAuthStatusSignedIn, view.OAuth.Status)
assert.Equal(t, as.URL, view.OAuth.Issuer)
assert.Equal(t, []string{"mcp:read", "offline_access"}, view.OAuth.Scopes)
// The short-lived token was refreshed on first use, with the resource
// indicator, and the rotated refresh token was persisted.
as.mu.Lock()
refreshes := as.refreshes
as.mu.Unlock()
assert.GreaterOrEqual(t, refreshes, 1)
stored, err := store.get(ctx, "notion")
require.NoError(t, err)
assert.Equal(t, "refresh-2", stored.RefreshToken)
assert.Equal(t, "access-2", stored.AccessToken)
result, err := service.manager.CallTool(ctx, "notion", "echo", json.RawMessage(`{"x":1}`))
require.NoError(t, err)
assert.False(t, result.IsError)
require.NoError(t, service.SignOutOAuth(ctx, "notion"))
assert.Equal(t, int32(2), notified.Load(), "a sign-out is announced to other gateway instances")
_, err = store.get(ctx, "notion")
require.ErrorIs(t, err, ErrNotFound)
view = waitForStatus(t, service, "notion", StatusDegraded)
assert.Nil(t, view.OAuth)
assert.True(t, view.AuthRequired)
}
func TestOAuthCallbackRejections(t *testing.T) {
as := newFakeAuthServer(t, true)
serverURL := newProtectedUpstream(t, as)
service, store := newOAuthTestService(t, serverURL)
ctx := context.Background()
tests := []struct {
name string
mutate func(*OAuthCallback)
wantErr string
}{
{name: "other browser", mutate: func(cb *OAuthCallback) { cb.Binding = "other" }, wantErr: "different browser"},
{name: "unexpected issuer", mutate: func(cb *OAuthCallback) { cb.Issuer = "https://evil.example" }, wantErr: "unexpected issuer"},
{name: "denied", mutate: func(cb *OAuthCallback) { cb.Code = ""; cb.Error = "access_denied" }, wantErr: "access_denied"},
{name: "wrong code", mutate: func(cb *OAuthCallback) { cb.Code = "forged" }, wantErr: "invalid_grant"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
start, err := service.StartOAuth(ctx, "notion", OAuthStartRequest{RedirectURI: testRedirectURI})
require.NoError(t, err)
callback := authorize(t, as, start.AuthorizationURL, start.Binding)
tt.mutate(&callback)
name, err := service.CompleteOAuth(ctx, callback)
require.ErrorContains(t, err, tt.wantErr)
assert.Equal(t, "notion", name)
_, err = store.get(ctx, "notion")
require.ErrorIs(t, err, ErrNotFound, "nothing is stored")
view := serverView(t, service, "notion")
require.NotNil(t, view.OAuth)
assert.Equal(t, OAuthStatusError, view.OAuth.Status)
assert.Contains(t, view.OAuth.LastError, tt.wantErr)
// The state is single-use.
_, err = service.CompleteOAuth(ctx, callback)
require.ErrorIs(t, err, ErrOAuthFlowUnknown)
})
}
_, err := service.CompleteOAuth(ctx, OAuthCallback{State: "unknown", Code: "code-1"})
require.ErrorIs(t, err, ErrOAuthFlowUnknown)
}
func TestOAuthStartWithoutDynamicRegistration(t *testing.T) {
as := newFakeAuthServer(t, false)
serverURL := newProtectedUpstream(t, as)
service, _ := newOAuthTestService(t, serverURL)
ctx := context.Background()
_, err := service.StartOAuth(ctx, "notion", OAuthStartRequest{RedirectURI: testRedirectURI})
require.ErrorIs(t, err, ErrOAuthClientRequired)
start, err := service.StartOAuth(ctx, "notion", OAuthStartRequest{RedirectURI: testRedirectURI, ClientID: "pre-registered"})
require.NoError(t, err)
authURL, err := url.Parse(start.AuthorizationURL)
require.NoError(t, err)
assert.Equal(t, "pre-registered", authURL.Query().Get("client_id"))
_, err = service.StartOAuth(ctx, "missing", OAuthStartRequest{RedirectURI: testRedirectURI})
require.ErrorIs(t, err, ErrNotFound)
_, err = service.StartOAuth(ctx, "notion", OAuthStartRequest{RedirectURI: "https://evil.example/steal"})
require.ErrorIs(t, err, ErrOAuthInvalidRequest)
}
func TestValidateOAuthRedirectURI(t *testing.T) {
tests := []struct {
uri string
valid bool
}{
{uri: "https://gomodel.example.com/admin/mcp-oauth/callback", valid: true},
{uri: "http://localhost:8080/admin/mcp-oauth/callback", valid: true},
{uri: "https://example.com/gomodel/admin/mcp-oauth/callback", valid: true},
{uri: "/admin/mcp-oauth/callback"},
{uri: "ftp://example.com/admin/mcp-oauth/callback"},
{uri: "https://example.com/admin/mcp-oauth/callback?next=x"},
{uri: "https://example.com/admin/mcp-oauth/callback#x"},
{uri: "https://user:pass@example.com/admin/mcp-oauth/callback"},
{uri: "https://example.com/other"},
}
for _, tt := range tests {
t.Run(tt.uri, func(t *testing.T) {
err := ValidateOAuthRedirectURI(tt.uri)
if tt.valid {
require.NoError(t, err)
return
}
require.ErrorIs(t, err, ErrOAuthInvalidRequest)
})
}
}
func TestValidateAuthorizationIssuer(t *testing.T) {
tests := []struct {
name string
iss string
supported bool
wantErr bool
}{
{name: "advertised and matching", iss: "https://as.example", supported: true},
{name: "advertised but absent", supported: true, wantErr: true},
{name: "advertised and different", iss: "https://other.example", supported: true, wantErr: true},
{name: "not advertised and absent"},
{name: "not advertised but matching", iss: "https://as.example"},
{name: "not advertised and different", iss: "https://other.example", wantErr: true},
{name: "no normalization", iss: "https://as.example/", wantErr: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := validateAuthorizationIssuer(tt.iss, "https://as.example", tt.supported)
if tt.wantErr {
require.Error(t, err)
return
}
require.NoError(t, err)
})
}
}
func TestProtectedResourceMetadataURLs(t *testing.T) {
tests := []struct {
name string
metadataURL string
serverURL string
want []prmCandidate
}{
{
name: "path endpoint",
serverURL: "https://mcp.example.com/public/mcp",
want: []prmCandidate{
{url: "https://mcp.example.com/.well-known/oauth-protected-resource/public/mcp", resource: "https://mcp.example.com/public/mcp"},
{url: "https://mcp.example.com/.well-known/oauth-protected-resource", resource: "https://mcp.example.com"},
},
},
{
name: "escaped separator",
serverURL: "https://mcp.example.com/a%2Fb/mcp",
want: []prmCandidate{
{url: "https://mcp.example.com/.well-known/oauth-protected-resource/a%2Fb/mcp", resource: "https://mcp.example.com/a%2Fb/mcp"},
{url: "https://mcp.example.com/.well-known/oauth-protected-resource", resource: "https://mcp.example.com"},
},
},
{
name: "challenge first",
metadataURL: "https://mcp.example.com/prm",
serverURL: "https://mcp.example.com",
want: []prmCandidate{
{url: "https://mcp.example.com/prm", resource: "https://mcp.example.com"},
{url: "https://mcp.example.com/.well-known/oauth-protected-resource", resource: "https://mcp.example.com"},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert.Equal(t, tt.want, protectedResourceMetadataURLs(tt.metadataURL, tt.serverURL))
})
}
}
func TestBearerRoundTripperStaysOnServerOrigin(t *testing.T) {
var seen []string
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
seen = append(seen, req.Header.Get("Authorization"))
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Request: req}, nil
})
record := testOAuthRecord()
transport := &bearerRoundTripper{
base: base,
source: newOAuthSource("notion", record, nil, nil),
origin: "https://mcp.example.com",
}
for _, target := range []string{"https://mcp.example.com/mcp", "https://elsewhere.example.com/mcp"} {
req, err := http.NewRequest(http.MethodGet, target, nil)
require.NoError(t, err)
req.Header.Set("Authorization", "Bearer static")
resp, err := transport.RoundTrip(req)
require.NoError(t, err)
_ = resp.Body.Close()
}
assert.Equal(t, []string{"Bearer access-secret", "Bearer static"}, seen)
}
type roundTripFunc func(*http.Request) (*http.Response, error)
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) }
func TestOAuthTokenReplacesStaticAuthorizationHeader(t *testing.T) {
got := make(chan string, 1)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
got <- r.Header.Get("Authorization")
w.WriteHeader(http.StatusNoContent)
}))
t.Cleanup(server.Close)
record := testOAuthRecord()
record.ServerURL = server.URL
spec := testSpec("notion", server.URL, func(spec *ServerSpec) {
spec.Headers = map[string]string{"Authorization": "Bearer static", "X-Team": "alpha"}
spec.oauth = newOAuthSource("notion", record, nil, nil)
})
u := newUpstream(spec, server.Client())
client := u.dialClient(&connectProbe{method: http.MethodPost}, context.Background())
resp, err := client.Get(server.URL)
require.NoError(t, err)
_ = resp.Body.Close()
assert.Equal(t, "Bearer access-secret", <-got)
}
func TestOAuthAttachDropsKeptSourceWhenURLChanged(t *testing.T) {
record := testOAuthRecord()
manager := newOAuthManager(nil, nil)
current := newOAuthSource("notion", record, nil, nil)
manager.sources["notion"] = current
specs := []ServerSpec{testSpec("notion", record.ServerURL, nil)}
manager.attach(specs, nil, true, nil)
assert.Same(t, current, specs[0].oauth, "an unreadable record keeps the current tokens")
specs = []ServerSpec{testSpec("notion", "https://other.example.com/mcp", nil)}
manager.attach(specs, nil, true, nil)
assert.Nil(t, specs[0].oauth, "tokens never follow a server to a new URL")
}
func TestRegistrationAuthMethod(t *testing.T) {
tests := []struct {
name string
supported []string
want string
}{
{name: "not listed", want: "none"},
{name: "public clients allowed", supported: []string{"client_secret_basic", "none"}, want: "none"},
{name: "basic only", supported: []string{"client_secret_basic"}, want: "client_secret_basic"},
{name: "post only", supported: []string{"client_secret_post"}, want: "client_secret_post"},
{name: "unknown methods", supported: []string{"private_key_jwt"}, want: "none"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert.Equal(t, tt.want, registrationAuthMethod(tt.supported))
})
}
}
func TestResolveOAuthClientComparesIssuersExactly(t *testing.T) {
stored := testOAuthRecord()
stored.Issuer = "https://as.example/"
stored.RedirectURI = testRedirectURI
req := OAuthStartRequest{RedirectURI: testRedirectURI}
asm := &oauthex.AuthServerMeta{Issuer: "https://as.example/"}
client, err := resolveOAuthClient(context.Background(), nil, asm, req, &stored)
require.NoError(t, err)
assert.Equal(t, stored.ClientID, client.id, "the same issuer reuses the stored client")
asm = &oauthex.AuthServerMeta{Issuer: "https://as.example"}
_, err = resolveOAuthClient(context.Background(), nil, asm, req, &stored)
require.ErrorIs(t, err, ErrOAuthClientRequired, "another issuer string never gets the stored client")
}
func TestOAuthSignOutDropsKeptSource(t *testing.T) {
record := testOAuthRecord()
manager := newOAuthManager(nil, nil)
manager.sources["notion"] = newOAuthSource("notion", record, nil, nil)
manager.signedOut("notion")
specs := []ServerSpec{testSpec("notion", record.ServerURL, nil)}
manager.attach(specs, nil, true, nil)
assert.Nil(t, specs[0].oauth, "a reload that cannot read the store keeps no signed-out tokens")
}
func TestOAuthFailureFollowsServerURL(t *testing.T) {
manager := newOAuthManager(nil, nil)
manager.fail("notion", "https://old.example.com/mcp", assert.AnError)
view := manager.view(testSpec("notion", "https://old.example.com/mcp", nil))
require.NotNil(t, view)
assert.Equal(t, OAuthStatusError, view.Status)
assert.Nil(t, manager.view(testSpec("notion", "https://new.example.com/mcp", nil)),
"a failure from the previous URL is not shown")
}
func TestOAuthSecureURL(t *testing.T) {
tests := []struct {
url string
want bool
}{
{url: "https://mcp.example.com/mcp", want: true},
{url: "http://localhost:8080/mcp", want: true},
{url: "http://LOCALHOST:8080/mcp", want: true},
{url: "http://127.0.0.1:8080/mcp", want: true},
{url: "http://[::1]:8080/mcp", want: true},
{url: "http://mcp.example.com/mcp"},
{url: "http://192.168.1.10/mcp"},
{url: "ftp://mcp.example.com/mcp"},
{url: "://bad"},
}
for _, tt := range tests {
t.Run(tt.url, func(t *testing.T) {
assert.Equal(t, tt.want, oauthSecureURL(tt.url))
})
}
}
func TestOAuthAttachSkipsPlainHTTPServers(t *testing.T) {
record := testOAuthRecord()
record.ServerURL = "http://mcp.example.com/mcp"
manager := newOAuthManager(nil, nil)
specs := []ServerSpec{testSpec("notion", record.ServerURL, nil)}
manager.attach(specs, map[string]oauthRecord{"notion": record}, false, nil)
assert.Nil(t, specs[0].oauth, "tokens never go to a non-loopback http server")
}
func TestConnectProbeFollowsRedirectToUnauthorized(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/mcp" {
http.Redirect(w, r, "/mcp/", http.StatusTemporaryRedirect)
return
}
http.Error(w, "unauthorized", http.StatusUnauthorized)
}))
t.Cleanup(server.Close)
u := newUpstream(testSpec("notion", server.URL+"/mcp", nil), server.Client())
probe := &connectProbe{method: http.MethodPost}
client := u.dialClient(probe, context.Background())
resp, err := client.Post(server.URL+"/mcp", "application/json", strings.NewReader("{}"))
require.NoError(t, err)
_ = resp.Body.Close()
assert.True(t, probe.unauthorized(), "the 401 after a redirect counts, not the redirect")
}
func TestOAuthSignOutCancelsPendingSignIn(t *testing.T) {
as := newFakeAuthServer(t, true)
serverURL := newProtectedUpstream(t, as)
service, store := newOAuthTestService(t, serverURL)
ctx := context.Background()
pending, err := service.StartOAuth(ctx, "notion", OAuthStartRequest{RedirectURI: testRedirectURI})
require.NoError(t, err)
callback := authorize(t, as, pending.AuthorizationURL, pending.Binding)
require.NoError(t, service.SignOutOAuth(ctx, "notion"))
_, err = service.CompleteOAuth(ctx, callback)
require.ErrorIs(t, err, ErrOAuthFlowUnknown, "a sign-out drops pending sign-ins")
// A flow taken before the sign-out, whose exchange finishes after it.
running, err := service.StartOAuth(ctx, "notion", OAuthStartRequest{RedirectURI: testRedirectURI})
require.NoError(t, err)
flow, ok := service.oauth.takeFlow(authorize(t, as, running.AuthorizationURL, running.Binding).State)
require.True(t, ok)
service.oauth.signedOut("notion")
err = service.completeOAuth(ctx, flow, OAuthCallback{Code: "code-1", Binding: running.Binding})
require.ErrorContains(t, err, "signed out, removed, or changed")
_, err = store.get(ctx, "notion")
require.ErrorIs(t, err, ErrNotFound, "the late callback does not sign the server back in")
// A sign-out while discovery ran, before the flow was registered.
generation := service.oauth.generation("notion")
service.oauth.signedOut("notion")
spec, ok := service.manager.spec("notion")
require.True(t, ok)
early, err := service.startOAuth(ctx, spec, OAuthStartRequest{RedirectURI: testRedirectURI}, generation)
require.NoError(t, err)
_, err = service.CompleteOAuth(ctx, authorize(t, as, early.AuthorizationURL, early.Binding))
require.ErrorContains(t, err, "signed out, removed, or changed")
_, err = store.get(ctx, "notion")
require.ErrorIs(t, err, ErrNotFound)
}
func TestOAuthAttachDoesNotHoldManagerLockOnBusySource(t *testing.T) {
record := testOAuthRecord()
manager := newOAuthManager(nil, nil)
busy := newOAuthSource("notion", record, nil, nil)
manager.sources["notion"] = busy
busy.mu.Lock() // a token refresh in flight
attached := make(chan struct{})
go func() {
manager.attach([]ServerSpec{testSpec("notion", record.ServerURL, nil)}, map[string]oauthRecord{"notion": record}, false, nil)
close(attached)
}()
answered := make(chan struct{})
go func() {
manager.generation("notion")
close(answered)
}()
select {
case <-answered:
case <-time.After(5 * time.Second):
assert.Fail(t, "the manager lock was held while waiting for a busy source")
}
busy.mu.Unlock()
<-attached
}
func TestOAuthAttachIgnoresRecordsReadBeforeSignOut(t *testing.T) {
record := testOAuthRecord()
manager := newOAuthManager(nil, nil)
// A reload read the record, then a sign-out landed before it attached.
generations := manager.generationsSnapshot()
manager.signedOut("notion")
specs := []ServerSpec{testSpec("notion", record.ServerURL, nil)}
manager.attach(specs, map[string]oauthRecord{"notion": record}, false, generations)
assert.Nil(t, specs[0].oauth, "a stale snapshot never brings signed-out tokens back")
// The next reload reads after the sign-out and attaches normally.
specs = []ServerSpec{testSpec("notion", record.ServerURL, nil)}
manager.attach(specs, map[string]oauthRecord{"notion": record}, false, manager.generationsSnapshot())
assert.NotNil(t, specs[0].oauth)
}