diff --git a/service/authorization/v2/authorization.go b/service/authorization/v2/authorization.go index e49e6a1216..ce97a73cef 100644 --- a/service/authorization/v2/authorization.go +++ b/service/authorization/v2/authorization.go @@ -277,6 +277,8 @@ func (as *Service) GetDecisionBulk(ctx context.Context, req *connect.Request[aut return nil, statusifyError(ctx, as.logger, errors.Join(ErrFailedToInitPDP, err)) } + pdp = pdp.WithRequestReuse() + multiRequests := req.Msg.GetDecisionRequests() decisionResponses := make([]*authzV2.GetDecisionMultiResourceResponse, len(multiRequests)) @@ -285,8 +287,6 @@ func (as *Service) GetDecisionBulk(ctx context.Context, req *connect.Request[aut return nil, statusifyError(ctx, as.logger, err) } - // TODO: revisit performance of this loop after introduction of caching and registered resource values within decisioning, - // as the same entity in multiple requests should only be resolved JIT once, not once per request if the same in each. for idx, request := range multiRequests { entityIdentifier := request.GetEntityIdentifier() action := request.GetAction() diff --git a/service/internal/access/v2/just_in_time_pdp.go b/service/internal/access/v2/just_in_time_pdp.go index 61e97b6be6..d41d567e99 100644 --- a/service/internal/access/v2/just_in_time_pdp.go +++ b/service/internal/access/v2/just_in_time_pdp.go @@ -38,6 +38,7 @@ var ( ) type JustInTimePDP struct { + reuse *requestReuse logger *logger.Logger sdk *otdfSDK.SDK // embedded obligations PDP @@ -346,7 +347,7 @@ func (p *JustInTimePDP) GetEntitlements( // buildInnerPDP fetches the entitleable attributes for the provided value FQNs and constructs a // request-scoped PolicyDecisionPoint from them plus the fully-loaded registered resources and // dynamic value mappings. -func (p *JustInTimePDP) buildInnerPDP(ctx context.Context, valueFQNs []string) (*PolicyDecisionPoint, error) { +func (p *JustInTimePDP) buildInnerPDPUncached(ctx context.Context, valueFQNs []string) (*PolicyDecisionPoint, error) { // Direct entitlements / dynamic value mappings require the full policy load (see NewJustInTimePDP). if p.fullPolicyPDP != nil { return p.fullPolicyPDP, nil @@ -457,7 +458,7 @@ func (p *JustInTimePDP) getMatchedSubjectMappings( } // resolveEntitiesFromEntityChain roundtrips caller-provided entity chains through ERS. -func (p *JustInTimePDP) resolveEntitiesFromEntityChain( +func (p *JustInTimePDP) resolveEntitiesFromEntityChainUncached( ctx context.Context, entityChain *entity.EntityChain, skipEnvironmentEntities bool, @@ -548,7 +549,7 @@ func entityRepresentationsFromResolvedChain(entityChain *entity.EntityChain, ski // resolveEntitiesFromToken roundtrips to ERS to resolve the provided token // and optionally skips environment entities (which is expected behavior in decision flow) -func (p *JustInTimePDP) resolveEntitiesFromToken( +func (p *JustInTimePDP) resolveEntitiesFromTokenUncached( ctx context.Context, token *entity.Token, skipEnvironmentEntities bool, diff --git a/service/internal/access/v2/request_reuse.go b/service/internal/access/v2/request_reuse.go new file mode 100644 index 0000000000..48b9b169d5 --- /dev/null +++ b/service/internal/access/v2/request_reuse.go @@ -0,0 +1,95 @@ +package access + +import ( + "context" + "fmt" + "slices" + "strings" + + authzV2 "github.com/opentdf/platform/protocol/go/authorization/v2" + "github.com/opentdf/platform/protocol/go/entity" + entityresolutionV2 "github.com/opentdf/platform/protocol/go/entityresolution/v2" + attrs "github.com/opentdf/platform/protocol/go/policy/attributes" + "google.golang.org/protobuf/proto" +) + +const ( + maxRequestPolicyCacheEntries = 8 + maxRequestResolutionCacheEntries = 32 +) + +type resolutionCacheKey struct { + messageType string + request string + skipEnvironment bool +} + +type requestReuse struct { + policies map[string]*PolicyDecisionPoint + resolutions map[resolutionCacheKey][]*entityresolutionV2.EntityRepresentation +} + +// WithRequestReuse returns a PDP for sequential subrequests within one bulk RPC. +// Its bounded memo tables must never be shared across RPCs or caller contexts. +func (p *JustInTimePDP) WithRequestReuse() *JustInTimePDP { + scoped := *p + scoped.reuse = &requestReuse{policies: make(map[string]*PolicyDecisionPoint), resolutions: make(map[resolutionCacheKey][]*entityresolutionV2.EntityRepresentation)} + return &scoped +} + +func (p *JustInTimePDP) buildInnerPDP(ctx context.Context, fqns []string) (*PolicyDecisionPoint, error) { + if p.reuse == nil || p.fullPolicyPDP != nil { + return p.buildInnerPDPUncached(ctx, fqns) + } + normalized := make([]string, len(fqns)) + for i, fqn := range fqns { + normalized[i] = strings.ToLower(fqn) + } + slices.Sort(normalized) + normalized = slices.Compact(normalized) + encoded, err := proto.Marshal(&attrs.GetEntitleableAttributesByFqnsRequest{Fqns: normalized}) + if err != nil { + return nil, fmt.Errorf("encode policy lookup: %w", err) + } + key := string(encoded) + if cached, ok := p.reuse.policies[key]; ok { + return cached, nil + } + pdp, err := p.buildInnerPDPUncached(ctx, normalized) + if err == nil && len(p.reuse.policies) < maxRequestPolicyCacheEntries { + p.reuse.policies[key] = pdp + } + return pdp, err +} + +func (p *JustInTimePDP) reuseResolution(request proto.Message, skipEnvironment bool, resolve func() ([]*entityresolutionV2.EntityRepresentation, error)) ([]*entityresolutionV2.EntityRepresentation, error) { + if p.reuse == nil { + return resolve() + } + encoded, err := proto.MarshalOptions{Deterministic: true}.Marshal(request) + if err != nil { + return nil, fmt.Errorf("encode entity resolution request: %w", err) + } + key := resolutionCacheKey{messageType: string(request.ProtoReflect().Descriptor().FullName()), request: string(encoded), skipEnvironment: skipEnvironment} + if cached, ok := p.reuse.resolutions[key]; ok { + return cached, nil + } + representations, err := resolve() + if err == nil && len(p.reuse.resolutions) < maxRequestResolutionCacheEntries { + p.reuse.resolutions[key] = representations + } + return representations, err +} + +func (p *JustInTimePDP) resolveEntitiesFromEntityChain(ctx context.Context, chain *entity.EntityChain, skipEnvironment bool) ([]*entityresolutionV2.EntityRepresentation, error) { + return p.reuseResolution(chain, skipEnvironment, func() ([]*entityresolutionV2.EntityRepresentation, error) { + return p.resolveEntitiesFromEntityChainUncached(ctx, chain, skipEnvironment) + }) +} + +func (p *JustInTimePDP) resolveEntitiesFromToken(ctx context.Context, token *entity.Token, skipEnvironment bool, resources []*authzV2.Resource) ([]*entityresolutionV2.EntityRepresentation, error) { + request := &entityresolutionV2.CreateEntityChainsFromTokensRequest{Tokens: []*entity.Token{token}, Resources: resources} + return p.reuseResolution(request, skipEnvironment, func() ([]*entityresolutionV2.EntityRepresentation, error) { + return p.resolveEntitiesFromTokenUncached(ctx, token, skipEnvironment, resources) + }) +} diff --git a/service/internal/access/v2/request_reuse_test.go b/service/internal/access/v2/request_reuse_test.go new file mode 100644 index 0000000000..e2c94a50fe --- /dev/null +++ b/service/internal/access/v2/request_reuse_test.go @@ -0,0 +1,72 @@ +package access + +import ( + "fmt" + "testing" + + authzV2 "github.com/opentdf/platform/protocol/go/authorization/v2" + "github.com/opentdf/platform/protocol/go/entity" + entityresolutionV2 "github.com/opentdf/platform/protocol/go/entityresolution/v2" + otdfSDK "github.com/opentdf/platform/sdk" + "github.com/opentdf/platform/service/logger" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" +) + +func TestRequestReuseCachesPolicySetsAndStaysRequestScoped(t *testing.T) { + const def = "https://example.com/attr/department" + const fqn = def + "/value/engineering" + client := decisionAttrFake(def, fqn, "abc") + original := &JustInTimePDP{logger: logger.CreateTestLogger(), sdk: &otdfSDK.SDK{Attributes: client}} + scoped := original.WithRequestReuse() + first, err := scoped.buildInnerPDP(t.Context(), []string{fqn}) + require.NoError(t, err) + second, err := scoped.buildInnerPDP(t.Context(), []string{fqn, fqn}) + require.NoError(t, err) + require.Same(t, first, second) + require.Len(t, client.requests, 1) + fresh := original.WithRequestReuse() + third, err := fresh.buildInnerPDP(t.Context(), []string{fqn}) + require.NoError(t, err) + require.NotSame(t, first, third) + require.Len(t, client.requests, 2) + for i := 0; i < maxRequestPolicyCacheEntries+1; i++ { + _, err := scoped.buildInnerPDP(t.Context(), []string{fmt.Sprintf("%s/value/%d", def, i)}) + require.NoError(t, err) + } + require.Len(t, scoped.reuse.policies, maxRequestPolicyCacheEntries) +} + +func TestRequestReuseIncludesTokenResourcesAndEnvironmentMode(t *testing.T) { + client := &recordingERSV2Client{createResponse: &entityresolutionV2.CreateEntityChainsFromTokensResponse{ + EntityChains: []*entity.EntityChain{{Entities: []*entity.Entity{{EphemeralId: "subject", Category: entity.Entity_CATEGORY_SUBJECT, EntityType: &entity.Entity_Claims{Claims: claimsAnyForTest(t, map[string]interface{}{"name": "alice"})}}}}}, + }} + scoped := testJITPDP(client).WithRequestReuse() + token := &entity.Token{Jwt: "token", EphemeralId: "token-id"} + resources := attrValueResource("https://example.com/attr/a/value/one") + _, err := scoped.resolveEntitiesFromToken(t.Context(), token, true, resources) + require.NoError(t, err) + _, err = scoped.resolveEntitiesFromToken(t.Context(), proto.CloneOf(token), true, []*authzV2.Resource{proto.CloneOf(resources[0])}) + require.NoError(t, err) + require.Equal(t, 1, client.createCalls) + _, err = scoped.resolveEntitiesFromToken(t.Context(), token, true, attrValueResource("https://example.com/attr/a/value/two")) + require.NoError(t, err) + require.Equal(t, 2, client.createCalls) + _, err = scoped.resolveEntitiesFromToken(t.Context(), token, false, resources) + require.NoError(t, err) + require.Equal(t, 3, client.createCalls) + _, err = scoped.resolveEntitiesFromToken(t.Context(), &entity.Token{Jwt: "different", EphemeralId: "token-id"}, true, resources) + require.NoError(t, err) + require.Equal(t, 4, client.createCalls) +} + +func TestRequestReuseResolvesIdenticalEntityChainOnce(t *testing.T) { + client := &recordingERSV2Client{resolveResponse: &entityresolutionV2.ResolveEntitiesResponse{EntityRepresentations: []*entityresolutionV2.EntityRepresentation{entityRepWithClientID("abc")}}} + scoped := testJITPDP(client).WithRequestReuse() + chain := entityChainIdentifier().GetEntityChain() + for range 3 { + _, err := scoped.resolveEntitiesFromEntityChain(t.Context(), chain, true) + require.NoError(t, err) + } + require.Equal(t, 1, client.resolveCalls) +}