cli/auth: add CellClientFactory — one subject, one token per jurisdiction

64230af·

Soph·1w ago·2 files·+123 added/-10 removed

NewEntireAPICellClient resolved the stored login subject (discovery + login refresh) and ran the RFC 8693 exchange on every call. For a single-cell command that's fine, but a multi-cell fan-out (one request per cell hosting the caller's repos, the BFF's code-search pattern) would pay all of it once per cell — even though identity tokens are per-jurisdiction, not per-cell: every cell in a jurisdiction accepts the same token.

CellClientFactory resolves the subject once at construction and caches minted identity tokens by jurisdiction; ClientFor(target) reuses them across cells. NewEntireAPICellClient stays as the single-cell wrapper (factory of one), so existing callers are unchanged.

Co-Authored-By: Claude Fable 5 noreply@anthropic.com

Sessions

935f91398e25View transcript

Changes

2

• cmd/entire/cli/auth

• Mcell_data_api.go+61/-10

• Mcell_data_api_test.go+62

11 unmodified lines

12
13
14
15
16
17
18
79 unmodified lines

98
99
100
100
101
102
103
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
109
137
138
139
140
141
142
143
144
145
114
146
147
148
149
2 unmodified lines

152
153
154
123
155
156
157
158
1 unmodified line

160
161
162
131
132
163
164
134
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193

11 unmodified lines

"os"
    "regexp"
    "strings"
    "sync"
    "time"

"github.com/entireio/cli/cmd/entire/cli/api"
79 unmodified lines

//   - otherwise the data host is a BFF/apex: resolve the caller's home-cell
//     apiUrl from the cluster catalog (home-jurisdiction fallback).
func NewEntireAPICellClient(ctx context.Context, insecureHTTP bool, target *CellTarget) (*api.Client, error) {
    // NewEntireAPICellClient deliberately does NOT consult ENTIRE_TOKEN: it
    // resolves the active stored login context (like every other cell/data-API
    // command). Only `entire auth token --jurisdiction` (JurisdictionToken) adds
    // the env-token path.
    factory, err := NewEntireAPICellClientFactory(ctx, insecureHTTP)
    if err != nil {
        return nil, err
    }
    return factory.ClientFor(ctx, target)
}

// CellClientFactory builds entire-api cell clients from a single resolved
// exchange subject, minting at most one jurisdictional identity token per
// jurisdiction. Identity tokens are per-jurisdiction, not per-cell — every cell
// in a jurisdiction accepts the same token — so a caller dialing several cells
// in one operation (multi-cell fan-out over the caller's repos) should build
// one factory and reuse it for every cell, instead of paying discovery + login
// refresh + RFC 8693 exchange once per cell via NewEntireAPICellClient.
//
// A factory is safe for concurrent use, and holds credentials resolved at
// construction time — build it per operation, don't store it long-term. Like
// NewEntireAPICellClient it deliberately does NOT consult ENTIRE_TOKEN.
type CellClientFactory struct {
    subject cellSubject

mu     sync.Mutex
    tokens map[string]string // jurisdiction -> minted identity token
}

// NewEntireAPICellClientFactory resolves the exchange subject (active stored
// login context) once, for building clients aimed at several cells. See
// NewEntireAPICellClient for the single-cell convenience wrapper.
func NewEntireAPICellClientFactory(ctx context.Context, insecureHTTP bool) (*CellClientFactory, error) {
    subject, err := resolveStoredCellSubject(ctx, insecureHTTP)
    if err != nil {
        return nil, err
    }
    return &CellClientFactory{subject: subject, tokens: make(map[string]string)}, nil
}

jurisdiction, err := targetJurisdiction(target, subject.loginJWT)
// ClientFor returns an authenticated client for the given cell target (nil
// falls back to home-jurisdiction routing), reusing an already-minted identity
// token when the target's jurisdiction matches an earlier call.
func (f *CellClientFactory) ClientFor(ctx context.Context, target *CellTarget) (*api.Client, error) {
    jurisdiction, err := targetJurisdiction(target, f.subject.loginJWT)
    if err != nil {
        return nil, err
    }

coreURL := jurisdictionCoreURL(jurisdiction, subject.dataOrigin, subject.discoveredCore)
    coreURL := jurisdictionCoreURL(jurisdiction, f.subject.dataOrigin, f.subject.discoveredCore)
    if err := requireSafeExchangeURL("entire-core", coreURL); err != nil {
        return nil, err
    }
2 unmodified lines

// which is signed by the discovered login core — so list there, not at the
    // templated jurisdiction core (coreURL), which in a multi-core setup could
    // differ and reject the token. coreURL still governs the token exchange below.
    cellBaseURL, err := resolveTargetCellBaseURL(ctx, target, subject.dataOrigin, jurisdiction, subject.discoveredCore, subject.loginJWT, subject.httpClient)
    cellBaseURL, err := resolveTargetCellBaseURL(ctx, target, f.subject.dataOrigin, jurisdiction, f.subject.discoveredCore, f.subject.loginJWT, f.subject.httpClient)
    if err != nil {
        return nil, err
    }
1 unmodified line

audience := jurisdictionAudience(jurisdiction, subject.dataOrigin, subject.discoveredCore)
    token, err := exchangeJurisdictionToken(ctx, coreURL, subject.loginJWT, audience, subject.httpClient)
    token, err := f.tokenFor(ctx, jurisdiction, coreURL)
    if err != nil {
        return nil, fmt.Errorf("exchange jurisdictional identity token: %w", err)
        return nil, err
    }

return api.NewClientWithBaseURL(token, cellBaseURL), nil
}

// tokenFor returns the cached identity token for jurisdiction, minting it on
// first use. The mutex is held across the mint: concurrent callers for the
// same jurisdiction wait for one exchange instead of duplicating it, at the
// cost of serializing cross-jurisdiction mints (fine for the handful of
// jurisdictions a fan-out touches).
func (f *CellClientFactory) tokenFor(ctx context.Context, jurisdiction, coreURL string) (string, error) {
    f.mu.Lock()
    defer f.mu.Unlock()
    if token, ok := f.tokens[jurisdiction]; ok {
        return token, nil
    }
    audience := jurisdictionAudience(jurisdiction, f.subject.dataOrigin, f.subject.discoveredCore)
    token, err := exchangeJurisdictionToken(ctx, coreURL, f.subject.loginJWT, audience, f.subject.httpClient)
    if err != nil {
        return "", fmt.Errorf("exchange jurisdictional identity token: %w", err)
    }
    f.tokens[jurisdiction] = token
    return token, nil
}

// JurisdictionToken mints and returns a jurisdictional identity token
// (scope=openid, aud=jurisdiction host) for `jurisdiction`, for authenticating
// against that jurisdiction's entire-api cells (e.g.

Mcmd/entire/cli/auth/cell_data_api.go+61/-10

541 unmodified lines

542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606

541 unmodified lines

t.Errorf("home-fallback audience = %q, want https://us.entire.io", got) }

// TestCellClientFactory_ReusesTokenPerJurisdiction pins the factory's core // contract: identity tokens are per-jurisdiction, not per-cell, so building // clients for several cells must mint one token per distinct jurisdiction and // reuse it across cells. Not parallel: manipulates env + token store. func TestCellClientFactory_ReusesTokenPerJurisdiction(t *testing.T) { t.Setenv("ENTIRE_CONFIG_DIR", t.TempDir()) t.Setenv("ENTIRE_API_BASE_URL", "https://entire.io") t.Setenv("ENTIRE_API_AUDIENCE_TEMPLATE", "") t.Setenv("ENTIRE_CORE_BASE_URL_TEMPLATE", "")

restore := tokenstore.UseFileBackendForTesting(filepath.Join(t.TempDir(), "tokens.json")) t.Cleanup(restore)

var exchangeCount int var audiences []string coreSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != oauthTokenPath { http.NotFound(w, r) return } _ = r.ParseForm() //nolint:errcheck // test handler exchangeCount++ audiences = append(audiences, r.FormValue("audience")) tw.Header().Set("Content-Type", "application/json") _, _ = fmt.Fprintf(w, {"access_token":"identity-token-%d","token_type":"Bearer","expires_in":3600}, exchangeCount) })) defer coreSrv.Close()

svc := tokenstore.CoreKeyringService(coreSrv.URL) loginJWT := makeJWT(t, fmt.Sprintf({"iss":%q,"home_jurisdiction":"us","exp":%d}, coreSrv.URL, time.Now().Add(2*time.Hour).Unix())) if err := tokenstore.Set(svc, "me", tokenstore.EncodeTokenWithExpiration(loginJWT, 7200)); err != nil { t.Fatalf("seed token: %v", err) } ctxObj := &contexts.Context{Name: "me@core", CoreURL: coreSrv.URL, Handle: "me", KeychainService: svc} t.Cleanup(SetResolveContextForCellAPIForTest(t, func(context.Context, string, string, string, *http.Client, clusterdiscovery.DebugFunc) (*contexts.Context, error) { return ctxObj, nil })) t.Cleanup(SetCellExchangeTransportForTest(t, coreSrv.Client().Transport))

factory, err := NewEntireAPICellClientFactory(context.Background(), false) if err != nil { t.Fatalf("NewEntireAPICellClientFactory: %v", err) }

// Two eu cells then one us cell: two exchanges total, the eu token reused // for the second eu cell. for _, target := range []*CellTarget{ {BaseURL: "https://cell-a.api.example", Jurisdiction: "eu"}, {BaseURL: "https://cell-b.api.example", Jurisdiction: "eu"}, {BaseURL: "https://cell-c.api.example", Jurisdiction: "us"}, } { if _, err := factory.ClientFor(context.Background(), target); err != nil { t.Fatalf("ClientFor(%s): %v", target.BaseURL, err) } } if exchangeCount != 2 { t.Fatalf("exchange count = %d, want 2 (one per distinct jurisdiction)", exchangeCount) } if got, want := strings.Join(audiences, ","), "https://eu.entire.io,"+usEntireAudience; got != want { t.Fatalf("exchange audiences = %q, want %q", got, want) } }