Enforce Cedar Policies on Upstream IDP Token

[tgrunnagle]

Problem

When the embedded auth server is active, Cedar policies are evaluated against
claims from the ToolHive-issued JWT. Upstream IDP tokens (GitHub, Okta, etc.)
are fetched and stored in identity.UpstreamTokens but are never read by the
Cedar authorizer. This means policies cannot reference upstream-specific claims
(e.g. GitHub login, Okta groups) that are absent from the ToolHive-issued token.

Two related gaps make group-based policies impossible even after switching to
upstream claims:

• Identity.Groups is never populated. The comment at pkg/auth/context.go:60
  says "Authorization logic MUST extract groups from the Claims map", but no code
  does this extraction. Groups is always nil.
• The Cedar entity hierarchy is never used. pkg/authz/authorizers/cedar/entity.go
  creates entities with empty Parents sets, so the Client → Group → Role
  hierarchy is never exercised. Policies like principal in THVGroup::"engineering"
  cannot match anything today.

Goal

When the embedded auth server is configured, use the access token from the
first configured upstream provider as the source of JWT claims for Cedar policy
evaluation, instead of the ToolHive-issued token's claims. Additionally, extract
group memberships from those claims, populate Identity.Groups, and wire group
entities into the Cedar entity hierarchy so that group-membership policies work.

Background: Current Data Flow

HTTP request
  │  Authorization: Bearer <ToolHive JWT>
  ▼
pkg/auth token.go – TokenValidator.Middleware()
  ├─ ExtractBearerToken()        → raw ToolHive JWT string
  ├─ ValidateToken()             → jwt.MapClaims (ToolHive claims)
  ├─ claimsToIdentity()          → Identity{Claims: <ToolHive claims>}
  └─ loadUpstreamTokens()        → Identity{UpstreamTokens: {"default": "<upstream token>"}}
       (uses tsid claim → storage lookup)
  │  WithIdentity(ctx, identity)
  ▼
pkg/authz/authorizers/cedar/core.go – AuthorizeWithJWTClaims()
  ├─ auth.IdentityFromContext()  → Identity
  ├─ jwt.MapClaims(identity.Claims)  ← ToolHive claims used here
  ├─ extractClientIDFromClaims() → sub from ToolHive JWT
  └─ authorize*(clientID, ...)   → Cedar decision
       identity.UpstreamTokens  ← NEVER READ
       identity.Groups           ← ALWAYS EMPTY
       entity.Parents            ← ALWAYS EMPTY SET

Proposed Change

Add a PrimaryUpstreamProvider option to the Cedar authorizer. When set,
AuthorizeWithJWTClaims reads the named upstream token from
identity.UpstreamTokens, parses its JWT claims (without re-validating the
signature — it was already validated during the OAuth exchange), and uses those
claims as the Cedar principal and context instead.

Additionally, extract group claims from the resolved claims (whichever source),
populate Identity.Groups, and create Group entities in the Cedar entity map
as parents of the principal entity so that Cedar's in operator works for
group membership.

---

Implementation Steps

Step 1 — Add PrimaryUpstreamProvider to Cedar ConfigOptions

File: pkg/authz/authorizers/cedar/core.go

Add the field to ConfigOptions:

type ConfigOptions struct {
    Policies                []string `json:"policies" yaml:"policies"`
    EntitiesJSON            string   `json:"entities_json" yaml:"entities_json"`
    // PrimaryUpstreamProvider names the upstream IDP provider whose access
    // token should be used as the source of JWT claims for Cedar evaluation.
    // When empty, claims from the ToolHive-issued token are used (current behaviour).
    // Must match an entry in identity.UpstreamTokens (e.g. "default", "github").
    PrimaryUpstreamProvider string `json:"primary_upstream_provider,omitempty" yaml:"primary_upstream_provider,omitempty"`
}

Add the field to the Authorizer struct:

type Authorizer struct {
    policySet               *cedar.PolicySet
    entities                cedar.EntityMap
    entityFactory           *EntityFactory
    mu                      sync.RWMutex
    primaryUpstreamProvider string  // new
}

Store it in NewCedarAuthorizer:

func NewCedarAuthorizer(options ConfigOptions) (authorizers.Authorizer, error) {
    authorizer := &Authorizer{
        ...
        primaryUpstreamProvider: options.PrimaryUpstreamProvider,
    }
    ...
}

Step 2 — Add upstream JWT parsing helper

File: pkg/authz/authorizers/cedar/core.go

Add a private helper that parses a JWT string without signature verification.
Signature verification is intentionally skipped because the token was already
validated by the upstream IDP during the OAuth 2.0 code exchange. We only need
the claims.

// parseUpstreamJWTClaims parses JWT claims from an upstream access token without
// verifying the signature. The token was already validated by the upstream IDP
// during the OAuth 2.0 code exchange; we only need its claims for Cedar evaluation.
// Returns an error if the token is not a parseable JWT (e.g. opaque token).
func parseUpstreamJWTClaims(tokenStr string) (jwt.MapClaims, error) {
    parser := jwt.NewParser()
    token, _, err := parser.ParseUnverified(tokenStr, jwt.MapClaims{})
    if err != nil {
        return nil, fmt.Errorf("upstream token is not a parseable JWT: %w", err)
    }
    claims, ok := token.Claims.(jwt.MapClaims)
    if !ok {
        return nil, fmt.Errorf("upstream token has unexpected claims type")
    }
    return claims, nil
}

Step 3 — Branch in AuthorizeWithJWTClaims

File: pkg/authz/authorizers/cedar/core.go — AuthorizeWithJWTClaims (line ~581)

Replace the unconditional ToolHive claim extraction with a branch:

func (a *Authorizer) AuthorizeWithJWTClaims(
    ctx context.Context,
    feature authorizers.MCPFeature,
    operation authorizers.MCPOperation,
    resourceID string,
    arguments map[string]interface{},
) (bool, error) {
    identity, ok := auth.IdentityFromContext(ctx)
    if !ok {
        return false, ErrMissingPrincipal
    }

    var claims jwt.MapClaims

    if a.primaryUpstreamProvider != "" {
        // Embedded auth server path: use the upstream IDP token's claims.
        upstreamToken, tokenFound := identity.UpstreamTokens[a.primaryUpstreamProvider]
        if !tokenFound || upstreamToken == "" {
            // The upstream token must be present if the authorizer is configured to use it.
            // Missing token means the session has no upstream credential; deny.
            return false, fmt.Errorf("upstream token for provider %q not found in identity",
                a.primaryUpstreamProvider)
        }
        parsedClaims, err := parseUpstreamJWTClaims(upstreamToken)
        if err != nil {
            return false, fmt.Errorf("failed to parse upstream token for provider %q: %w",
                a.primaryUpstreamProvider, err)
        }
        claims = parsedClaims
    } else {
        // Default path: use ToolHive-issued token claims.
        claims = jwt.MapClaims(identity.Claims)
    }

    clientID, ok := extractClientIDFromClaims(claims)
    if !ok {
        return false, ErrMissingPrincipal
    }

    processedClaims := preprocessClaims(claims)
    processedArgs := preprocessArguments(arguments)

    switch {
    case feature == authorizers.MCPFeatureTool && operation == authorizers.MCPOperationCall:
        return a.authorizeToolCall(ctx, clientID, resourceID, processedClaims, processedArgs)
    case feature == authorizers.MCPFeaturePrompt && operation == authorizers.MCPOperationGet:
        return a.authorizePromptGet(clientID, resourceID, processedClaims, processedArgs)
    case feature == authorizers.MCPFeatureResource && operation == authorizers.MCPOperationRead:
        return a.authorizeResourceRead(clientID, resourceID, processedClaims, processedArgs)
    case operation == authorizers.MCPOperationList:
        return a.authorizeFeatureList(clientID, feature, processedClaims, processedArgs)
    default:
        return false, fmt.Errorf("unsupported feature/operation combination: %s/%s", feature, operation)
    }
}

Step 4 — Extract groups from claims and populate Identity.Groups

File: pkg/auth/context.go

The comment at line 60 says groups must be extracted by authorization logic, but no code
does this. Add a helper that checks common group claim names (groups, roles,
cognito:groups, etc.) and a function that updates Identity.Groups from the resolved
claims (ToolHive or upstream). This should be called from the same code path that selects
which claims to use in AuthorizeWithJWTClaims (Step 3), or alternatively extracted
into the Cedar authorizer since group claim names may be provider-specific.

// Common group claim names across popular identity providers.
var groupClaimNames = []string{"groups", "roles", "cognito:groups", "https://example.com/groups"}

// extractGroupsFromClaims looks for well-known group claim names in claims
// and returns the first non-empty match as a string slice.
func extractGroupsFromClaims(claims map[string]interface{}) []string {
    for _, name := range groupClaimNames {
        val, ok := claims[name]
        if !ok {
            continue
        }
        switch v := val.(type) {
        case []interface{}:
            groups := make([]string, 0, len(v))
            for _, g := range v {
                if s, ok := g.(string); ok {
                    groups = append(groups, s)
                }
            }
            if len(groups) > 0 {
                return groups
            }
        case []string:
            if len(v) > 0 {
                return v
            }
        }
    }
    return nil
}

The extracted groups should be set on Identity.Groups before the Cedar authorizer
runs so that the entity factory can use them (Step 5). Whether this happens inside
AuthorizeWithJWTClaims or in the token middleware depends on whether the group claim
name needs to be configurable; a reasonable default is to extract on every authorization
call using the resolved claims.

Note: Consider adding a GroupClaimName field to ConfigOptions to let operators
specify the exact claim name rather than guessing from a fixed list, especially when
upstream providers use custom claim URIs.

Step 5 — Wire group entities into the Cedar entity hierarchy

File: pkg/authz/authorizers/cedar/entity.go

CreatePrincipalEntity currently creates entities with empty Parents sets, making
Cedar's principal in THVGroup::"engineering" operator always false. Update it to
accept a groups slice and create Group entities as parents:

// CreatePrincipalEntity creates a principal entity with the given ID, attributes,
// and group memberships. Each group is added as a parent entity so that Cedar's
// in operator works for group-based policies.
func (*EntityFactory) CreatePrincipalEntity(
    principalType, principalID string,
    attributes map[string]interface{},
    groups []string,
) (cedar.EntityUID, cedar.Entity, []cedar.Entity) {
    uid := cedar.NewEntityUID(cedar.EntityType(principalType), cedar.String(principalID))
    attrs := convertMapToCedarRecord(attributes)

    parents := cedar.NewEntityUIDSet()
    var groupEntities []cedar.Entity
    for _, g := range groups {
        groupUID := cedar.NewEntityUID("THVGroup", cedar.String(g))
        parents.Add(groupUID)
        groupEntities = append(groupEntities, cedar.Entity{
            UID:        groupUID,
            Parents:    cedar.NewEntityUIDSet(),
            Attributes: cedar.NewRecord(cedar.RecordMap{}),
            Tags:       cedar.NewRecord(cedar.RecordMap{}),
        })
    }

    entity := cedar.Entity{
        UID:        uid,
        Parents:    parents,
        Attributes: attrs,
        Tags:       cedar.NewRecord(cedar.RecordMap{}),
    }

    return uid, entity, groupEntities
}

Update CreateEntitiesForRequest to accept a groups slice, call the updated
CreatePrincipalEntity, and add the returned group entities to the EntityMap.

Step 6 — Wire the upstream provider name in the proxy runner

File: pkg/runner/middleware.go (~line 148)

When building authzParams, if the embedded auth server is configured, inject
the first upstream provider name into the Cedar config before passing it to
the authz factory:

if config.AuthzConfig != nil {
    authzConfig := config.AuthzConfig
    // When the embedded auth server is active, enrich the Cedar config with
    // the primary upstream provider name so Cedar uses upstream IDP claims.
    if config.EmbeddedAuthServerConfig != nil && len(config.EmbeddedAuthServerConfig.Upstreams) > 0 {
        providerName := authserver.ResolveUpstreamName(config.EmbeddedAuthServerConfig.Upstreams[0].Name)
        enrichedConfig, err := injectUpstreamProvider(authzConfig, providerName)
        if err != nil {
            slog.Warn("could not inject upstream provider into Cedar config, using ToolHive claims",
                "error", err)
        } else {
            authzConfig = enrichedConfig
        }
    }
    authzParams := authz.FactoryMiddlewareParams{
        ConfigPath: config.AuthzConfigPath,
        ConfigData: authzConfig,
    }
    ...
}

Add a helper function injectUpstreamProvider in pkg/runner/middleware.go (or
a new pkg/runner/authz_helper.go):

// injectUpstreamProvider returns a copy of cfg with the Cedar
// PrimaryUpstreamProvider field set to providerName. It returns an error
// if cfg is not a Cedar config (other authorizer types are passed through
// unchanged since they don't use this field).
func injectUpstreamProvider(cfg *authz.Config, providerName string) (*authz.Config, error) {
    var raw cedar.Config
    if err := json.Unmarshal(cfg.RawConfig(), &raw); err != nil || raw.Options == nil {
        // Not a Cedar config or unparseable — return unchanged.
        return cfg, nil
    }
    raw.Options.PrimaryUpstreamProvider = providerName
    return authz.NewConfig(raw)
}

This function is a no-op for non-Cedar configs and is therefore safe to call
unconditionally whenever the embedded auth server is active.

Note on config_builder.go: addAuthzMiddleware (line ~672) is a separate
code path used by the Kubernetes operator. It should receive the same treatment
— pass the provider name through when the builder also has an embedded auth
server config. This can be done by adding a providerName string parameter to
addAuthzMiddleware.

Step 7 — Wire the upstream provider name in vMCP

File: pkg/vmcp/auth/factory/incoming.go

newCedarAuthzMiddleware builds cedar.ConfigOptions directly. It needs the
upstream provider name. The cleanest approach is to add an optional
primaryUpstreamProvider field to config.AuthzConfig (vMCP config struct)
and read it here:

// In newCedarAuthzMiddleware:
cedarConfig := cedar.Config{
    Version: "1.0",
    Type:    cedar.ConfigType,
    Options: &cedar.ConfigOptions{
        Policies:                authzCfg.Policies,
        EntitiesJSON:            "[]",
        PrimaryUpstreamProvider: authzCfg.PrimaryUpstreamProvider, // new
    },
}

File: pkg/vmcp/config/config.go — AuthzConfig struct:

type AuthzConfig struct {
    Type                    string   `json:"type" yaml:"type"`
    Policies                []string `json:"policies" yaml:"policies"`
    // PrimaryUpstreamProvider is the name of the upstream IDP whose token
    // claims are used for Cedar policy evaluation. Leave empty to use the
    // ToolHive-issued token claims.
    PrimaryUpstreamProvider string `json:"primary_upstream_provider,omitempty" yaml:"primary_upstream_provider,omitempty"`
}

---

Error Handling

Condition	Behaviour
primaryUpstreamProvider is set but identity.UpstreamTokens has no entry for it	Return error → deny. This indicates a misconfiguration or a missing upstream credential.
Upstream token present but not a parseable JWT (opaque token)	Return error → deny. The operator must configure an OIDC upstream (which issues JWT access tokens) for this feature. Do not fall back silently to ToolHive claims.
Upstream token present and parseable but has no sub claim	Return ErrMissingPrincipal → deny (same as today for ToolHive tokens).
primaryUpstreamProvider is empty	Use existing ToolHive claims path — no behaviour change.
Group claims present in upstream token	Populate Identity.Groups and create Group entities in Cedar entity map.
No recognizable group claim in upstream token	Identity.Groups remains nil; no Group parent entities added. Policies using principal in THVGroup::"..." evaluate to false (not an error).

The strict no-fallback policy is intentional: silent fallback would mean a
misconfigured deployment silently authorises against different claims than the
operator intended.

---

Testing

Unit tests — pkg/authz/authorizers/cedar/

• parseUpstreamJWTClaims with a valid JWT, an opaque token, and a malformed string
• AuthorizeWithJWTClaims when primaryUpstreamProvider is:
  • Empty (existing tests must still pass unchanged)
  • Set and token present with valid claims → Cedar decision uses upstream sub
  • Set and token missing from identity.UpstreamTokens → error
  • Set and token is an opaque string → error
• CreatePrincipalEntity with a non-empty groups slice → Parents set contains Group UIDs, group entities are returned
• CreateEntitiesForRequest with groups → Group entities present in the returned EntityMap
• End-to-end Cedar policy evaluation using principal in THVGroup::"engineering" → allow/deny correctly

Unit tests — pkg/auth/

• extractGroupsFromClaims with groups, roles, cognito:groups, and unknown claim names
• extractGroupsFromClaims with empty and missing claim values → returns nil

Unit tests — pkg/runner/

• injectUpstreamProvider with a Cedar config → provider name injected
• injectUpstreamProvider with a non-Cedar config → returns config unchanged
• middleware.go middleware-build path: when both AuthzConfig and
  EmbeddedAuthServerConfig are set, the constructed FactoryMiddlewareParams
  carries the enriched Cedar config with provider name set

Unit tests — pkg/vmcp/auth/factory/

• newCedarAuthzMiddleware propagates PrimaryUpstreamProvider from
  config.AuthzConfig into cedar.ConfigOptions

---

Files Changed

File	Change
pkg/authz/authorizers/cedar/core.go	Add PrimaryUpstreamProvider to ConfigOptions and Authorizer; add parseUpstreamJWTClaims; branch in AuthorizeWithJWTClaims
pkg/authz/authorizers/cedar/entity.go	Update CreatePrincipalEntity to accept groups, create Group parent entities; update CreateEntitiesForRequest
pkg/auth/context.go	Add extractGroupsFromClaims helper; call it during authorization to populate Identity.Groups
pkg/runner/middleware.go	Call injectUpstreamProvider when embedded auth server + authz are both configured
pkg/runner/middleware.go or pkg/runner/authz_helper.go	Add injectUpstreamProvider helper
pkg/runner/config_builder.go	Propagate provider name through addAuthzMiddleware
pkg/vmcp/config/config.go	Add PrimaryUpstreamProvider to AuthzConfig
pkg/vmcp/auth/factory/incoming.go	Pass PrimaryUpstreamProvider into cedar.ConfigOptions
Test files for each of the above	New cases per Testing section

---

Non-goals / Out of Scope

• Re-validating the upstream token's signature. The signature was already
  checked by the IDP at exchange time; re-checking requires the IDP's JWKS and
  adds latency for no security gain in this trust model.
• Supporting opaque upstream tokens (OAuth 2.0 providers that don't issue JWT
  access tokens). Those would require token introspection — a separate feature.
• Multi-upstream round-robin or fallback. Only the first configured
  upstream is used; the config field is named PrimaryUpstreamProvider to
  signal this is a deliberate single-provider selection.