revised OAuthHandler

Author: jba

auth/client.go

@@ -6,137 +6,97 @@ package auth
 
 import (
 	"context"
-	"fmt"
+	"log"
 	"net/http"
 	"sync"
 
 	"github.com/modelcontextprotocol/go-sdk/internal/oauthex"
 	"golang.org/x/oauth2"
-	"golang.org/x/oauth2/authhandler"
 )
 
+// An OAuthHandler conducts an OAuth flow and returns a [oauth2.TokenSource] if the authorization
+// is approved, or an error if not.
+type OAuthHandler func(context.Context, OAuthHandlerArgs) (oauth2.TokenSource, error)
+
+// OAuthHandlerArgs are arguments to an [OAuthHandler].
+type OAuthHandlerArgs struct {
+	// The URL to fetch protected resource metadata, extracted from the WWW-Authenticate header.
+	// Empty if not present or there was an error obtaining it.
+	ResourceMetadataURL string
+}
+
 // HTTPTransport is an [http.RoundTripper] that follows the MCP
 // OAuth protocol when it encounters a 401 Unauthorized response.
 type HTTPTransport struct {
-	mu   sync.Mutex
-	opts HTTPTransportConfig
+	handler OAuthHandler
+	mu      sync.Mutex // protects opts.Base
+	opts    HTTPTransportOptions
 }
 
-func NewHTTPTransport(opts *HTTPTransportConfig) (*HTTPTransport, error) {
+// NewHTTPTransport returns a new [*HTTPTransport].
+// The handler is invoked when an HTTP request results in a 401 Unauthorized status.
+// It is called only once per transport. Once a TokenSource is obtained, it is used
+// for the lifetime of the transport; subsequent 401s are not processed.
+func NewHTTPTransport(handler OAuthHandler, opts *HTTPTransportOptions) (*HTTPTransport, error) {
 	t := &HTTPTransport{}
 	if opts != nil {
 		t.opts = *opts
 	}
 	if t.opts.Base == nil {
 		t.opts.Base = http.DefaultTransport
 	}
-	if t.opts.OAuthClient == nil {
-		t.opts.OAuthClient = http.DefaultClient
-	}
 	return t, nil
 }
 
-type HTTPTransportConfig struct {
-	// OAuthHandler is conducts the OAuth flow, using information obtained from the
-	// MCP server and the auth server that it refers to.
-	OAuthHandler func(context.Context, *oauthex.ProtectedResourceMetadata, *oauthex.AuthServerMeta) (oauth2.TokenSource, error)
-	// Base is the [http.RoundTripper] to use initially, before credentials are obtained.
-	// (After the OAuth flow is completed, an [oauth2.Transport] with the resulting
-	// [oauth2.TokenSource] is used.)
+// HTTPTransportOptions are options to [NewHTTPTransport].
+type HTTPTransportOptions struct {
+	// Base is the [http.RoundTripper] to use.
 	// If nil, [http.DefaultTransport] is used.
 	Base http.RoundTripper
-	// OAuth is used for HTTP requests that are part of the OAuth protocol,
-	// such as requests to the authorization server. If nil, http.DefaultClient
-	// is used.
-	OAuthClient *http.Client
 }
 
 func (t *HTTPTransport) RoundTrip(req *http.Request) (*http.Response, error) {
-	// baseRoundTrip calls RoundTrip on the base transport.
-	// If we should do OAuth to fix a 401 Unauthorized, it returns nil, nil.
-	baseRoundTrip := func() (*http.Response, error) {
-		t.mu.Lock()
-		base := t.opts.Base
-		_, haveTokenSource := base.(*oauth2.Transport)
-		t.mu.Unlock()
+	t.mu.Lock()
+	base := t.opts.Base
+	_, haveTokenSource := base.(*oauth2.Transport)
+	t.mu.Unlock()
 
-		resp, err := base.RoundTrip(req)
-		if err != nil {
-			return nil, err
-		}
-		if resp.StatusCode != http.StatusUnauthorized {
-			return resp, nil
-		}
-		if haveTokenSource {
-			// We failed to authorize even with a token source; give up.
-			return resp, nil
-		}
-		return nil, nil
+	resp, err := base.RoundTrip(req)
+	if err != nil {
+		return nil, err
 	}
-
-	resp, err := baseRoundTrip()
-	if resp != nil || err != nil {
-		return resp, err
+	if resp.StatusCode != http.StatusUnauthorized {
+		return resp, nil
+	}
+	if haveTokenSource {
+		// We failed to authorize even with a token source; give up.
+		return resp, nil
 	}
-
 	// Try to authorize.
 	t.mu.Lock()
 	// If we don't have a token source, get one by following the OAuth flow.
 	// (We may have obtained one while t.mu was not held above.)
 	if _, ok := t.opts.Base.(*oauth2.Transport); !ok {
-		ts, err := t.doOauth(req.Context(), resp.Header)
+		authHeaders := resp.Header[http.CanonicalHeaderKey("WWW-Authenticate")]
+		ts, err := t.handler(req.Context(), OAuthHandlerArgs{
+			ResourceMetadataURL: extractResourceMetadataURL(authHeaders),
+		})
 		if err != nil {
 			t.mu.Unlock()
 			return nil, err
 		}
 		t.opts.Base = &oauth2.Transport{Base: t.opts.Base, Source: ts}
 	}
 	t.mu.Unlock()
-	// This will not return (nil, nil), because once we have a TokenSource we never lose it.
-	return baseRoundTrip()
+	// Only one level of recursion, because we now have a token source.
+	return t.RoundTrip(req)
 }
 
-// doOauth runs the OAuth 2.1 flow for MCP as described in
-// https://modelcontextprotocol.io/specification/2025-06-18/basic/authorization.
-// It returns the resulting TokenSource.
-func (t *HTTPTransport) doOauth(ctx context.Context, header http.Header) (oauth2.TokenSource, error) {
-	prm, err := oauthex.GetProtectedResourceMetadataFromHeader(ctx, header, c)
-	if err != nil {
-		return nil, err
-	}
-	if len(prm.AuthorizationServers) == 0 {
-		return nil, fmt.Errorf("resource %s provided no authorization servers", prm.Resource)
-	}
-	// TODO: try more than one?
-	authServer := prm.AuthorizationServers[0]
-	// TODO: which scopes to ask for? All of them?
-	scopes := prm.ScopesSupported
-	asm, err := oauthex.GetAuthServerMeta(ctx, authServer, c)
+func extractResourceMetadataURL(authHeaders []string) string {
+	cs, err := oauthex.ParseWWWAuthenticate(authHeaders)
 	if err != nil {
-		return nil, err
-	}
-	// TODO: register the client with the auth server if not registered yet,
-	// or find another way to get the client ID and secret.
-
-	// Get an access token from the auth server.
-	config := &oauth2.Config{
-		ClientID:     "TODO: from registration",
-		ClientSecret: "TODO: from registration",
-		Endpoint: oauth2.Endpoint{
-			AuthURL:  asm.AuthorizationEndpoint,
-			TokenURL: asm.TokenEndpoint,
-			// DeviceAuthURL: "",
-			// AuthStyle: "from auth meta?",
-		},
-		RedirectURL: "", // ???
-		Scopes:      scopes,
-	}
-	v := oauth2.GenerateVerifier()
-	pkceParams := authhandler.PKCEParams{
-		ChallengeMethod: "S256",
-		Challenge:       oauth2.S256ChallengeFromVerifier(v),
-		Verifier:        v,
+		log.Printf("parsing auth headers %q: %v", authHeaders, err)
+		return ""
 	}
-	state := randText()
-	return authhandler.TokenSourceWithPKCE(ctx, config, state, oauthHandler, &pkceParams), nil
+	return oauthex.ResourceMetadataURL(cs)
 }

internal/oauthex/resource_meta.go

@@ -149,11 +149,11 @@ func GetProtectedResourceMetadataFromHeader(ctx context.Context, header http.Hea
 	if len(authHeaders) == 0 {
 		return nil, nil
 	}
-	cs, err := parseWWWAuthenticate(authHeaders)
+	cs, err := ParseWWWAuthenticate(authHeaders)
 	if err != nil {
 		return nil, err
 	}
-	url := resourceMetadataURL(cs)
+	url := ResourceMetadataURL(cs)
 	if url == "" {
 		return nil, nil
 	}
@@ -210,9 +210,9 @@ type challenge struct {
 	Params map[string]string
 }
 
-// resourceMetadataURL returns a resource metadata URL from the given challenges,
+// ResourceMetadataURL returns a resource metadata URL from the given challenges,
 // or the empty string if there is none.
-func resourceMetadataURL(cs []challenge) string {
+func ResourceMetadataURL(cs []challenge) string {
 	for _, c := range cs {
 		if u := c.Params["resource_metadata"]; u != "" {
 			return u
@@ -221,11 +221,11 @@ func resourceMetadataURL(cs []challenge) string {
 	return ""
 }
 
-// parseWWWAuthenticate parses a WWW-Authenticate header string.
+// ParseWWWAuthenticate parses a WWW-Authenticate header string.
 // The header format is defined in RFC 9110, Section 11.6.1, and can contain
 // one or more challenges, separated by commas.
 // It returns a slice of challenges or an error if one of the headers is malformed.
-func parseWWWAuthenticate(headers []string) ([]challenge, error) {
+func ParseWWWAuthenticate(headers []string) ([]challenge, error) {
 	// GENERATED BY GEMINI 2.5 (human-tweaked)
 	var challenges []challenge
 	for _, h := range headers {