fix(authz): optimize subject mapping evaluation performance

lib/flattening/flatten.go

@@ -7,6 +7,8 @@ import (
 
 type Flattened struct {
 	Items []Item `json:"flattened"`
+	// index provides O(1) selector lookups; populated by Flatten
+	index map[string][]interface{}
 }
 
 type Item struct {
@@ -15,6 +17,14 @@ type Item struct {
 }
 
 func GetFromFlattened(flat Flattened, selector string) []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 for backwards compatibility (e.g., if Flattened was created without Flatten)
 	itemsToReturn := []interface{}{}
 	for _, item := range flat.Items {
 		if item.Key == selector {
@@ -25,43 +35,52 @@ func GetFromFlattened(flat Flattened, selector string) []interface{} {
 }
 
 func Flatten(m map[string]interface{}) (Flattened, error) {
-	flattened := Flattened{}
-	items, err := flattenInterface(m)
+	idx := make(map[string][]interface{})
+	items, err := flattenInterface(m, idx)
 	if err != nil {
 		return Flattened{}, err
 	}
-	flattened.Items = items
-	return flattened, nil
+	return Flattened{
+		Items: items,
+		index: idx,
+	}, nil
 }
 
-func flattenInterface(i interface{}) ([]Item, error) {
+func flattenInterface(i interface{}, idx map[string][]interface{}) ([]Item, error) {
 	o := []Item{}
 	switch child := i.(type) {
 	case map[string]interface{}:
 		for k, v := range child {
-			nm, err := flattenInterface(v)
+			nm, err := flattenInterface(v, idx)
 			if err != nil {
 				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})
+				idx[key] = append(idx[key], item.Value)
 			}
 		}
 	case []interface{}:
-		for idx, item := range child {
-			k := fmt.Sprintf("[%v]", idx)
-			k2 := "[]"
-			flattenedItem, err := flattenInterface(item)
+		for index, item := range child {
+			kIdx := fmt.Sprintf("[%v]", index)
+			kAny := "[]"
+			flattenedItem, err := flattenInterface(item, idx)
 			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})
+				idx[keyIdx] = append(idx[keyIdx], it.Value)
+				idx[keyAny] = append(idx[keyAny], it.Value)
 			}
 		}
 	case bool, int, string, float64, float32:
 		o = append(o, Item{Key: "", Value: child})
+		idx[""] = append(idx[""], child)
 	default:
 		return nil, errors.New("unrecognized item in json")
 	}

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)
+		}
+	}
+}

lib/flattening/flatten_test.go

@@ -344,7 +344,8 @@ func TestFlattenInterfaceNoPanic(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			require.NotPanics(t, func() {
-				_, _ = flattenInterface(tc.value)
+				idx := make(map[string][]interface{})
+				_, _ = flattenInterface(tc.value, idx)
 			})
 		})
 	}

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