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() +}