From 5e75c8052c49aa1eb028065754948c0febbcafa3 Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Wed, 4 Nov 2026 22:20:45 +0000 Subject: [PATCH] feat: add OAuth token flows, refresh and request authorization Grants are obtained with the authorization code flow with PKCE and a loopback redirect to 127.0.0.1/oauth/callback (RFC 7636, RFC 8252), whose URL is printed for the user, or with the device authorization grant (RFC 8628) when requested or when the server has no authorization endpoint. Every request to the authorization server carries the agent URL as the resource indicator (RFC 8707). An a2a client interceptor attaches stored grants before the first request when the card declares OAuth, and an http.RoundTripper handles 401s: it refreshes a rejected grant, or discovers the server from the challenge, and retries once. Only a card that requires OAuth makes a missing grant an error; otherwise a stored grant is attached when one covers the agent, and problems finding it are only logged. Tokens are sent only over https or to loopback hosts, and requests that already carry an Authorization header are left alone. When a grant has to be obtained interactively, both report an InteractionRequiredError so the caller can sign in outside the request timeout. Grants can also be revoked (RFC 7009). Refs #98 --- internal/oauth/flows.go | 273 ++++++++++++++++++ internal/oauth/interceptor.go | 120 ++++++++ internal/oauth/interceptor_test.go | 185 ++++++++++++ internal/oauth/manager.go | 440 +++++++++++++++++++++++++++++ internal/oauth/manager_test.go | 415 +++++++++++++++++++++++++++ internal/oauth/transport.go | 141 +++++++++ 6 files changed, 1574 insertions(+) create mode 100644 internal/oauth/flows.go create mode 100644 internal/oauth/interceptor.go create mode 100644 internal/oauth/interceptor_test.go create mode 100644 internal/oauth/manager.go create mode 100644 internal/oauth/manager_test.go create mode 100644 internal/oauth/transport.go diff --git a/internal/oauth/flows.go b/internal/oauth/flows.go new file mode 100644 index 0000000..b596dd1 --- /dev/null +++ b/internal/oauth/flows.go @@ -0,0 +1,273 @@ +// Copyright 2026 The A2A Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package oauth + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "crypto/subtle" + "encoding/base64" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/url" + "strings" + "sync" + "time" +) + +// CallbackPath is the loopback redirect path, as listed in the a2a-cli Client ID +// Metadata Document. +const CallbackPath = "/oauth/callback" + +const defaultDeviceInterval = 5 * time.Second + +// Prompter shows the user what to do to authorize. +type Prompter struct { + // Out receives the instructions, normally stderr. + Out io.Writer + // Open, if set, is called with the URL the user has to visit. + Open func(url string) +} + +func (p Prompter) printf(format string, args ...any) { + if p.Out != nil { + _, _ = fmt.Fprintf(p.Out, format, args...) + } +} + +func (p Prompter) open(u string) { + if p.Open != nil { + p.Open(u) + } +} + +func randomString(n int) string { + b := make([]byte, n) + _, _ = rand.Read(b) + return base64.RawURLEncoding.EncodeToString(b) +} + +// pkceChallenge returns the S256 challenge for a verifier (RFC 7636 ยง4.2). +func pkceChallenge(verifier string) string { + sum := sha256.Sum256([]byte(verifier)) + return base64.RawURLEncoding.EncodeToString(sum[:]) +} + +type callbackResult struct { + code string + err error +} + +// authorizationCode runs the authorization code flow with PKCE and a loopback +// redirect (RFC 7636, RFC 8252): it prints the authorization URL, waits for the +// browser to come back to 127.0.0.1, and exchanges the code. +func (f fetcher) authorizationCode(ctx context.Context, e Endpoints, client ClientChoice, resource string, scopes []string, p Prompter) (*tokenResponse, error) { + if e.Authorization == "" { + return nil, fmt.Errorf("authorization server %s has no authorization endpoint; use --device", e.Issuer) + } + authURL, err := url.Parse(e.Authorization) + if err != nil { + return nil, fmt.Errorf("invalid authorization endpoint: %w", err) + } + + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return nil, fmt.Errorf("starting the loopback listener: %w", err) + } + redirectURI := fmt.Sprintf("http://%s%s", lis.Addr().String(), CallbackPath) + state := randomString(32) + verifier := randomString(48) + + results := make(chan callbackResult, 1) + var once sync.Once + deliver := func(r callbackResult) { once.Do(func() { results <- r }) } + mux := http.NewServeMux() + mux.HandleFunc(CallbackPath, func(w http.ResponseWriter, r *http.Request) { + q := r.URL.Query() + if subtle.ConstantTimeCompare([]byte(q.Get("state")), []byte(state)) != 1 { + http.Error(w, "Authorization failed: unexpected state. You can close this tab.", http.StatusBadRequest) + return + } + if code := q.Get("error"); code != "" { + deliver(callbackResult{err: &Error{Code: code, Description: q.Get("error_description")}}) + http.Error(w, "Authorization was not granted. You can close this tab.", http.StatusForbidden) + return + } + code := q.Get("code") + if code == "" { + http.Error(w, "Authorization failed: no code. You can close this tab.", http.StatusBadRequest) + return + } + deliver(callbackResult{code: code}) + w.Header().Set("Content-Type", "text/plain; charset=utf-8") + _, _ = io.WriteString(w, "Authorized. You can close this tab and return to the terminal.\n") + }) + srv := &http.Server{Handler: mux, ReadHeaderTimeout: 10 * time.Second} + go func() { _ = srv.Serve(lis) }() + defer func() { + shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = srv.Shutdown(shutdownCtx) + }() + + q := authURL.Query() + q.Set("response_type", "code") + q.Set("client_id", client.ClientID) + q.Set("redirect_uri", redirectURI) + q.Set("state", state) + q.Set("code_challenge", pkceChallenge(verifier)) + q.Set("code_challenge_method", "S256") + if len(scopes) > 0 { + q.Set("scope", strings.Join(scopes, " ")) + } + if resource != "" { + q.Set("resource", resource) + } + authURL.RawQuery = q.Encode() + + p.printf("To authorize, open:\n%s\n", authURL) + p.open(authURL.String()) + + var res callbackResult + select { + case <-ctx.Done(): + return nil, fmt.Errorf("waiting for authorization: %w", ctx.Err()) + case res = <-results: + } + if res.err != nil { + return nil, res.err + } + + form := tokenForm(grantTypeAuthorizationCode, client.ClientID, resource) + form.Set("code", res.code) + form.Set("redirect_uri", redirectURI) + form.Set("code_verifier", verifier) + return f.token(ctx, e.Token, form) +} + +type deviceAuthorization struct { + DeviceCode string `json:"device_code"` + UserCode string `json:"user_code"` + VerificationURI string `json:"verification_uri"` + VerificationURL string `json:"verification_url"` + VerificationURIComplete string `json:"verification_uri_complete"` + ExpiresIn flexInt `json:"expires_in"` + Interval flexInt `json:"interval"` +} + +// sleepFunc waits for d or until ctx is done. +type sleepFunc func(ctx context.Context, d time.Duration) error + +func sleepContext(ctx context.Context, d time.Duration) error { + t := time.NewTimer(d) + defer t.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-t.C: + return nil + } +} + +// deviceCode runs the device authorization grant (RFC 8628): it prints the +// verification URI and user code, then polls the token endpoint. +func (f fetcher) deviceCode(ctx context.Context, e Endpoints, client ClientChoice, resource string, scopes []string, p Prompter, sleep sleepFunc) (*tokenResponse, error) { + if e.DeviceAuthorization == "" { + return nil, fmt.Errorf("authorization server %s does not support the device flow", e.Issuer) + } + form := url.Values{"client_id": {client.ClientID}} + if len(scopes) > 0 { + form.Set("scope", strings.Join(scopes, " ")) + } + if resource != "" { + form.Set("resource", resource) + } + var da deviceAuthorization + if err := f.postForm(ctx, e.DeviceAuthorization, form, &da); err != nil { + return nil, fmt.Errorf("starting the device flow: %w", err) + } + verification := da.VerificationURI + if verification == "" { + verification = da.VerificationURL + } + if da.DeviceCode == "" || da.UserCode == "" || verification == "" { + return nil, errors.New("device authorization response is missing device_code, user_code or verification_uri") + } + p.printf("To authorize, open: %s\nEnter code: %s\n", verification, da.UserCode) + if da.VerificationURIComplete != "" { + p.open(da.VerificationURIComplete) + } else { + p.open(verification) + } + + interval := defaultDeviceInterval + if da.Interval > 0 { + interval = time.Duration(da.Interval) * time.Second + } + if da.ExpiresIn > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, time.Duration(da.ExpiresIn)*time.Second) + defer cancel() + } + for { + if err := sleep(ctx, interval); err != nil { + return nil, fmt.Errorf("waiting for device authorization: %w", err) + } + poll := tokenForm(grantTypeDeviceCode, client.ClientID, resource) + poll.Set("device_code", da.DeviceCode) + t, err := f.token(ctx, e.Token, poll) + switch { + case err == nil: + return t, nil + case isOAuthError(err, "authorization_pending"): + continue + case isOAuthError(err, "slow_down"): + interval += 5 * time.Second + continue + default: + return nil, err + } + } +} + +// refresh exchanges a refresh token for a new access token. +func (f fetcher) refresh(ctx context.Context, g *Grant) (*tokenResponse, error) { + form := tokenForm(grantTypeRefreshToken, g.Key.ClientID, g.Key.Resource) + form.Set("refresh_token", g.RefreshToken) + if len(g.Key.Scopes) > 0 { + form.Set("scope", strings.Join(g.Key.Scopes, " ")) + } + return f.token(ctx, g.Endpoints.Token, form) +} + +// revoke revokes the grant's refresh token, or its access token when it has none +// (RFC 7009). +func (f fetcher) revoke(ctx context.Context, g *Grant) error { + if g.Endpoints.Revocation == "" { + return errRevocationUnsupported + } + token, hint := g.RefreshToken, "refresh_token" + if token == "" { + token, hint = g.AccessToken, "access_token" + } + form := url.Values{"token": {token}, "token_type_hint": {hint}, "client_id": {g.Key.ClientID}} + return f.postForm(ctx, g.Endpoints.Revocation, form, nil) +} + +var errRevocationUnsupported = errors.New("the authorization server has no revocation endpoint") diff --git a/internal/oauth/interceptor.go b/internal/oauth/interceptor.go new file mode 100644 index 0000000..fbf95b8 --- /dev/null +++ b/internal/oauth/interceptor.go @@ -0,0 +1,120 @@ +// Copyright 2026 The A2A Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package oauth + +import ( + "context" + "fmt" + "strings" + + "github.com/a2aproject/a2a-go/v2/a2aclient" +) + +type attachedGrantKey struct{} + +// withAttached records on ctx the grant whose token a request carries, so the +// transport can tell its own token from an explicit credential. +func withAttached(ctx context.Context, g *Grant) context.Context { + return context.WithValue(ctx, attachedGrantKey{}, g) +} + +func attachedFrom(ctx context.Context) (*Grant, bool) { + g, ok := ctx.Value(attachedGrantKey{}).(*Grant) + return g, ok +} + +// Interceptor returns an a2a client interceptor that attaches grants to calls. +// When the agent card requires OAuth for the call's interface, the call fails +// with *InteractionRequiredError if no grant is stored. Otherwise a stored grant +// is attached when there is one, and problems finding it are only logged, so +// agents that accept other credentials or none are not affected. Calls that +// already carry an Authorization service parameter (from --auth, --svc-param or +// a credential helper) are left alone. +func (m *Manager) Interceptor() a2aclient.CallInterceptor { + return &interceptor{m: m} +} + +type interceptor struct { + a2aclient.PassthroughInterceptor + m *Manager +} + +var _ a2aclient.CallInterceptor = (*interceptor)(nil) + +// Before implements [a2aclient.CallInterceptor]. +func (i *interceptor) Before(ctx context.Context, req *a2aclient.Request) (context.Context, any, error) { + if len(req.ServiceParams.Get("authorization")) > 0 { + return ctx, nil, nil + } + var r *Requirement + if req.Card != nil { + var err error + r, err = i.m.RequirementForCard(ctx, req.Card, req.BaseURL) + if err != nil { + if !acceptsOtherCredentials(req.Card) { + return ctx, nil, err + } + i.m.logf("not using OAuth for %s: %v", req.BaseURL, err) + } + } + if r == nil || r.Optional { + g, err := i.optionalGrant(ctx, req.BaseURL, r) + if err != nil || g == nil { + if err != nil { + i.m.logf("not using a stored OAuth grant for %s: %v", req.BaseURL, err) + } + return ctx, nil, nil + } + req.ServiceParams.Append("Authorization", g.header()) + return withAttached(ctx, g), nil, nil + } + if err := checkTokenDestination(req.BaseURL); err != nil { + return ctx, nil, failure("%w", err) + } + g, err := i.m.Grant(ctx, r) + if err != nil { + return ctx, nil, err + } + if g == nil { + return ctx, nil, &InteractionRequiredError{Requirement: r} + } + req.ServiceParams.Append("Authorization", g.header()) + return withAttached(ctx, g), nil, nil +} + +// optionalGrant returns a stored grant for a call that doesn't require one: for +// the optional requirement r, or else for the call's resource. +func (i *interceptor) optionalGrant(ctx context.Context, baseURL string, r *Requirement) (*Grant, error) { + if err := checkTokenDestination(baseURL); err != nil { + return nil, err + } + if r != nil { + return i.m.Grant(ctx, r) + } + return i.m.GrantForResource(ctx, baseURL) +} + +// checkTokenDestination refuses to send a bearer token where it would travel in +// cleartext: it allows https, http to a loopback host, and gRPC addresses +// (host:port), which use TLS unless --insecure is set. +func checkTokenDestination(target string) error { + if !strings.Contains(target, "://") { + return nil + } + if err := checkSecureURL(target); err != nil { + return fmt.Errorf("refusing to send an OAuth token to %s over plain HTTP", target) + } + return nil +} diff --git a/internal/oauth/interceptor_test.go b/internal/oauth/interceptor_test.go new file mode 100644 index 0000000..b646d65 --- /dev/null +++ b/internal/oauth/interceptor_test.go @@ -0,0 +1,185 @@ +// Copyright 2026 The A2A Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package oauth_test + +import ( + "context" + "errors" + "os" + "path/filepath" + "testing" + "time" + + "github.com/a2aproject/a2a-cli/internal/oauth" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" +) + +var inlineEndpoints = oauth.Endpoints{ + Issuer: "https://as.example.com", + Authorization: "https://as.example.com/authorize", + Token: "https://as.example.com/token", +} + +// inlineCard declares the inline oauth2 scheme, and an API key alternative when +// withAPIKey is set. +func inlineCard(withAPIKey bool) *a2a.AgentCard { + card := &a2a.AgentCard{ + Name: "inline", + SecuritySchemes: a2a.NamedSecuritySchemes{ + "oauth": a2a.OAuth2SecurityScheme{Flows: a2a.AuthorizationCodeOAuthFlow{AuthorizationURL: inlineEndpoints.Authorization, TokenURL: inlineEndpoints.Token}}, + "key": a2a.APIKeySecurityScheme{Name: "X-API-Key", Location: "header"}, + }, + SecurityRequirements: a2a.SecurityRequirementsOptions{{"oauth": nil}}, + } + if withAPIKey { + card.SecurityRequirements = append(card.SecurityRequirements, a2a.SecurityRequirements{"key": nil}) + } + return card +} + +// managerWithGrant returns a manager whose store holds one grant for resource +// and scopes, with the given tokens. +func managerWithGrant(t *testing.T, resource string, scopes []string, access, refresh string) *oauth.Manager { + t.Helper() + store := &oauth.FileStore{Dir: t.TempDir()} + if resource != "" { + g := oauth.Grant{ + Key: oauth.NewGrantKey(inlineEndpoints.Issuer, "c", resource, scopes), + AccessToken: access, + RefreshToken: refresh, + Endpoints: inlineEndpoints, + } + if access != "" { + g.ExpiresAt = time.Now().Add(time.Hour) + } + if err := store.Put(g); err != nil { + t.Fatal(err) + } + } + return &oauth.Manager{Clients: []oauth.Registration{{Authority: inlineEndpoints.Issuer, ClientID: "c"}}, Store: store} +} + +// before runs the interceptor and returns the Authorization values it attached. +func before(t *testing.T, m *oauth.Manager, baseURL string, card *a2a.AgentCard) ([]string, error) { + t.Helper() + req := &a2aclient.Request{BaseURL: baseURL, Card: card, ServiceParams: a2aclient.ServiceParams{}} + _, _, err := m.Interceptor().Before(context.Background(), req) + return req.ServiceParams.Get("authorization"), err +} + +func TestInterceptorOptionalOAuth(t *testing.T) { + t.Parallel() + const agent = "https://agent.example.com" + + // A card that also accepts an API key goes out without a grant. + got, err := before(t, managerWithGrant(t, "", nil, "", ""), agent, inlineCard(true)) + if err != nil || len(got) != 0 { + t.Fatalf("Before() without a grant = %q, %v, want no token and no error", got, err) + } + // It still gets a stored grant. + got, err = before(t, managerWithGrant(t, agent, nil, "at", ""), agent, inlineCard(true)) + if err != nil || len(got) != 1 || got[0] != "Bearer at" { + t.Fatalf("Before() with a grant = %q, %v, want the stored token", got, err) + } + // A card that only accepts OAuth requires sign-in. + _, err = before(t, managerWithGrant(t, "", nil, "", ""), agent, inlineCard(false)) + var ir *oauth.InteractionRequiredError + if !errors.As(err, &ir) { + t.Fatalf("Before() for an OAuth-only card error = %v, want InteractionRequiredError", err) + } + // Without a matching client, only the OAuth-only card fails. + m := managerWithGrant(t, "", nil, "", "") + m.Clients = nil + if got, err := before(t, m, agent, inlineCard(true)); err != nil || len(got) != 0 { + t.Fatalf("Before() for an optional card without a client = %q, %v, want no token and no error", got, err) + } + var nc *oauth.NoClientError + if _, err := before(t, m, agent, inlineCard(false)); !errors.As(err, &nc) { + t.Fatalf("Before() for an OAuth-only card without a client error = %v, want NoClientError", err) + } +} + +func TestInterceptorUnreadableStore(t *testing.T) { + t.Parallel() + m := managerWithGrant(t, "", nil, "", "") + dir := m.Store.(*oauth.FileStore).Dir + if err := os.WriteFile(filepath.Join(dir, "grants.json"), []byte("{not json"), 0o600); err != nil { + t.Fatal(err) + } + // Agents that don't require OAuth are not affected. + if got, err := before(t, m, "https://agent.example.com", nil); err != nil || len(got) != 0 { + t.Fatalf("Before() without a requirement = %q, %v, want no token and no error", got, err) + } + if _, err := before(t, m, "https://agent.example.com", inlineCard(false)); err == nil { + t.Fatal("Before() for an OAuth-only card error = nil, want the store error") + } +} + +func TestInterceptorRefusesCleartext(t *testing.T) { + t.Parallel() + const agent = "http://agent.example.com" + m := managerWithGrant(t, agent, nil, "secret", "") + got, err := before(t, m, agent, inlineCard(false)) + var fe *oauth.FailureError + if !errors.As(err, &fe) || len(got) != 0 { + t.Fatalf("Before() for an http agent = %q, %v, want FailureError and no token", got, err) + } + if got, err := before(t, m, agent, nil); err != nil || len(got) != 0 { + t.Fatalf("Before() for an http agent without a requirement = %q, %v, want no token and no error", got, err) + } + if got, err := before(t, managerWithGrant(t, "http://127.0.0.1:8080", nil, "at", ""), "http://127.0.0.1:8080", nil); err != nil || len(got) != 1 { + t.Fatalf("Before() for a loopback agent = %q, %v, want the token", got, err) + } +} + +func TestInterceptorGrantWithoutTokens(t *testing.T) { + t.Parallel() + const agent = "https://agent.example.com" + // A grant whose tokens are gone, such as a keychain entry that was removed, + // needs a new sign-in instead of an empty Bearer token. + got, err := before(t, managerWithGrant(t, agent, nil, "", ""), agent, inlineCard(false)) + var ir *oauth.InteractionRequiredError + if !errors.As(err, &ir) || len(got) != 0 { + t.Fatalf("Before() = %q, %v, want InteractionRequiredError and no token", got, err) + } + if got, err := before(t, managerWithGrant(t, agent, nil, "", ""), agent, nil); err != nil || len(got) != 0 { + t.Fatalf("Before() without a requirement = %q, %v, want no token and no error", got, err) + } +} + +func TestInterceptorCoveringGrant(t *testing.T) { + t.Parallel() + const iface = "https://agent.example.com/a2a/v1" + // A grant for the agent's origin, with the scopes auth login --scope asked + // for, serves calls to an interface below it. + m := managerWithGrant(t, "https://agent.example.com", []string{"a2a.invoke"}, "at", "") + for _, card := range []*a2a.AgentCard{inlineCard(false), nil} { + if got, err := before(t, m, iface, card); err != nil || len(got) != 1 || got[0] != "Bearer at" { + t.Fatalf("Before(card=%v) = %q, %v, want the origin grant", card != nil, got, err) + } + } + // A grant lacking a required scope is not used. + card := inlineCard(false) + card.SecurityRequirements = a2a.SecurityRequirementsOptions{{"oauth": {"a2a.admin"}}} + var ir *oauth.InteractionRequiredError + if _, err := before(t, m, iface, card); !errors.As(err, &ir) { + t.Fatalf("Before() needing another scope error = %v, want InteractionRequiredError", err) + } + // Another agent on the same host is not covered. + if got, err := before(t, managerWithGrant(t, "https://agent.example.com/a", nil, "at", ""), "https://agent.example.com/b", nil); err != nil || len(got) != 0 { + t.Fatalf("Before() for a sibling path = %q, %v, want no token", got, err) + } +} diff --git a/internal/oauth/manager.go b/internal/oauth/manager.go new file mode 100644 index 0000000..29f6b69 --- /dev/null +++ b/internal/oauth/manager.go @@ -0,0 +1,440 @@ +// Copyright 2026 The A2A Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package oauth + +import ( + "context" + "errors" + "net/http" + "net/url" + "slices" + "sync" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" +) + +const lockTimeout = 10 * time.Second + +// ErrRevocationUnsupported is returned by [Manager.Revoke] when the +// authorization server has no revocation endpoint. +var ErrRevocationUnsupported = errRevocationUnsupported + +// Manager finds, refreshes and obtains grants for agent requests. +type Manager struct { + // Clients are the configured client registrations. + Clients []Registration + // DefaultClientID is used as the client ID with authorization servers that + // support Client ID Metadata Documents when no registration matches, + // normally [DefaultClientID]. Empty disables this fallback. + DefaultClientID string + // Store keeps grants. + Store Store + // HTTPClient makes metadata, token and revocation requests. It must not add + // agent credentials. + HTTPClient *http.Client + // Logf, if set, receives diagnostics. Tokens are never logged. + Logf func(format string, args ...any) + // Now, if set, replaces time.Now. + Now func() time.Time + + mu sync.Mutex + cardReqs map[cardReqKey]*Requirement + learned map[string]GrantKey +} + +type cardReqKey struct { + card *a2a.AgentCard + baseURL string +} + +// AcquireOptions control an interactive authorization. +type AcquireOptions struct { + // Device forces the device authorization grant. + Device bool + // Scopes replace the scopes from the requirement and the registration. + Scopes []string + // Prompter shows the user what to do. + Prompter Prompter + // Sleep, if set, replaces the wait between device flow polls. + Sleep func(ctx context.Context, d time.Duration) error +} + +func (m *Manager) fetcher() fetcher { + client := m.HTTPClient + if client == nil { + client = http.DefaultClient + } + return fetcher{client: client} +} + +func (m *Manager) now() time.Time { + if m.Now != nil { + return m.Now() + } + return time.Now() +} + +func (m *Manager) logf(format string, args ...any) { + if m.Logf != nil { + m.Logf(format, args...) + } +} + +// RequirementForCard returns the OAuth requirement the card declares for calls to +// baseURL, or nil when it declares none. Results are cached per card and URL. +func (m *Manager) RequirementForCard(ctx context.Context, card *a2a.AgentCard, baseURL string) (*Requirement, error) { + key := cardReqKey{card: card, baseURL: baseURL} + m.mu.Lock() + r, ok := m.cardReqs[key] + m.mu.Unlock() + if ok { + return r, nil + } + r, err := m.fetcher().requirementFromCard(ctx, card, baseURL) + if err != nil { + return nil, failure("reading the authorization requirements of %s: %w", card.Name, err) + } + if r != nil { + m.logf("agent card requires authorization at %s for %s", r.Endpoints.Issuer, r.Resource) + } + m.mu.Lock() + if m.cardReqs == nil { + m.cardReqs = map[cardReqKey]*Requirement{} + } + m.cardReqs[key] = r + m.mu.Unlock() + return r, nil +} + +// RequirementForChallenge follows the Bearer challenges of a 401 response for +// target to its Protected Resource Metadata and the first listed authorization +// server whose metadata validates and that has a usable client. It returns +// (nil, nil) when the resource publishes no metadata. +func (m *Manager) RequirementForChallenge(ctx context.Context, target *url.URL, chs []Challenge) (*Requirement, error) { + f := m.fetcher() + ct, err := f.discoverFromChallenge(ctx, target, chs) + if errors.Is(err, errNoResourceMetadata) { + return nil, nil + } + if err != nil { + return nil, failure("discovering authorization for %s: %w", target.Redacted(), err) + } + m.logf("resource %s uses authorization servers %v", ct.resource.Resource, ct.resource.AuthorizationServers) + var errs []error + var noClient []string + for _, issuer := range ct.resource.AuthorizationServers { + md, err := f.serverMetadata(ctx, issuer) + if err != nil { + errs = append(errs, err) + continue + } + e := md.endpoints() + if err := e.validate(); err != nil { + errs = append(errs, err) + continue + } + if _, err := selectClient(m.Clients, m.DefaultClientID, e); err != nil { + var nc *NoClientError + if errors.As(err, &nc) { + noClient = append(noClient, issuer) + continue + } + return nil, err + } + return &Requirement{Endpoints: e, Resource: NormalizeResource(ct.resource.Resource), Scopes: ct.scopes}, nil + } + if len(noClient) > 0 { + return nil, &NoClientError{Issuers: noClient} + } + return nil, failure("no usable authorization server for %s: %w", ct.resource.Resource, errors.Join(errs...)) +} + +// Key returns the grant key and client for a requirement. Scopes come from +// override, else the requirement, else the registration. +func (m *Manager) Key(req *Requirement, override []string) (GrantKey, ClientChoice, error) { + choice, err := selectClient(m.Clients, m.DefaultClientID, req.Endpoints) + if err != nil { + return GrantKey{}, ClientChoice{}, err + } + scopes := override + if len(scopes) == 0 { + scopes = req.Scopes + } + if len(scopes) == 0 { + scopes = choice.Scopes + } + return NewGrantKey(req.Endpoints.Issuer, choice.ClientID, req.Resource, scopes), choice, nil +} + +// Grant returns the stored grant for a requirement, refreshed when it is about to +// expire, or nil when there is none that can be used without interaction. +func (m *Manager) Grant(ctx context.Context, req *Requirement) (*Grant, error) { + key, _, err := m.Key(req, nil) + if err != nil { + return nil, err + } + found, err := m.lookup(key) + if err != nil || found == nil { + return nil, err + } + return m.grant(ctx, *found) +} + +// lookup returns the key of the stored grant to use for key: key itself, or else +// the grant from the same issuer and client with the most specific resource +// covering key's resource and with all of key's scopes, such as one obtained +// with auth login --scope or for the agent's origin. It returns nil when there +// is none. +func (m *Manager) lookup(key GrantKey) (*GrantKey, error) { + g, err := m.Store.Get(key) + if err != nil { + return nil, err + } + if g != nil { + return &key, nil + } + grants, err := m.Store.List() + if err != nil { + return nil, err + } + var found *GrantKey + for _, g := range grants { + k := g.Key + if k.Issuer != key.Issuer || k.ClientID != key.ClientID || !resourceCovers(k.Resource, key.Resource) || !containsAll(k.Scopes, key.Scopes) { + continue + } + if found == nil || len(k.Resource) > len(found.Resource) { + found = &k + } + } + return found, nil +} + +func containsAll(have, want []string) bool { + for _, w := range want { + if !slices.Contains(have, w) { + return false + } + } + return true +} + +func (m *Manager) grant(ctx context.Context, key GrantKey) (*Grant, error) { + g, err := m.Store.Get(key) + if err != nil || g == nil { + return nil, err + } + if !g.needsRefresh(m.now()) { + return g, nil + } + return m.refreshIfStale(ctx, key) +} + +// refreshIfStale refreshes the grant for key when its access token is missing or +// about to expire. It returns nil when a new authorization is needed. +func (m *Manager) refreshIfStale(ctx context.Context, key GrantKey) (*Grant, error) { + return m.refresh(ctx, key, false, "") +} + +// refreshRejected refreshes the grant for key after the agent refused its +// access token, unless another process has already replaced that token. It +// returns nil when a new authorization is needed. +func (m *Manager) refreshRejected(ctx context.Context, key GrantKey, token string) (*Grant, error) { + return m.refresh(ctx, key, true, token) +} + +// refresh refreshes the grant for key under the store lock; see refreshIfStale +// and refreshRejected. +func (m *Manager) refresh(ctx context.Context, key GrantKey, rejected bool, token string) (*Grant, error) { + unlock, err := m.lock(ctx) + if err != nil { + return nil, err + } + defer unlock() + + g, err := m.Store.Get(key) + if err != nil || g == nil { + return nil, err + } + now := m.now() + if g.AccessToken == "" && g.RefreshToken == "" { + return nil, nil + } + if !rejected && !g.needsRefresh(now) { + return g, nil + } + if rejected && g.AccessToken != token && !g.needsRefresh(now) { + return g, nil + } + usable := g.AccessToken != "" && !g.expired(now) + if g.RefreshToken == "" { + if !rejected && usable { + return g, nil + } + return nil, nil + } + m.logf("refreshing the access token for %s", g.Key.Resource) + t, err := m.fetcher().refresh(ctx, g) + if isOAuthError(err, "invalid_grant") { + m.logf("the refresh token for %s was rejected; a new authorization is needed", g.Key.Resource) + if err := m.Store.Delete(key); err != nil { + return nil, err + } + return nil, nil + } + if err != nil { + if !rejected && usable { + m.logf("refreshing the access token for %s failed, using the current one: %v", g.Key.Resource, err) + return g, nil + } + return nil, failure("refreshing the access token for %s: %w", g.Key.Resource, err) + } + g.apply(t, now) + if err := m.Store.Put(*g); err != nil { + return nil, err + } + return g, nil +} + +// Acquire obtains a new grant interactively and stores it: with the device flow +// when requested or when the server or card offers no other, otherwise with the +// authorization code flow. +func (m *Manager) Acquire(ctx context.Context, req *Requirement, opts AcquireOptions) (*Grant, error) { + key, choice, err := m.Key(req, opts.Scopes) + if err != nil { + return nil, err + } + m.logf("authorizing %s at %s as client %s (%s)", req.Resource, req.Endpoints.Issuer, choice.ClientID, choice.Source) + f := m.fetcher() + var t *tokenResponse + if opts.Device || req.PreferDevice || req.Endpoints.Authorization == "" { + sleep := sleepFunc(sleepContext) + if opts.Sleep != nil { + sleep = opts.Sleep + } + t, err = f.deviceCode(ctx, req.Endpoints, choice, key.Resource, key.Scopes, opts.Prompter, sleep) + } else { + t, err = f.authorizationCode(ctx, req.Endpoints, choice, key.Resource, key.Scopes, opts.Prompter) + } + if err != nil { + return nil, err + } + now := m.now() + g := Grant{Key: key, Endpoints: req.Endpoints, CreatedAt: now.UTC(), ClientSource: choice.Source} + g.apply(t, now) + unlock, err := m.lock(ctx) + if err != nil { + return nil, err + } + defer unlock() + if err := m.Store.Put(g); err != nil { + return nil, err + } + m.learn(g.Key) + return &g, nil +} + +// GrantForResource returns the stored grant for the most specific resource +// covering resource, refreshed when needed, or nil when there is none or +// several grants share that resource. +func (m *Manager) GrantForResource(ctx context.Context, resource string) (*Grant, error) { + resource = NormalizeResource(resource) + m.mu.Lock() + key, ok := m.learned[resource] + m.mu.Unlock() + if ok { + return m.grant(ctx, key) + } + grants, err := m.Store.List() + if err != nil { + return nil, err + } + var found *GrantKey + ambiguous := false + for _, g := range grants { + k := g.Key + if !resourceCovers(k.Resource, resource) { + continue + } + switch { + case found == nil || len(k.Resource) > len(found.Resource): + found, ambiguous = &k, false + case len(k.Resource) == len(found.Resource): + ambiguous = true + } + } + if ambiguous { + m.logf("several grants for %s; not choosing one", found.Resource) + return nil, nil + } + if found == nil { + return nil, nil + } + return m.grant(ctx, *found) +} + +// learnedGrant returns a grant learned in this process for a resource covering target. +func (m *Manager) learnedGrant(ctx context.Context, target *url.URL) (*Grant, error) { + m.mu.Lock() + var key *GrantKey + for resource, k := range m.learned { + if resourceCovers(resource, target.String()) { + k := k + key = &k + break + } + } + m.mu.Unlock() + if key == nil { + return nil, nil + } + return m.grant(ctx, *key) +} + +func (m *Manager) learn(key GrantKey) { + m.mu.Lock() + defer m.mu.Unlock() + if m.learned == nil { + m.learned = map[string]GrantKey{} + } + m.learned[key.Resource] = key +} + +// List returns every stored grant. +func (m *Manager) List() ([]Grant, error) { + return m.Store.List() +} + +// Revoke revokes a grant at its authorization server. +func (m *Manager) Revoke(ctx context.Context, g *Grant) error { + return m.fetcher().revoke(ctx, g) +} + +// Delete removes a grant from the store. +func (m *Manager) Delete(ctx context.Context, key GrantKey) error { + unlock, err := m.lock(ctx) + if err != nil { + return err + } + defer unlock() + return m.Store.Delete(key) +} + +func (m *Manager) lock(ctx context.Context) (func(), error) { + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), lockTimeout) + defer cancel() + return m.Store.Lock(ctx) +} diff --git a/internal/oauth/manager_test.go b/internal/oauth/manager_test.go new file mode 100644 index 0000000..32d88b1 --- /dev/null +++ b/internal/oauth/manager_test.go @@ -0,0 +1,415 @@ +// Copyright 2026 The A2A Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package oauth_test + +import ( + "bytes" + "context" + "errors" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "github.com/google/go-cmp/cmp" + + "github.com/a2aproject/a2a-cli/internal/oauth" + "github.com/a2aproject/a2a-cli/internal/oauth/oauthtest" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient/agentcard" +) + +func newManager(t *testing.T, as *oauthtest.AuthServer) *oauth.Manager { + t.Helper() + return &oauth.Manager{ + Clients: []oauth.Registration{{Authority: as.URL, ClientID: "test-client"}}, + DefaultClientID: oauth.DefaultClientID, + Store: &oauth.FileStore{Dir: t.TempDir()}, + HTTPClient: http.DefaultClient, + } +} + +func fetchCard(t *testing.T, agentURL string) *a2a.AgentCard { + t.Helper() + card, err := agentcard.DefaultResolver.Resolve(t.Context(), agentURL) + if err != nil { + t.Fatalf("Resolve(%q) error = %v", agentURL, err) + } + return card +} + +func noSleep(context.Context, time.Duration) error { return nil } + +func TestRequirementForCard(t *testing.T) { + t.Parallel() + as := oauthtest.NewAuthServer(t, oauthtest.ServerOptions{}) + full := oauth.Endpoints{ + Issuer: as.URL, + Authorization: as.URL + "/authorize", + Token: as.URL + "/token", + DeviceAuthorization: as.URL + "/device", + Revocation: as.URL + "/revoke", + } + tests := []struct { + name string + card oauthtest.CardSecurity + want *oauth.Requirement + }{ + {name: "no security", card: oauthtest.CardNone}, + {name: "metadata url", card: oauthtest.CardMetadataURL, want: &oauth.Requirement{Endpoints: full, Scopes: []string{"a2a.invoke"}}}, + {name: "openid connect", card: oauthtest.CardOIDC, want: &oauth.Requirement{Endpoints: full, Scopes: []string{"a2a.invoke"}}}, + { + name: "inline authorization code", + card: oauthtest.CardAuthCode, + want: &oauth.Requirement{Endpoints: oauth.Endpoints{Issuer: as.URL, Authorization: as.URL + "/authorize", Token: as.URL + "/token"}, Scopes: []string{"a2a.invoke"}}, + }, + { + name: "inline device code", + card: oauthtest.CardDeviceCode, + want: &oauth.Requirement{Endpoints: oauth.Endpoints{Issuer: as.URL, DeviceAuthorization: as.URL + "/device", Token: as.URL + "/token"}, Scopes: []string{"a2a.invoke"}, PreferDevice: true}, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + agent := oauthtest.NewProtectedAgent(t, as, oauthtest.AgentOptions{Card: tc.card, Scopes: []string{"a2a.invoke"}}) + m := newManager(t, as) + got, err := m.RequirementForCard(t.Context(), fetchCard(t, agent.URL), agent.URL+"/") + if err != nil { + t.Fatalf("Manager.RequirementForCard() error = %v", err) + } + if tc.want != nil { + tc.want.Resource = agent.URL + } + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Fatalf("Manager.RequirementForCard() wrong result (-want +got) diff = %s", diff) + } + }) + } +} + +func TestRequirementForChallenge(t *testing.T) { + t.Parallel() + tests := []struct { + name string + agent oauthtest.AgentOptions + clients func(as *oauthtest.AuthServer) []oauth.Registration + server oauthtest.ServerOptions + noDefault bool + wantErr bool + noClient bool + }{ + {name: "resource_metadata in the challenge", agent: oauthtest.AgentOptions{Scopes: []string{"a2a.invoke"}}}, + {name: "well-known probe", agent: oauthtest.AgentOptions{OmitResourceMetadata: true}}, + {name: "OIDC discovery only", server: oauthtest.ServerOptions{OIDCOnly: true}}, + { + name: "no client", + clients: func(*oauthtest.AuthServer) []oauth.Registration { return nil }, + wantErr: true, + noClient: true, + }, + { + name: "CIMD fallback", + clients: func(*oauthtest.AuthServer) []oauth.Registration { return nil }, + server: oauthtest.ServerOptions{CIMD: true}, + }, + { + name: "CIMD fallback disabled", + clients: func(*oauthtest.AuthServer) []oauth.Registration { return nil }, + server: oauthtest.ServerOptions{CIMD: true}, + noDefault: true, + wantErr: true, + noClient: true, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + as := oauthtest.NewAuthServer(t, tc.server) + agent := oauthtest.NewProtectedAgent(t, as, tc.agent) + m := newManager(t, as) + if tc.clients != nil { + m.Clients = tc.clients(as) + } + if tc.noDefault { + m.DefaultClientID = "" + } + resp, err := http.Post(agent.URL+"/message:send", "application/json", strings.NewReader("{}")) + if err != nil { + t.Fatal(err) + } + _ = resp.Body.Close() + target, _ := url.Parse(agent.URL + "/message:send") + got, err := m.RequirementForChallenge(t.Context(), target, oauth.ParseChallenges(resp.Header.Values("WWW-Authenticate"))) + if (err != nil) != tc.wantErr { + t.Fatalf("Manager.RequirementForChallenge() error = %v, want error %v", err, tc.wantErr) + } + var nc *oauth.NoClientError + if errors.As(err, &nc) != tc.noClient { + t.Fatalf("Manager.RequirementForChallenge() error = %v, want NoClientError %v", err, tc.noClient) + } + if tc.wantErr { + return + } + if got.Resource != agent.URL || got.Endpoints.Issuer != as.URL || got.Endpoints.Token != as.URL+"/token" { + t.Fatalf("Manager.RequirementForChallenge() = %+v, want resource %s at %s", got, agent.URL, as.URL) + } + if diff := cmp.Diff(tc.agent.Scopes, got.Scopes); diff != "" { + t.Fatalf("Manager.RequirementForChallenge() scopes wrong result (-want +got) diff = %s", diff) + } + }) + } +} + +func TestRequirementForChallengeRejectsBadMetadata(t *testing.T) { + t.Parallel() + as := oauthtest.NewAuthServer(t, oauthtest.ServerOptions{}) + other := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"resource":"http://elsewhere.example","authorization_servers":["` + as.URL + `"]}`)) + })) + t.Cleanup(other.Close) + tests := []struct { + name string + handler http.HandlerFunc + }{ + { + name: "cross-origin redirect", + handler: func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, other.URL+"/prm", http.StatusFound) + }, + }, + { + name: "metadata for another resource", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"resource":"http://elsewhere.example","authorization_servers":["` + as.URL + `"]}`)) + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(tc.handler) + t.Cleanup(srv.Close) + target, _ := url.Parse(srv.URL + "/a2a") + chs := oauth.ParseChallenges([]string{`Bearer resource_metadata="` + srv.URL + `/prm"`}) + _, err := newManager(t, as).RequirementForChallenge(t.Context(), target, chs) + var fe *oauth.FailureError + if !errors.As(err, &fe) { + t.Fatalf("Manager.RequirementForChallenge() error = %v, want FailureError", err) + } + }) + } +} + +func TestAcquire(t *testing.T) { + t.Parallel() + tests := []struct { + name string + server oauthtest.ServerOptions + device bool + path string + }{ + {name: "authorization code", path: "/token"}, + {name: "device flow requested", device: true, path: "/device"}, + {name: "device flow without authorization endpoint", server: oauthtest.ServerOptions{NoAuthorizationEndpoint: true, PendingPolls: 2}, path: "/device"}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + as := oauthtest.NewAuthServer(t, tc.server) + agent := oauthtest.NewProtectedAgent(t, as, oauthtest.AgentOptions{Card: oauthtest.CardMetadataURL, Scopes: []string{"a2a.invoke"}}) + m := newManager(t, as) + req, err := m.RequirementForCard(t.Context(), fetchCard(t, agent.URL), agent.URL) + if err != nil { + t.Fatal(err) + } + var prompt bytes.Buffer + g, err := m.Acquire(t.Context(), req, oauth.AcquireOptions{ + Device: tc.device, + Prompter: oauth.Prompter{Out: &prompt, Open: oauthtest.Browser(t)}, + Sleep: noSleep, + }) + if err != nil { + t.Fatalf("Manager.Acquire() error = %v", err) + } + if !as.Valid(g.AccessToken, agent.URL, []string{"a2a.invoke"}) { + t.Fatalf("Manager.Acquire() returned a token the agent does not accept") + } + if !strings.Contains(prompt.String(), "To authorize, open:") { + t.Fatalf("Manager.Acquire() prompt = %q, want authorization instructions", prompt.String()) + } + if got := as.LastForm(tc.path).Get("resource"); got != agent.URL { + t.Fatalf("resource sent to %s = %q, want %q", tc.path, got, agent.URL) + } + stored, err := m.Grant(t.Context(), req) + if err != nil || stored == nil || stored.AccessToken != g.AccessToken { + t.Fatalf("Manager.Grant() after Acquire = %v, %v, want the acquired grant", stored, err) + } + }) + } +} + +func TestAcquireDenied(t *testing.T) { + t.Parallel() + for _, device := range []bool{false, true} { + as := oauthtest.NewAuthServer(t, oauthtest.ServerOptions{DenyAuthorization: true}) + agent := oauthtest.NewProtectedAgent(t, as, oauthtest.AgentOptions{Card: oauthtest.CardMetadataURL}) + m := newManager(t, as) + req, err := m.RequirementForCard(t.Context(), fetchCard(t, agent.URL), agent.URL) + if err != nil { + t.Fatal(err) + } + _, err = m.Acquire(t.Context(), req, oauth.AcquireOptions{Device: device, Prompter: oauth.Prompter{Open: oauthtest.Browser(t)}, Sleep: noSleep}) + var oe *oauth.Error + if !errors.As(err, &oe) || oe.Code != "access_denied" { + t.Fatalf("Manager.Acquire(device=%v) error = %v, want access_denied", device, err) + } + } +} + +func acquire(t *testing.T, m *oauth.Manager, agentURL string) (*oauth.Requirement, *oauth.Grant) { + t.Helper() + req, err := m.RequirementForCard(t.Context(), fetchCard(t, agentURL), agentURL) + if err != nil { + t.Fatal(err) + } + g, err := m.Acquire(t.Context(), req, oauth.AcquireOptions{Prompter: oauth.Prompter{Open: oauthtest.Browser(t)}}) + if err != nil { + t.Fatalf("Manager.Acquire() error = %v", err) + } + return req, g +} + +func TestGrantRefreshesNearExpiry(t *testing.T) { + t.Parallel() + for _, rotate := range []bool{false, true} { + as := oauthtest.NewAuthServer(t, oauthtest.ServerOptions{RotateRefreshTokens: rotate, TokenLifetime: 30 * time.Second}) + agent := oauthtest.NewProtectedAgent(t, as, oauthtest.AgentOptions{Card: oauthtest.CardMetadataURL}) + m := newManager(t, as) + req, first := acquire(t, m, agent.URL) + for i := range 2 { + g, err := m.Grant(t.Context(), req) + if err != nil { + t.Fatalf("Manager.Grant() #%d error = %v", i, err) + } + if g.AccessToken == first.AccessToken { + t.Fatalf("Manager.Grant() #%d did not refresh a token that expires in 30s", i) + } + if !as.Valid(g.AccessToken, agent.URL, nil) { + t.Fatalf("Manager.Grant() #%d returned an invalid token (rotate=%v)", i, rotate) + } + } + if rotate && as.Hits("/token") != 3 { + t.Fatalf("token endpoint hits = %d, want 3 (one exchange, two refreshes with rotated tokens)", as.Hits("/token")) + } + } +} + +func TestGrantDropsRevokedRefreshToken(t *testing.T) { + t.Parallel() + as := oauthtest.NewAuthServer(t, oauthtest.ServerOptions{TokenLifetime: 30 * time.Second}) + agent := oauthtest.NewProtectedAgent(t, as, oauthtest.AgentOptions{Card: oauthtest.CardMetadataURL}) + m := newManager(t, as) + req, _ := acquire(t, m, agent.URL) + as.RevokeRefreshTokens() + g, err := m.Grant(t.Context(), req) + if err != nil || g != nil { + t.Fatalf("Manager.Grant() with a revoked refresh token = %v, %v, want nil, nil", g, err) + } + if grants, _ := m.List(); len(grants) != 0 { + t.Fatalf("Manager.List() = %d grants, want the rejected grant deleted", len(grants)) + } +} + +func TestRevoke(t *testing.T) { + t.Parallel() + as := oauthtest.NewAuthServer(t, oauthtest.ServerOptions{}) + agent := oauthtest.NewProtectedAgent(t, as, oauthtest.AgentOptions{Card: oauthtest.CardMetadataURL}) + m := newManager(t, as) + _, g := acquire(t, m, agent.URL) + if err := m.Revoke(t.Context(), g); err != nil { + t.Fatalf("Manager.Revoke() error = %v", err) + } + if diff := cmp.Diff([]string{g.RefreshToken}, as.Revoked()); diff != "" { + t.Fatalf("revoked tokens wrong result (-want +got) diff = %s", diff) + } + if as.Valid(g.AccessToken, agent.URL, nil) { + t.Fatalf("access token still valid after revoking its refresh token") + } + + noRevoke := oauthtest.NewAuthServer(t, oauthtest.ServerOptions{NoRevocation: true}) + agent2 := oauthtest.NewProtectedAgent(t, noRevoke, oauthtest.AgentOptions{Card: oauthtest.CardMetadataURL}) + m2 := newManager(t, noRevoke) + _, g2 := acquire(t, m2, agent2.URL) + if err := m2.Revoke(t.Context(), g2); !errors.Is(err, oauth.ErrRevocationUnsupported) { + t.Fatalf("Manager.Revoke() without a revocation endpoint error = %v, want ErrRevocationUnsupported", err) + } +} + +func TestTransport(t *testing.T) { + t.Parallel() + as := oauthtest.NewAuthServer(t, oauthtest.ServerOptions{}) + agent := oauthtest.NewProtectedAgent(t, as, oauthtest.AgentOptions{}) + m := newManager(t, as) + client := &http.Client{Transport: m.Transport(nil)} + body := `{"message":{"messageId":"1","role":"ROLE_USER","parts":[{"text":"hi"}]}}` + + _, err := client.Post(agent.URL+"/message:send", "application/json", strings.NewReader(body)) + var ir *oauth.InteractionRequiredError + if !errors.As(err, &ir) { + t.Fatalf("POST without a grant error = %v, want InteractionRequiredError", err) + } + if _, err := m.Acquire(t.Context(), ir.Requirement, oauth.AcquireOptions{Prompter: oauth.Prompter{Open: oauthtest.Browser(t)}}); err != nil { + t.Fatal(err) + } + + post := func() int { + t.Helper() + resp, err := client.Post(agent.URL+"/message:send", "application/json", strings.NewReader(body)) + if err != nil { + t.Fatalf("POST error = %v", err) + } + _ = resp.Body.Close() + return resp.StatusCode + } + if got := post(); got != http.StatusOK { + t.Fatalf("POST with a stored grant status = %d, want 200", got) + } + as.Invalidate() + if got := post(); got != http.StatusOK { + t.Fatalf("POST after the token was invalidated status = %d, want 200 after a refresh", got) + } + if as.SawAuthorizationHeader() { + t.Fatalf("the authorization server received an Authorization header") + } + + explicit, err := http.NewRequest(http.MethodPost, agent.URL+"/message:send", strings.NewReader(body)) + if err != nil { + t.Fatal(err) + } + explicit.Header.Set("Authorization", "Bearer user-supplied") + resp, err := client.Do(explicit) + if err != nil { + t.Fatalf("POST with an explicit credential error = %v, want the 401 passed through", err) + } + _ = resp.Body.Close() + if resp.StatusCode != http.StatusUnauthorized { + t.Fatalf("POST with an explicit credential status = %d, want 401", resp.StatusCode) + } +} diff --git a/internal/oauth/transport.go b/internal/oauth/transport.go new file mode 100644 index 0000000..1b5990b --- /dev/null +++ b/internal/oauth/transport.go @@ -0,0 +1,141 @@ +// Copyright 2026 The A2A Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package oauth + +import ( + "io" + "net/http" +) + +// Transport returns an http.RoundTripper that handles 401 responses to requests +// made through base: it refreshes a rejected grant, or discovers the +// authorization server from the Bearer challenge and uses a stored grant, and +// retries the request once. When a grant has to be obtained interactively, the +// request fails with *InteractionRequiredError. Requests carrying an +// Authorization header the manager did not attach are passed through untouched. +func (m *Manager) Transport(base http.RoundTripper) http.RoundTripper { + if base == nil { + base = http.DefaultTransport + } + return &transport{base: base, m: m} +} + +type transport struct { + base http.RoundTripper + m *Manager +} + +// RoundTrip implements http.RoundTripper. +func (t *transport) RoundTrip(req *http.Request) (*http.Response, error) { + ctx := req.Context() + if req.Header.Get("Authorization") == "" && checkTokenDestination(req.URL.String()) == nil { + g, err := t.m.learnedGrant(ctx, req.URL) + if err != nil { + return nil, err + } + if g != nil { + req = withToken(req, g) + } + } + + resp, err := t.base.RoundTrip(req) + if err != nil || resp.StatusCode != http.StatusUnauthorized { + return resp, err + } + chs := ParseChallenges(resp.Header.Values("WWW-Authenticate")) + + if auth := req.Header.Get("Authorization"); auth != "" { + used, ok := attachedFrom(req.Context()) + if !ok || used.header() != auth { + return resp, nil + } + if bc, ok := bearerChallenge(chs); ok && bc.Params["error"] == "insufficient_scope" { + return resp, nil + } + t.m.logf("the agent rejected the access token for %s; refreshing", used.Key.Resource) + g, err := t.m.refreshRejected(ctx, used.Key, used.AccessToken) + if err != nil { + drain(resp) + return nil, err + } + if g == nil { + drain(resp) + return nil, &InteractionRequiredError{Requirement: used.requirement(), Reason: "the stored grant was rejected"} + } + return t.retry(req, resp, g) + } + + if !wantsBearer(chs) { + return resp, nil + } + r, err := t.m.RequirementForChallenge(ctx, req.URL, chs) + if err != nil { + drain(resp) + return nil, err + } + if r == nil { + return resp, nil + } + g, err := t.m.Grant(ctx, r) + if err != nil { + drain(resp) + return nil, err + } + if g == nil { + drain(resp) + return nil, &InteractionRequiredError{Requirement: r} + } + t.m.learn(g.Key) + return t.retry(req, resp, g) +} + +// retry repeats req once with the grant's token. A request whose body can't be +// replayed gets the original 401 back. +func (t *transport) retry(req *http.Request, resp *http.Response, g *Grant) (*http.Response, error) { + if req.Body != nil && req.Body != http.NoBody && req.GetBody == nil { + return resp, nil + } + drain(resp) + clone := withToken(req, g) + if req.GetBody != nil { + body, err := req.GetBody() + if err != nil { + return nil, err + } + clone.Body = body + } + return t.base.RoundTrip(clone) +} + +// wantsBearer reports whether a 401 may be OAuth-protected: it offers a Bearer +// challenge, or no challenge at all. +func wantsBearer(chs []Challenge) bool { + if len(chs) == 0 { + return true + } + _, ok := bearerChallenge(chs) + return ok +} + +func withToken(req *http.Request, g *Grant) *http.Request { + clone := req.Clone(withAttached(req.Context(), g)) + clone.Header.Set("Authorization", g.header()) + return clone +} + +func drain(resp *http.Response) { + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 64<<10)) + _ = resp.Body.Close() +}