fix(authz): optimize subject mapping evaluation performance

lib/flattening/flatten.go

@@ -5,8 +5,14 @@ import (
 	"fmt"
 )
 
+// indexThreshold is the minimum number of items before we build an index.
+// For smaller structures, linear scan is faster than map overhead.
+const indexThreshold = 8
+
 type Flattened struct {
 	Items []Item `json:"flattened"`
+	// index provides O(1) selector lookups; populated by Flatten for structures >= indexThreshold
+	index map[string][]interface{}
 }
 
 type Item struct {
@@ -15,7 +21,15 @@ type Item struct {
 }
 
 func GetFromFlattened(flat Flattened, selector string) []interface{} {
-	itemsToReturn := []interface{}{}
+	// Fast-path: use prebuilt index for O(1) lookup
+	if flat.index != nil {
+		if vals, ok := flat.index[selector]; ok {
+			return vals
+		}
+		return nil
+	}
+	// Fallback: linear scan for small structures or backwards compatibility
+	var itemsToReturn []interface{}
 	for _, item := range flat.Items {
 		if item.Key == selector {
 			itemsToReturn = append(itemsToReturn, item.Value)
@@ -24,18 +38,39 @@ func GetFromFlattened(flat Flattened, selector string) []interface{} {
 	return itemsToReturn
 }
 
+// Flatten returns a Flattened struct with an index for O(1) lookups via GetFromFlattened.
+// For small structures (< indexThreshold items), the index is skipped and lookups use linear scan.
 func Flatten(m map[string]interface{}) (Flattened, error) {
-	flattened := Flattened{}
 	items, err := flattenInterface(m)
 	if err != nil {
 		return Flattened{}, err
 	}
-	flattened.Items = items
-	return flattened, nil
+
+	// Build index in a separate pass, only for larger structures
+	idx := buildIndex(items)
+
+	return Flattened{
+		Items: items,
+		index: idx,
+	}, nil
+}
+
+// buildIndex constructs the selector index from flattened items.
+// Returns nil for small structures where linear scan is faster.
+func buildIndex(items []Item) map[string][]interface{} {
+	if len(items) < indexThreshold {
+		return nil
+	}
+
+	idx := make(map[string][]interface{}, len(items))
+	for _, it := range items {
+		idx[it.Key] = append(idx[it.Key], it.Value)
+	}
+	return idx
 }
 
 func flattenInterface(i interface{}) ([]Item, error) {
-	o := []Item{}
+	var o []Item
 	switch child := i.(type) {
 	case map[string]interface{}:
 		for k, v := range child {
@@ -44,20 +79,23 @@ func flattenInterface(i interface{}) ([]Item, error) {
 				return nil, err
 			}
 			for _, item := range nm {
-				o = append(o, Item{Key: "." + k + item.Key, Value: item.Value})
+				key := "." + k + item.Key
+				o = append(o, Item{Key: key, Value: item.Value})
 			}
 		}
 	case []interface{}:
-		for idx, item := range child {
-			k := fmt.Sprintf("[%v]", idx)
-			k2 := "[]"
+		for index, item := range child {
+			kIdx := fmt.Sprintf("[%v]", index)
+			kAny := "[]"
 			flattenedItem, err := flattenInterface(item)
 			if err != nil {
 				return nil, err
 			}
-			for _, item := range flattenedItem {
-				o = append(o, Item{Key: k + item.Key, Value: item.Value})
-				o = append(o, Item{Key: k2 + item.Key, Value: item.Value})
+			for _, it := range flattenedItem {
+				keyIdx := kIdx + it.Key
+				keyAny := kAny + it.Key
+				o = append(o, Item{Key: keyIdx, Value: it.Value})
+				o = append(o, Item{Key: keyAny, Value: it.Value})
 			}
 		}
 	case bool, int, string, float64, float32:

lib/flattening/flatten_bench_large_test.go

@@ -0,0 +1,73 @@
+package flattening
+
+import (
+	"fmt"
+	"testing"
+)
+
+// Benchmark with larger dataset to show lookup performance
+func BenchmarkGetFromFlattened_Large(b *testing.B) {
+	// Create a larger flattened structure via Flatten
+	largeInput := make(map[string]interface{})
+	for i := 0; i < 100; i++ {
+		largeInput[fmt.Sprintf("key%d", i)] = fmt.Sprintf("value%d", i)
+	}
+	flatInput, err := Flatten(largeInput)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	// Query for a key in the middle
+	queryString := ".key50"
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		_ = GetFromFlattened(flatInput, queryString)
+	}
+}
+
+// Benchmark multiple lookups on same flattened entity
+func BenchmarkGetFromFlattened_MultipleLookups(b *testing.B) {
+	largeInput := make(map[string]interface{})
+	for i := 0; i < 50; i++ {
+		largeInput[fmt.Sprintf("attr%d", i)] = fmt.Sprintf("value%d", i)
+	}
+	flatInput, err := Flatten(largeInput)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	queries := []string{".attr0", ".attr10", ".attr25", ".attr40", ".attr49"}
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		for _, q := range queries {
+			_ = GetFromFlattened(flatInput, q)
+		}
+	}
+}
+
+// Benchmark with nested structure (common in entity representations)
+func BenchmarkFlatten_NestedEntity(b *testing.B) {
+	nestedInput := map[string]interface{}{
+		"user": map[string]interface{}{
+			"id":    "user123",
+			"name":  "Test User",
+			"email": "[email protected]",
+			"attributes": map[string]interface{}{
+				"department": "Engineering",
+				"level":      "Senior",
+				"groups":     []interface{}{"group1", "group2", "group3"},
+			},
+		},
+		"roles": []interface{}{
+			map[string]interface{}{"name": "admin", "scope": "global"},
+			map[string]interface{}{"name": "reader", "scope": "local"},
+		},
+	}
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		_, err := Flatten(nestedInput)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}

service/internal/access/v2/evaluate.go

@@ -367,25 +367,37 @@ func hierarchyRule(
 		}
 	}
 
-	// Check if the entitlements contain any values with index <= lowestValueFQNIndex
-	// This checks the requested value and any hierarchically higher values in a single pass - O(e) where e is entitlements count
-	for entitlementFQN, entitledActions := range entitlements {
-		// Check if this entitlement FQN has a valid index in the hierarchy
-		if idx, exists := valueFQNToIndex[entitlementFQN]; exists && idx <= lowestValueFQNIndex {
-			// Check if the required action is entitled
+	// If we didn't find any matching value indices, fail with all as missing
+	if lowestValueFQNIndex == len(attrValues) {
+		failures := make([]EntitlementFailure, 0, len(resourceValueFQNs))
+		for _, fqn := range resourceValueFQNs {
+			failures = append(failures, EntitlementFailure{
+				AttributeValueFQN: fqn,
+				ActionName:        actionName,
+			})
+		}
+		return failures
+	}
+
+	// Check entitlement at or above the hierarchy level:
+	// Scan only the FQNs in the definition up to the highest requested index (O(h))
+	// This is more efficient than scanning all entitlements (O(e)) when e >> h
+	for i := 0; i <= lowestValueFQNIndex; i++ {
+		candidateFQN := attrValues[i].GetFqn()
+		if entitledActions, ok := entitlements[candidateFQN]; ok {
 			for _, entitledAction := range entitledActions {
 				if strings.EqualFold(entitledAction.GetName(), actionName) {
 					l.DebugContext(ctx, "hierarchy rule satisfied",
 						slog.Group("entitled_by_value",
-							slog.String("FQN", entitlementFQN),
-							slog.Int("index", idx),
+							slog.String("FQN", candidateFQN),
+							slog.Int("index", i),
 						),
 						slog.Group("resource_highest_hierarchy_value",
 							slog.String("FQN", attrValues[lowestValueFQNIndex].GetFqn()),
 							slog.Int("index", lowestValueFQNIndex),
 						),
 					)
-					return nil // Found an entitled action at or above the hierarchy level, no failures
+					return nil
 				}
 			}
 		}

service/internal/access/v2/evaluate_test.go

@@ -630,6 +630,72 @@ func (s *EvaluateTestSuite) TestHierarchyRule() {
 	}
 }
 
+// TestHierarchyRule_DefinitionMissingValues tests the defensive code path where
+// resourceValueFQNs contains values not present in attrDefinition.GetValues().
+// This guards against data inconsistency between the policy service and attribute definitions.
+func (s *EvaluateTestSuite) TestHierarchyRule_DefinitionMissingValues() {
+	tests := []struct {
+		name              string
+		attrDefinition    *policy.Attribute
+		resourceValueFQNs []string
+		entitlements      subjectmappingbuiltin.AttributeValueFQNsToActions
+		expectedFailures  int
+	}{
+		{
+			name: "empty definition values - all resource FQNs should fail",
+			attrDefinition: &policy.Attribute{
+				Fqn:    levelFQN,
+				Rule:   policy.AttributeRuleTypeEnum_ATTRIBUTE_RULE_TYPE_ENUM_HIERARCHY,
+				Values: []*policy.Value{}, // Empty - simulates data inconsistency
+			},
+			resourceValueFQNs: []string{levelMidFQN, levelLowerMidFQN},
+			entitlements:      subjectmappingbuiltin.AttributeValueFQNsToActions{},
+			expectedFailures:  2,
+		},
+		{
+			name: "definition missing requested values",
+			attrDefinition: &policy.Attribute{
+				Fqn:  levelFQN,
+				Rule: policy.AttributeRuleTypeEnum_ATTRIBUTE_RULE_TYPE_ENUM_HIERARCHY,
+				Values: []*policy.Value{
+					// Only has highest - missing the values we're requesting
+					{Fqn: levelHighestFQN, Value: "highest"},
+				},
+			},
+			resourceValueFQNs: []string{levelMidFQN, levelLowerMidFQN}, // Not in definition
+			entitlements: subjectmappingbuiltin.AttributeValueFQNsToActions{
+				levelHighestFQN: []*policy.Action{actionRead},
+			},
+			expectedFailures: 2,
+		},
+	}
+
+	for _, tc := range tests {
+		s.Run(tc.name, func() {
+			failures := hierarchyRule(
+				s.T().Context(),
+				s.logger,
+				tc.entitlements,
+				s.action,
+				tc.resourceValueFQNs,
+				tc.attrDefinition,
+			)
+
+			s.Len(failures, tc.expectedFailures, "Expected %d failures but got %d", tc.expectedFailures, len(failures))
+
+			// Verify each resource FQN is reported as a failure
+			failedFQNs := make(map[string]bool)
+			for _, f := range failures {
+				failedFQNs[f.AttributeValueFQN] = true
+				s.Equal(s.action.GetName(), f.ActionName)
+			}
+			for _, fqn := range tc.resourceValueFQNs {
+				s.True(failedFQNs[fqn], "Expected FQN %s to be in failures", fqn)
+			}
+		})
+	}
+}
+
 // Test cases for evaluateDefinition
 func (s *EvaluateTestSuite) TestEvaluateDefinition() {
 	tests := []struct {

service/internal/access/v2/just_in_time_pdp.go

@@ -331,6 +331,7 @@ func (p *JustInTimePDP) getMatchedSubjectMappings(
 			}
 			for _, item := range flattened.Items {
 				if _, ok := subjectPropertySet[item.Key]; !ok {
+					subjectPropertySet[item.Key] = struct{}{} // Track key to avoid duplicates
 					subjectProperties = append(subjectProperties, &policy.SubjectProperty{
 						ExternalSelectorValue: item.Key,
 					})