Skip to content

Commit

Permalink
feat(core): improve entitlements performance (#1271)
Browse files Browse the repository at this point in the history 
### Context:

While developing COP, we found that GetEntitlements could take around
5.5 seconds to return a response for the federal dataset. This forced us
to increase the server timeout, increase the grpc message size, and
implement caches.

After investigating the latency, we determined the two main causes:
[excess
logs](https://github.com/opentdf/platform/blob/f275e25e4f986455fc536d2c93f7e9535f8519ab/service/authorization/authorization.go#L393)
(33%) and [excess database
queries](https://github.com/opentdf/platform/blob/cc15f25af2c3e839d7ad45283b7bd298a80e8728/service/policy/db/attribute_fqn.go#L188-L195)
(66%).

### Proposed Solution:

#### Primary
1. In the case of logging subject mappings, we now log their count
instead of their content.
2. In the case of database calls, we now list attributes, list subject
mappings, and match them based on their values. The database calls were
O(n) time complexity because they were dependent on the number of
attribute values. Now the database calls are constant time. We still
loop through all the values to match their subject mappings; however we
were already doing that in the `prepareValues` values function, so the
new approach is strictly better (especially due to our ubiquitous use of
maps).

#### Rego Query Optimization:
Yet if we simply match subject mappings and attribute values, the rego
query becomes massive (65 mb). It takes 3 seconds to build (20%) and
evaluate (80%). To optimize for not only time but also space, we remove
unrelated values for each fqn/attribute pair in the rego input (unless
the attribute rule is hierarchical).

After all the optimizations, fetching entitlements using the federal
dataset now takes about 125 ms. This is a **_latency reduction of
98%_**.

resolves: #1259

---------

Co-authored-by: Sean Trantalis <[email protected]>
  • Loading branch information
@alkalescent @strantalis
alkalescent and strantalis authored Aug 12, 2024
1 parent de5be3c commit f6a1b26
Show file tree
Hide file tree
Showing 3 changed files with 118 additions and 46 deletions.
124 changes: 95 additions & 29 deletions service/authorization/authorization.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import (
"github.com/opentdf/platform/protocol/go/entityresolution"
"github.com/opentdf/platform/protocol/go/policy"
attr "github.com/opentdf/platform/protocol/go/policy/attributes"
"github.com/opentdf/platform/protocol/go/policy/subjectmapping"
otdf "github.com/opentdf/platform/sdk"
"github.com/opentdf/platform/service/internal/access"
"github.com/opentdf/platform/service/internal/entitlements"
Expand Down Expand Up @@ -375,40 +376,107 @@ func (as *AuthorizationService) GetDecisions(ctx context.Context, req *authoriza
return rsp, nil
}

func (as *AuthorizationService) GetEntitlements(ctx context.Context, req *authorization.GetEntitlementsRequest) (*authorization.GetEntitlementsResponse, error) {
as.logger.DebugContext(ctx, "getting entitlements")
request := attr.GetAttributeValuesByFqnsRequest{
WithValue: &policy.AttributeValueSelector{
WithSubjectMaps: true,
},
// makeSubMapsByValLookup creates a lookup map of subject mappings by attribute value ID.
func makeSubMapsByValLookup(subjectMappings []*policy.SubjectMapping) map[string][]*policy.SubjectMapping {
// map keys will be attribute value IDs
lookup := make(map[string][]*policy.SubjectMapping)
for _, sm := range subjectMappings {
val := sm.GetAttributeValue()
id := val.GetId()
// if attribute value ID exists
if id != "" {
// append the subject mapping to the slice of subject mappings for the attribute value ID
lookup[id] = append(lookup[id], sm)
}
}
// Lack of scope has impacts on performance
// https://github.com/opentdf/platform/issues/365
if req.GetScope() == nil {
// TODO: Reomve and use MatchSubjectMappings instead later in the flow
listAttributeResp, err := as.sdk.Attributes.ListAttributes(ctx, &attr.ListAttributesRequest{})
if err != nil {
as.logger.ErrorContext(ctx, "failed to list attributes", slog.String("error", err.Error()))
return nil, status.Error(codes.Internal, "failed to list attributes")
return lookup
}

// updateValsWithSubMaps updates the subject mappings of values using the lookup map.
func updateValsWithSubMaps(values []*policy.Value, subMapsByVal map[string][]*policy.SubjectMapping) []*policy.Value {
for i, v := range values {
// if subject mappings exist for the value
if subjectMappings, ok := subMapsByVal[v.GetId()]; ok {
// update the subject mappings of the value
values[i].SubjectMappings = subjectMappings
}
var attributeFqns []string
for _, attr := range listAttributeResp.GetAttributes() {
for _, val := range attr.GetValues() {
attributeFqns = append(attributeFqns, val.GetFqn())
}
}
return values
}

// updateValsByFqnLookup updates the lookup map with attribute values by FQN.
func updateValsByFqnLookup(attribute *policy.Attribute, scopeMap map[string]bool, fqnAttrVals map[string]*attr.GetAttributeValuesByFqnsResponse_AttributeAndValue) map[string]*attr.GetAttributeValuesByFqnsResponse_AttributeAndValue {
rule := attribute.GetRule()
for _, v := range attribute.GetValues() {
// if scope exists and current attribute value FQN is not in scope
if !(scopeMap == nil || scopeMap[v.GetFqn()]) {
// skip
continue
}
// trim attribute values (by default only keep single value relevant to FQN)
// This is key to minimizing the rego query size.
values := []*policy.Value{v}
if rule == policy.AttributeRuleTypeEnum_ATTRIBUTE_RULE_TYPE_ENUM_HIERARCHY {
// restore ALL attribute values if attribute rule is hierarchical
// This is key to honoring comprehensive hierarchy.
values = attribute.GetValues()
}
request.Fqns = attributeFqns
} else {
// get subject mappings
request.Fqns = req.GetScope().GetAttributeValueFqns()
// only clone relevant fields for attribute
a := &policy.Attribute{Rule: rule, Values: values}
fqnAttrVals[v.GetFqn()] = &attr.GetAttributeValuesByFqnsResponse_AttributeAndValue{Attribute: a, Value: v}
}
return fqnAttrVals
}

// makeValsByFqnsLookup creates a lookup map of attribute values by FQN.
func makeValsByFqnsLookup(attrs []*policy.Attribute, subMapsByVal map[string][]*policy.SubjectMapping, scopeMap map[string]bool) map[string]*attr.GetAttributeValuesByFqnsResponse_AttributeAndValue {
fqnAttrVals := make(map[string]*attr.GetAttributeValuesByFqnsResponse_AttributeAndValue)
for i := range attrs {
// add subject mappings to attribute values
attrs[i].Values = updateValsWithSubMaps(attrs[i].GetValues(), subMapsByVal)
// update the lookup map with attribute values by FQN
fqnAttrVals = updateValsByFqnLookup(attrs[i], scopeMap, fqnAttrVals)
}
return fqnAttrVals
}

// makeScopeMap creates a map of attribute value FQNs.
func makeScopeMap(scope *authorization.ResourceAttribute) map[string]bool {
// if scope not defined, return nil pointer
if scope == nil {
return nil
}
scopeMap := make(map[string]bool)
// add attribute value FQNs from scope to the map
for _, fqn := range scope.GetAttributeValueFqns() {
scopeMap[fqn] = true
}
return scopeMap
}

func (as *AuthorizationService) GetEntitlements(ctx context.Context, req *authorization.GetEntitlementsRequest) (*authorization.GetEntitlementsResponse, error) {
as.logger.DebugContext(ctx, "getting entitlements")
attrsRes, err := as.sdk.Attributes.ListAttributes(ctx, &attr.ListAttributesRequest{})
if err != nil {
as.logger.ErrorContext(ctx, "failed to list attributes", slog.String("error", err.Error()))
return nil, status.Error(codes.Internal, "failed to list attributes")
}
avf, err := as.sdk.Attributes.GetAttributeValuesByFqns(ctx, &request)
subMapsRes, err := as.sdk.SubjectMapping.ListSubjectMappings(ctx, &subjectmapping.ListSubjectMappingsRequest{})
if err != nil {
as.logger.ErrorContext(ctx, "failed to get attribute values by fqns", slog.String("error", err.Error()))
return nil, status.Error(codes.Internal, "failed to get attribute values by fqns")
as.logger.ErrorContext(ctx, "failed to list subject mappings", slog.String("error", err.Error()))
return nil, status.Error(codes.Internal, "failed to list subject mappings")
}
// create a lookup map of attribute value FQNs (based on request scope)
scopeMap := makeScopeMap(req.GetScope())
// create a lookup map of subject mappings by attribute value ID
subMapsByVal := makeSubMapsByValLookup(subMapsRes.GetSubjectMappings())
// create a lookup map of attribute values by FQN (for rego query)
fqnAttrVals := makeValsByFqnsLookup(attrsRes.GetAttributes(), subMapsByVal, scopeMap)
avf := &attr.GetAttributeValuesByFqnsResponse{
FqnAttributeValues: fqnAttrVals,
}
subjectMappings := avf.GetFqnAttributeValues()
as.logger.DebugContext(ctx, "retrieved from subject mappings service", slog.Any("subject_mappings: ", subjectMappings))
as.logger.DebugContext(ctx, fmt.Sprintf("retrieved %d subject mappings", len(subjectMappings)))
// TODO: this could probably be moved to proto validation https://github.com/opentdf/platform/issues/1057
if req.Entities == nil {
as.logger.ErrorContext(ctx, "requires entities")
Expand All @@ -431,7 +499,6 @@ func (as *AuthorizationService) GetEntitlements(ctx context.Context, req *author
as.logger.ErrorContext(ctx, "failed to build rego input", slog.String("error", err.Error()))
return nil, status.Error(codes.Internal, "failed to build rego input")
}
as.logger.DebugContext(ctx, "entitlements", "input", fmt.Sprintf("%+v", in))

results, err := as.eval.Eval(ctx,
rego.EvalInput(in),
Expand Down Expand Up @@ -463,7 +530,6 @@ func (as *AuthorizationService) GetEntitlements(ctx context.Context, req *author
return rsp, nil
}
as.logger.DebugContext(ctx, "opa results", "results", fmt.Sprintf("%+v", results))

for idx, entity := range req.GetEntities() {
// Ensure the entity has an ID
entityID := entity.GetId()
Expand Down
31 changes: 23 additions & 8 deletions service/authorization/authorization_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import (
"github.com/opentdf/platform/protocol/go/entityresolution"
"github.com/opentdf/platform/protocol/go/policy"
attr "github.com/opentdf/platform/protocol/go/policy/attributes"
sm "github.com/opentdf/platform/protocol/go/policy/subjectmapping"
otdf "github.com/opentdf/platform/sdk"
"github.com/opentdf/platform/service/logger"
"github.com/stretchr/testify/assert"
Expand All @@ -25,6 +26,7 @@ import (
var (
getAttributesByValueFqnsResponse attr.GetAttributeValuesByFqnsResponse
listAttributeResp attr.ListAttributesResponse
listSubjectMappings sm.ListSubjectMappingsResponse
createEntityChainResp entityresolution.CreateEntityChainFromJwtResponse
resolveEntitiesResp entityresolution.ResolveEntitiesResponse
mockNamespace = "www.example.org"
Expand All @@ -51,6 +53,14 @@ type myERSClient struct {
entityresolution.EntityResolutionServiceClient
}

type mySubjectMappingClient struct {
sm.SubjectMappingServiceClient
}

func (*mySubjectMappingClient) ListSubjectMappings(_ context.Context, _ *sm.ListSubjectMappingsRequest, _ ...grpc.CallOption) (*sm.ListSubjectMappingsResponse, error) {
return &listSubjectMappings, nil
}

func (*myERSClient) CreateEntityChainFromJwt(_ context.Context, _ *entityresolution.CreateEntityChainFromJwtRequest, _ ...grpc.CallOption) (*entityresolution.CreateEntityChainFromJwtResponse, error) {
return &createEntityChainResp, nil
}
Expand Down Expand Up @@ -214,7 +224,8 @@ func Test_GetDecisionsAllOf_Pass(t *testing.T) {
}}

as := AuthorizationService{logger: logger, sdk: &otdf.SDK{
Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
SubjectMapping: &mySubjectMappingClient{},
Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
tokenSource: &testTokenSource, eval: prepared}

resp, err := as.GetDecisions(ctxb, &req)
Expand Down Expand Up @@ -372,7 +383,8 @@ func Test_GetDecisions_AllOf_Fail(t *testing.T) {
require.NoError(t, err)

as := AuthorizationService{logger: logger, sdk: &otdf.SDK{
Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
SubjectMapping: &mySubjectMappingClient{},
Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
tokenSource: &testTokenSource, eval: prepared}

resp, err := as.GetDecisions(ctxb, &req)
Expand Down Expand Up @@ -467,7 +479,8 @@ func Test_GetDecisionsAllOfWithEnvironmental_Pass(t *testing.T) {
}}

as := AuthorizationService{logger: logger, sdk: &otdf.SDK{
Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
SubjectMapping: &mySubjectMappingClient{},
Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
tokenSource: &testTokenSource, eval: prepared}

resp, err := as.GetDecisions(ctxb, &req)
Expand Down Expand Up @@ -560,7 +573,8 @@ func Test_GetDecisionsAllOfWithEnvironmental_Fail(t *testing.T) {
}}

as := AuthorizationService{logger: logger, sdk: &otdf.SDK{
Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
SubjectMapping: &mySubjectMappingClient{},
Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
tokenSource: &testTokenSource, eval: prepared}

resp, err := as.GetDecisions(ctxb, &req)
Expand Down Expand Up @@ -634,7 +648,8 @@ func Test_GetEntitlementsSimple(t *testing.T) {
require.NoError(t, err)

as := AuthorizationService{logger: logger, sdk: &otdf.SDK{
Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
SubjectMapping: &mySubjectMappingClient{},
Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
tokenSource: &testTokenSource, eval: prepared}

req := authorization.GetEntitlementsRequest{
Expand All @@ -653,8 +668,6 @@ func Test_GetEntitlementsSimple(t *testing.T) {

func Test_GetEntitlementsWithComprehensiveHierarchy(t *testing.T) {
logger := logger.CreateTestLogger()

listAttributeResp = attr.ListAttributesResponse{}
attrDef := policy.Attribute{
Name: mockAttrName,
Namespace: &policy.Namespace{
Expand All @@ -672,6 +685,7 @@ func Test_GetEntitlementsWithComprehensiveHierarchy(t *testing.T) {
},
},
}
listAttributeResp.Attributes = []*policy.Attribute{&attrDef}
getAttributesByValueFqnsResponse = attr.GetAttributeValuesByFqnsResponse{FqnAttributeValues: map[string]*attr.GetAttributeValuesByFqnsResponse_AttributeAndValue{
"https://www.example.org/attr/foo/value/value1": {
Attribute: &attrDef,
Expand Down Expand Up @@ -712,7 +726,8 @@ func Test_GetEntitlementsWithComprehensiveHierarchy(t *testing.T) {
prepared, err := rego.PrepareForEval(ctxb)
require.NoError(t, err)
as := AuthorizationService{logger: logger, sdk: &otdf.SDK{
Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
SubjectMapping: &mySubjectMappingClient{},
Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
tokenSource: &testTokenSource, eval: prepared}

withHierarchy := true
Expand Down
9 changes: 0 additions & 9 deletions service/policy/db/subject_mappings.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -189,15 +189,6 @@ func subjectMappingHydrateItem(row pgx.Row, logger *logger.Logger) (*policy.Subj
&scsJSON,
&attributeValueJSON,
)
logger.Debug(
"subjectMappingHydrateItem",
slog.Any("row", row),
slog.String("id", id),
slog.String("actionsJSON", string(actionsJSON)),
slog.String("metadataJSON", string(metadataJSON)),
slog.String("scsJSON", string(scsJSON)),
slog.String("attributeValueJSON", string(attributeValueJSON)),
)
if err != nil {
return nil, db.WrapIfKnownInvalidQueryErr(err)
}
Expand Down

0 comments on commit f6a1b26

Please sign in to comment.