selfheal-lab-apps/ai/internal/anthropic/federation.go

package anthropic

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"
)

// Federation authenticates with Workload Identity Federation instead of an API key (platform.claude.com/docs/en/
// manage-claude/wif-providers/gcp): the VM's Google-signed identity token (metadata server, audience
// https://api.anthropic.com, format=full so the email claim is present) is exchanged at POST /v1/oauth/token
// (RFC 7523 jwt-bearer grant) for a short-lived sk-ant-oat01-... access token bound to an Anthropic service account.
// The four ids are not secrets; they come from ai.env. The token is cached and re-exchanged 120 s before expiry
// (the SDKs' advisory refresh), and a 401 from the API triggers one re-exchange. A fresh identity token is fetched
// for every exchange (identity tokens with a jti claim are single-use on the Anthropic side).
type Federation struct {
	RuleID           string // fdrl_...
	OrgID            string // Anthropic organization id (uuid)
	ServiceAccountID string // svac_...
	WorkspaceID      string // wrkspc_... (optional when the rule covers one workspace)
	IdentityTokenURL string // metadata server URL (tests override it)
	HTTP             *http.Client

	mu     sync.Mutex
	token  string
	expiry time.Time
	nex    int // exchanges performed (tests)
}

const defaultIdentityTokenURL = "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity?audience=https://api.anthropic.com&format=full"

const refreshBefore = 120 * time.Second

// Configured reports whether the mandatory ids are set (then federation takes precedence over the key file).
func (f *Federation) Configured() bool {
	return f != nil && f.RuleID != "" && f.OrgID != "" && f.ServiceAccountID != ""
}

func (f *Federation) identityToken(ctx context.Context) (string, error) {
	u := f.IdentityTokenURL
	if u == "" {
		u = defaultIdentityTokenURL
	}
	req, err := http.NewRequestWithContext(ctx, "GET", u, nil)
	if err != nil {
		return "", err
	}
	req.Header.Set("Metadata-Flavor", "Google")
	resp, err := f.http().Do(req)
	if err != nil {
		return "", fmt.Errorf("identity token: %w", err)
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 64<<10))
	tok := strings.TrimSpace(string(b))
	if resp.StatusCode != 200 || strings.Count(tok, ".") != 2 {
		return "", fmt.Errorf("identity token: status %d: %s", resp.StatusCode, truncate(tok, 200))
	}
	return tok, nil
}

func (f *Federation) http() *http.Client {
	if f.HTTP != nil {
		return f.HTTP
	}
	return &http.Client{Timeout: 30 * time.Second}
}

// AccessToken returns a cached access token, exchanging a fresh identity token when none is cached or the cached
// one is within refreshBefore of its expiry.
func (f *Federation) AccessToken(ctx context.Context, baseURL string) (string, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.token != "" && time.Until(f.expiry) > refreshBefore {
		return f.token, nil
	}
	tok, exp, err := f.exchange(ctx, baseURL)
	if err != nil {
		return "", err
	}
	f.token, f.expiry = tok, exp
	f.nex++
	return tok, nil
}

// Invalidate drops the cached token (after a 401) so the next call exchanges again.
func (f *Federation) Invalidate() {
	f.mu.Lock()
	f.token, f.expiry = "", time.Time{}
	f.mu.Unlock()
}

// Exchanges counts the token exchanges performed (tests).
func (f *Federation) Exchanges() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.nex
}

type tokenResponse struct {
	AccessToken string `json:"access_token"`
	TokenType   string `json:"token_type"`
	ExpiresIn   int    `json:"expires_in"`
	Error       string `json:"error"`
	ErrorDesc   string `json:"error_description"`
	ErrorObj    *struct {
		Type    string `json:"type"`
		Message string `json:"message"`
	} `json:"error_obj,omitempty"`
}

func (f *Federation) exchange(ctx context.Context, baseURL string) (string, time.Time, error) {
	jwt, err := f.identityToken(ctx)
	if err != nil {
		return "", time.Time{}, err
	}
	body := map[string]string{
		"grant_type":         "urn:ietf:params:oauth:grant-type:jwt-bearer",
		"assertion":          jwt,
		"federation_rule_id": f.RuleID,
		"organization_id":    f.OrgID,
		"service_account_id": f.ServiceAccountID,
	}
	if f.WorkspaceID != "" {
		body["workspace_id"] = f.WorkspaceID
	}
	b, _ := json.Marshal(body)
	req, err := http.NewRequestWithContext(ctx, "POST", strings.TrimRight(baseURL, "/")+"/v1/oauth/token", bytes.NewReader(b))
	if err != nil {
		return "", time.Time{}, err
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := f.http().Do(req)
	if err != nil {
		return "", time.Time{}, fmt.Errorf("token exchange: %w", err)
	}
	defer resp.Body.Close()
	rb, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	var tr tokenResponse
	_ = json.Unmarshal(rb, &tr)
	if resp.StatusCode/100 != 2 || tr.AccessToken == "" {
		// the API answers a failed exchange with an opaque 401 "Authentication failed"; the deny reason is on the
		// Console authentication-history page (platform.claude.com/settings/workload-identity-federation?tab=history)
		return "", time.Time{}, fmt.Errorf("token exchange: status %d: %s", resp.StatusCode, truncate(strings.TrimSpace(string(rb)), 300))
	}
	ttl := time.Duration(tr.ExpiresIn) * time.Second
	if ttl <= 0 {
		ttl = 60 * time.Second
	}
	return tr.AccessToken, time.Now().Add(ttl), nil
}

// Probe performs one exchange and reports the outcome (selfcheck): the setup can be verified without a diagnosis.
func (f *Federation) Probe(ctx context.Context, baseURL string) map[string]any {
	out := map[string]any{"rule_id": f.RuleID, "service_account_id": f.ServiceAccountID}
	f.Invalidate()
	start := time.Now()
	tok, err := f.AccessToken(ctx, baseURL)
	out["ms"] = time.Since(start).Milliseconds()
	if err != nil {
		out["ok"] = false
		out["error"] = err.Error()
		return out
	}
	f.mu.Lock()
	out["ok"] = true
	out["token_prefix"] = truncate(tok, 12)
	out["expires_in_s"] = int(time.Until(f.expiry).Seconds())
	f.mu.Unlock()
	return out
}