From 684e70ba607947f01241b254cdce5dedf65dd8f4 Mon Sep 17 00:00:00 2001 From: seonNoh Date: Thu, 15 Oct 2026 11:48:26 +0900 Subject: [PATCH] feat(ai): authenticate to the Claude API with Workload Identity Federation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 정적 API 키 대신 Workload Identity Federation을 쓴다. 에이전트는 AI VM의 Google 서명 ID 토큰(메타데이터 서버, audience https://api.anthropic.com, format=full)을 POST /v1/oauth/token(jwt-bearer)으로 교환해 단기 access token을 받고 Authorization: Bearer로 호출한다. 토큰은 캐시하고 만료 120초 전에 재교환하며, 401이면 한 번 재교환한다. 규칙·조직·서비스 계정·워크스페이스 ID 4개는 비밀이 아니므로 ai.env 환경 변수로 받고, 비어 있으면 기존 키 파일 경로로 동작한다. /selfcheck가 auth 모드와 실제 교환 결과를 보고해 캠페인 없이 설정을 검증할 수 있다. 테스트: 교환 1회 재사용·401 재교환·만료 전 갱신·거부 응답·불완전 ID 무시. --- ai/cmd/agent/main.go | 12 ++ ai/internal/anthropic/client.go | 81 +++++++++-- ai/internal/anthropic/federation.go | 177 +++++++++++++++++++++++ ai/internal/anthropic/federation_test.go | 130 +++++++++++++++++ 4 files changed, 387 insertions(+), 13 deletions(-) create mode 100644 ai/internal/anthropic/federation.go create mode 100644 ai/internal/anthropic/federation_test.go diff --git a/ai/cmd/agent/main.go b/ai/cmd/agent/main.go index effe92b..7f5d8d1 100644 --- a/ai/cmd/agent/main.go +++ b/ai/cmd/agent/main.go @@ -99,7 +99,15 @@ func main() { // Claude Messages API for claude-* model ids chosen per campaign; the key may be registered later (empty = diagnoses with // a claude-* model fail with "anthropic api key is empty", recorded in the harness diagnosis event). anth := anthropic.New(env("ANTHROPIC_BASE_URL", ""), keyReader(env("ANTHROPIC_KEY_FILE", "/run/secrets/anthropic-api-key"))) + // Workload Identity Federation (preferred, no static key): the four ids are plain settings in ai.env; the VM's + // Google identity token is exchanged for a short-lived access token (internal/anthropic/federation.go). + anth.SetFederation(&anthropic.Federation{ + RuleID: env("ANTHROPIC_FEDERATION_RULE_ID", ""), OrgID: env("ANTHROPIC_ORGANIZATION_ID", ""), + ServiceAccountID: env("ANTHROPIC_SERVICE_ACCOUNT_ID", ""), WorkspaceID: env("ANTHROPIC_WORKSPACE_ID", ""), + IdentityTokenURL: env("ANTHROPIC_IDENTITY_TOKEN_URL", ""), + }) ag.SetAnthropic(anth) + log.Printf("anthropic auth: %s", anth.AuthMode()) jevMode := env("JEV_MODE", "off") // off | shadow | enforce (shadow records Jev answers next to the diagnosis; enforce also applies the two review rules) var jc *jev.Client if jevMode == "shadow" || jevMode == "enforce" { @@ -136,6 +144,10 @@ func main() { key, kerr := keyReader(env("OPENAI_KEY_FILE", "/run/secrets/openai-api-key"))() out["openai_key_present"] = kerr == nil && key != "" out["anthropic_key_present"] = anth.KeyPresent() + out["anthropic_auth"] = anth.AuthMode() + if f := anth.Federation(); f != nil { + out["anthropic_federation"] = f.Probe(c, anth.BaseURL) // one real token exchange; the Console history page shows the deny reason on failure + } out["jev_mode"] = jevMode if jc != nil { out["jev_key_present"] = jc.KeyPresent() diff --git a/ai/internal/anthropic/client.go b/ai/internal/anthropic/client.go index 0150523..26a656e 100644 --- a/ai/internal/anthropic/client.go +++ b/ai/internal/anthropic/client.go @@ -35,6 +35,7 @@ type Client struct { BaseURL string HTTP *http.Client keyFn func() (string, error) + fed *Federation // nil = API key from keyFn; set = Workload Identity Federation (takes precedence) mu sync.Mutex conv map[string]conversation // response id -> messages (and system prompt) up to and including that assistant turn @@ -60,6 +61,48 @@ func (c *Client) KeyPresent() bool { return err == nil && k != "" } +// SetFederation switches the client to Workload Identity Federation (ignored when the ids are incomplete). +func (c *Client) SetFederation(f *Federation) { + if f.Configured() { + c.fed = f + } +} + +// AuthMode is "federation", "api_key" or "none" (reported by /selfcheck). +func (c *Client) AuthMode() string { + if c.fed != nil { + return "federation" + } + if c.KeyPresent() { + return "api_key" + } + return "none" +} + +// Federation returns the active federation config (nil when the API key path is used). +func (c *Client) Federation() *Federation { return c.fed } + +// authorize sets the credential header: a federated bearer token, or the API key. +func (c *Client) authorize(ctx context.Context, req *http.Request) error { + if c.fed != nil { + tok, err := c.fed.AccessToken(ctx, c.BaseURL) + if err != nil { + return fmt.Errorf("anthropic federation: %w", err) + } + req.Header.Set("Authorization", "Bearer "+tok) + return nil + } + key, err := c.keyFn() + if err != nil { + return fmt.Errorf("anthropic key: %w", err) + } + if key == "" { + return fmt.Errorf("anthropic api key is empty (secret selfheal-anthropic-api-key not registered yet) and federation is not configured") + } + req.Header.Set("x-api-key", key) + return nil +} + type message struct { Role string `json:"role"` Content any `json:"content"` // string or []block @@ -114,12 +157,12 @@ const forcedNudge = "You must call the submit_diagnosis tool now with your final // CreateResponse translates one loop turn to a Messages API call and the answer back to Responses-API items. func (c *Client) CreateResponse(ctx context.Context, req openai.Request) (*openai.Response, error) { - key, err := c.keyFn() - if err != nil { - return nil, fmt.Errorf("anthropic key: %w", err) - } - if key == "" { - return nil, fmt.Errorf("anthropic api key is empty (secret selfheal-anthropic-api-key not registered yet)") + if c.fed == nil { + if key, err := c.keyFn(); err != nil { + return nil, fmt.Errorf("anthropic key: %w", err) + } else if key == "" { + return nil, fmt.Errorf("anthropic api key is empty (secret selfheal-anthropic-api-key not registered yet) and federation is not configured") + } } var msgs []message var system []string @@ -187,16 +230,28 @@ func (c *Client) CreateResponse(ctx context.Context, req openai.Request) (*opena } b, _ := json.Marshal(body) - hreq, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+"/v1/messages", bytes.NewReader(b)) + send := func() (*http.Response, error) { + hreq, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+"/v1/messages", bytes.NewReader(b)) + if err != nil { + return nil, err + } + if err := c.authorize(ctx, hreq); err != nil { + return nil, err + } + hreq.Header.Set("anthropic-version", apiVersion) + hreq.Header.Set("Content-Type", "application/json") + return c.HTTP.Do(hreq) + } + resp, err := send() if err != nil { return nil, err } - hreq.Header.Set("x-api-key", key) - hreq.Header.Set("anthropic-version", apiVersion) - hreq.Header.Set("Content-Type", "application/json") - resp, err := c.HTTP.Do(hreq) - if err != nil { - return nil, err + if resp.StatusCode == 401 && c.fed != nil { // token revoked/expired early: exchange once more + resp.Body.Close() + c.fed.Invalidate() + if resp, err = send(); err != nil { + return nil, err + } } defer resp.Body.Close() rb, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) diff --git a/ai/internal/anthropic/federation.go b/ai/internal/anthropic/federation.go new file mode 100644 index 0000000..eb4356a --- /dev/null +++ b/ai/internal/anthropic/federation.go @@ -0,0 +1,177 @@ +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 +} diff --git a/ai/internal/anthropic/federation_test.go b/ai/internal/anthropic/federation_test.go new file mode 100644 index 0000000..70f8dbc --- /dev/null +++ b/ai/internal/anthropic/federation_test.go @@ -0,0 +1,130 @@ +package anthropic + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai" +) + +// Federation: the identity token from the (fake) metadata server is exchanged once at /v1/oauth/token with the +// jwt-bearer grant and the four ids; /v1/messages gets "Authorization: Bearer" and no x-api-key; the cached token is +// reused; a 401 triggers exactly one re-exchange; an expired cache re-exchanges before the next call. +func TestFederationExchangeAndReuse(t *testing.T) { + var idCalls, exchanges, msgCalls int32 + var fail401 atomic.Bool + md := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Metadata-Flavor") != "Google" || !strings.Contains(r.URL.RawQuery, "audience=https://api.anthropic.com") || !strings.Contains(r.URL.RawQuery, "format=full") { + t.Errorf("bad metadata request: %s %v", r.URL.String(), r.Header) + } + atomic.AddInt32(&idCalls, 1) + w.Write([]byte("eyJhbGciOiJSUzI1NiJ9.eyJpc3MiOiJodHRwczovL2FjY291bnRzLmdvb2dsZS5jb20ifQ.sig")) + })) + defer md.Close() + api := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/v1/oauth/token": + var b map[string]string + _ = json.NewDecoder(r.Body).Decode(&b) + if b["grant_type"] != "urn:ietf:params:oauth:grant-type:jwt-bearer" || !strings.HasPrefix(b["assertion"], "eyJ") || b["federation_rule_id"] != "fdrl_1" || b["organization_id"] != "org-1" || b["service_account_id"] != "svac_1" || b["workspace_id"] != "wrkspc_1" { + t.Errorf("bad exchange body: %v", b) + } + n := atomic.AddInt32(&exchanges, 1) + w.Write([]byte(`{"access_token":"sk-ant-oat01-tok` + string(rune('0'+n)) + `","token_type":"Bearer","expires_in":600}`)) + case "/v1/messages": + atomic.AddInt32(&msgCalls, 1) + if r.Header.Get("x-api-key") != "" || !strings.HasPrefix(r.Header.Get("Authorization"), "Bearer sk-ant-oat01-tok") { + t.Errorf("bad auth headers: %v", r.Header) + } + if fail401.Swap(false) { + w.WriteHeader(401) + w.Write([]byte(`{"type":"error","error":{"type":"authentication_error","message":"Authentication failed"}}`)) + return + } + w.Write([]byte(`{"id":"msg_f","type":"message","model":"claude-sonnet-5-5","stop_reason":"end_turn","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":3,"output_tokens":1}}`)) + default: + t.Errorf("unexpected path %s", r.URL.Path) + } + })) + defer api.Close() + c := New(api.URL, func() (string, error) { return "", nil }) // no key file: federation must carry the call + fed := &Federation{RuleID: "fdrl_1", OrgID: "org-1", ServiceAccountID: "svac_1", WorkspaceID: "wrkspc_1", IdentityTokenURL: md.URL + "/identity?audience=https://api.anthropic.com&format=full"} + c.SetFederation(fed) + if c.AuthMode() != "federation" { + t.Fatalf("auth mode %s", c.AuthMode()) + } + req := openai.Request{Model: "claude-sonnet-5-5", Input: []openai.Item{{Type: "message", Role: "user", Content: "hi"}}} + for i := 0; i < 2; i++ { + if _, err := c.CreateResponse(context.Background(), req); err != nil { + t.Fatal(err) + } + } + if atomic.LoadInt32(&exchanges) != 1 || atomic.LoadInt32(&idCalls) != 1 || atomic.LoadInt32(&msgCalls) != 2 { + t.Fatalf("expected one exchange for two calls, got exchanges=%d id=%d msgs=%d", exchanges, idCalls, msgCalls) + } + // 401 -> one re-exchange and a successful retry + fail401.Store(true) + if _, err := c.CreateResponse(context.Background(), req); err != nil { + t.Fatal(err) + } + if atomic.LoadInt32(&exchanges) != 2 || atomic.LoadInt32(&msgCalls) != 4 { + t.Fatalf("401 must re-exchange once: exchanges=%d msgs=%d", exchanges, msgCalls) + } + // cached token close to expiry -> advisory refresh before the next call + fed.mu.Lock() + fed.expiry = time.Now().Add(60 * time.Second) + fed.mu.Unlock() + if _, err := c.CreateResponse(context.Background(), req); err != nil { + t.Fatal(err) + } + if atomic.LoadInt32(&exchanges) != 3 { + t.Fatalf("near-expiry token must be refreshed: exchanges=%d", exchanges) + } + // probe reports a successful exchange + p := fed.Probe(context.Background(), api.URL) + if p["ok"] != true || p["expires_in_s"].(int) < 500 || !strings.HasPrefix(p["token_prefix"].(string), "sk-ant-oat01") { + t.Fatalf("probe: %v", p) + } +} + +func TestFederationFailuresAndPrecedence(t *testing.T) { + // incomplete ids are ignored: the key path stays active + c := New("http://127.0.0.1:1", func() (string, error) { return "k", nil }) + c.SetFederation(&Federation{RuleID: "fdrl_1"}) + if c.AuthMode() != "api_key" || c.Federation() != nil { + t.Fatalf("incomplete federation must not activate: %s", c.AuthMode()) + } + // exchange denied (opaque 401) -> clear error, no message call + api := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/v1/oauth/token" { + w.WriteHeader(401) + w.Write([]byte(`{"type":"error","error":{"type":"authentication_error","message":"Authentication failed"}}`)) + return + } + t.Errorf("messages must not be called when the exchange fails") + })) + defer api.Close() + md := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Write([]byte("a.b.c")) })) + defer md.Close() + c2 := New(api.URL, func() (string, error) { return "", nil }) + fed := &Federation{RuleID: "fdrl_1", OrgID: "org-1", ServiceAccountID: "svac_1", IdentityTokenURL: md.URL} + c2.SetFederation(fed) + _, err := c2.CreateResponse(context.Background(), openai.Request{Model: "claude-haiku-5-5", Input: []openai.Item{{Type: "message", Role: "user", Content: "x"}}}) + if err == nil || !strings.Contains(err.Error(), "token exchange: status 401") { + t.Fatalf("expected exchange failure, got %v", err) + } + if p := fed.Probe(context.Background(), api.URL); p["ok"] != false || !strings.Contains(p["error"].(string), "status 401") { + t.Fatalf("probe must report the denial: %v", p) + } + // metadata server unreachable + fed2 := &Federation{RuleID: "fdrl_1", OrgID: "org-1", ServiceAccountID: "svac_1", IdentityTokenURL: "http://127.0.0.1:1/identity", HTTP: &http.Client{Timeout: 2 * time.Second}} + if _, err := fed2.AccessToken(context.Background(), api.URL); err == nil || !strings.Contains(err.Error(), "identity token") { + t.Fatalf("expected identity token error, got %v", err) + } +}