From 5d45603a55794076d1e092e6c9d7ef58e6767adf Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Wed, 29 Nov 2023 10:22:29 -0800
Subject: [PATCH] `seccomp`: Make `extractRepeatedMatchers` more efficient.

This does the following:

- Only allocate maps once.
- Check whether the filter can run before doing any expensive allocation
  or map modifications.
- Recursively optimize other arguments earlier on.
- Replace `PerArg.Copy` with a specialized version.

This helps make this function more efficient in `gotsan` mode.

PiperOrigin-RevId: 586382284
---
 pkg/seccomp/seccomp_optimizer.go | 39 +++++++++++++++++---------------
 pkg/seccomp/seccomp_rules.go     | 11 +++++++--
 2 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/pkg/seccomp/seccomp_optimizer.go b/pkg/seccomp/seccomp_optimizer.go
index d71fd921f..3da7c7144 100644
--- a/pkg/seccomp/seccomp_optimizer.go
+++ b/pkg/seccomp/seccomp_optimizer.go
@@ -298,11 +298,12 @@ func extractRepeatedMatchers(rule SyscallRule) (SyscallRule, bool) {
 		}
 	}
 
+	allOtherMatchersSigs := make(map[string]struct{}, len(orRule))
+	argExprToOtherMatchersSigs := make(map[string]map[string]struct{}, len(orRule))
 	for argNum := 0; argNum < len(orRule[0].(PerArg)); argNum++ {
-		// Check if `argNum` takes on a set of matchers common for all
-		// combinations of all other matchers.
-		allOtherMatchersSigs := make(map[string]struct{}, len(orRule))
-		argExprToOtherMatchersSigs := make(map[string]map[string]struct{}, len(orRule))
+		// Check if this argNum is always AnyValue,
+		// or if all other arguments are always AnyValue.
+		// If either of that is true, there is nothing for this filter to do.
 		allArgNumMatchersAreAnyValue := true
 		allOtherMatchersAreAnyValue := true
 		for _, subRule := range orRule {
@@ -315,8 +316,19 @@ func extractRepeatedMatchers(rule SyscallRule) (SyscallRule, bool) {
 					allOtherMatchersAreAnyValue = allOtherMatchersAreAnyValue && isAnyValue
 				}
 			}
+		}
+		if allArgNumMatchersAreAnyValue || allOtherMatchersAreAnyValue {
+			// Cannot optimize.
+			continue
+		}
+		// Check if `argNum` takes on a set of matchers common for all
+		// combinations of all other matchers.
+		clear(allOtherMatchersSigs)
+		clear(argExprToOtherMatchersSigs)
+		for _, subRule := range orRule {
+			perArg := subRule.(PerArg)
 			repr := perArg[argNum].Repr()
-			otherMatchers := perArg.Copy().(PerArg)
+			otherMatchers := perArg.clone()
 			otherMatchers[argNum] = invalidValueMatcher{}
 			otherMatchersSig := otherMatchers.signature()
 			allOtherMatchersSigs[otherMatchersSig] = struct{}{}
@@ -325,10 +337,6 @@ func extractRepeatedMatchers(rule SyscallRule) (SyscallRule, bool) {
 			}
 			argExprToOtherMatchersSigs[repr][otherMatchersSig] = struct{}{}
 		}
-		if allArgNumMatchersAreAnyValue || allOtherMatchersAreAnyValue {
-			// Cannot optimize.
-			continue
-		}
 		// Now check if each possible repr of `argNum` got the same set of
 		// signatures for other matchers as `allOtherMatchersSigs`.
 		sameOtherMatchers := true
@@ -352,19 +360,14 @@ func extractRepeatedMatchers(rule SyscallRule) (SyscallRule, bool) {
 			perArg := subRule.(PerArg)
 			onlyArg := PerArg{AnyValue{}, AnyValue{}, AnyValue{}, AnyValue{}, AnyValue{}, AnyValue{}, AnyValue{}}
 			onlyArg[argNum] = perArg[argNum]
-			allExceptArg := perArg.Copy().(PerArg)
+			allExceptArg := perArg.clone()
 			allExceptArg[argNum] = AnyValue{}
 			argNumMatch[i] = onlyArg
 			otherArgsMatch[i] = allExceptArg
 		}
-		// Do not attempt to see if other arguments are also eligible for the
-		// same optimization, as this would complicate the logic of this
-		// already-complicated function further, and will be caught in future
-		// iterations of the optimizer anyway.
-		// Additionally, `argNumMatch` and `otherArgsMatch` may well be
-		// single-item, so the other (simpler) optimizers should run on them
-		// first.
-		return And{argNumMatch, otherArgsMatch}, true
+		// Attempt to optimize the "other" arguments:
+		otherArgsMatchOpt, _ := extractRepeatedMatchers(otherArgsMatch)
+		return And{argNumMatch, otherArgsMatchOpt}, true
 	}
 	return rule, false
 }
diff --git a/pkg/seccomp/seccomp_rules.go b/pkg/seccomp/seccomp_rules.go
index 38f0dda58..050594725 100644
--- a/pkg/seccomp/seccomp_rules.go
+++ b/pkg/seccomp/seccomp_rules.go
@@ -650,8 +650,10 @@ type PerArg [7]ValueMatcher // 6 arguments + RIP
 // instruction pointer.
 const RuleIP = 6
 
-// Copy implements `SyscallRule.Copy`.
-func (pa PerArg) Copy() SyscallRule {
+// clone returns a copy of this `PerArg`.
+// It is more efficient than `Copy` because it returns a `PerArg`
+// directly, rather than a `SyscallRule` interface.
+func (pa PerArg) clone() PerArg {
 	return PerArg{
 		pa[0],
 		pa[1],
@@ -663,6 +665,11 @@ func (pa PerArg) Copy() SyscallRule {
 	}
 }
 
+// Copy implements `SyscallRule.Copy`.
+func (pa PerArg) Copy() SyscallRule {
+	return pa.clone()
+}
+
 // Render implements `SyscallRule.Render`.
 func (pa PerArg) Render(program *syscallProgram, labelSet *labelSet) {
 	for i, arg := range pa {