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

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

070412f·

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

11 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
    }

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
    }

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.

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


541 unmodified lines

```go
// 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"))
        w.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)
    }
}