From 09be6cec5e98985a3776286bdc29c26a979fafaa Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Thu, 5 Oct 2023 00:48:58 -0700
Subject: [PATCH] `bpf`: Use optimizer for seccomp-bpf programs.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This speeds up:

- The Sentry and the Gofer's seccomp-bpf programs
- User application seccomp-bpf programs
- Systrap and ptrace seccomp-bpf rules

Some benchmark results with Postgres:

```
                                   │    before     │                    opt                    │
                                   │    sec/op     │    sec/op     vs base                     │
SentrySystrap/Postgres/futex          88.82n ±  2%   81.29n ±  2%   -8.48% (p=0.000 n=519+510)
SentrySystrap/Postgres/nanosleep      116.9n ± 19%   115.9n ± 17%        ~ (p=0.859 n=350+317)
SentrySystrap/Postgres/sendmmsg       88.68n ±  1%   81.56n ±  1%   -8.04% (n=519+510)
SentrySystrap/Postgres/fstat          24.47n ±  3%   24.31n ±  6%        ~ (p=0.832 n=514+502)
[...]
SentrySystrap/Postgres-48             71.00n ±  8%   63.00n ±  6%  -11.27% (p=0.002 n=183+181)
```

PiperOrigin-RevId: 570931073
---
 pkg/bpf/interpreter.go      |  6 +++---
 pkg/bpf/interpreter_test.go | 13 +++++++++++++
 pkg/seccomp/seccomp.go      | 10 +++++++++-
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/pkg/bpf/interpreter.go b/pkg/bpf/interpreter.go
index de06bc694..cde771154 100644
--- a/pkg/bpf/interpreter.go
+++ b/pkg/bpf/interpreter.go
@@ -97,8 +97,8 @@ func (p Program) Length() int {
 	return len(p.instructions)
 }
 
-// Compile performs validation on a sequence of BPF instructions before
-// wrapping them in a Program.
+// Compile performs validation and optimization on a sequence of BPF
+// instructions before wrapping them in a Program.
 func Compile(insns []Instruction) (Program, error) {
 	if len(insns) == 0 || len(insns) > MaxInstructions {
 		return Program{}, Error{InvalidInstructionCount, len(insns)}
@@ -214,7 +214,7 @@ func Compile(insns []Instruction) (Program, error) {
 		}
 	}
 
-	return Program{insns}, nil
+	return Program{Optimize(insns)}, nil
 }
 
 // Input represents a source of input data for a BPF program. (BPF
diff --git a/pkg/bpf/interpreter_test.go b/pkg/bpf/interpreter_test.go
index 8052d2f46..cc5086e8f 100644
--- a/pkg/bpf/interpreter_test.go
+++ b/pkg/bpf/interpreter_test.go
@@ -705,6 +705,19 @@ func TestValidInstructions(t *testing.T) {
 			},
 			expectedRet: 2,
 		},
+		{
+			desc: "Optimizable program",
+			insns: []Instruction{
+				Stmt(Ld|Imm|W, 42),        // A = 42
+				Jump(Jmp|Jeq|K, 42, 0, 1), // if (A == 42) jmp 0 else 1
+				Jump(Jmp|Ja, 1, 0, 0),     // jmp 1
+				Jump(Jmp|Ja, 2, 0, 0),     // jmp 2
+				Stmt(Ld|Imm|W, 37),        // A = 37
+				Stmt(Ret|K, 0),            // return 0
+				Stmt(Ret|K, 1),            // return 1
+			},
+			expectedRet: 0,
+		},
 	} {
 		p, err := Compile(test.insns)
 		if err != nil {
diff --git a/pkg/seccomp/seccomp.go b/pkg/seccomp/seccomp.go
index 9f476f3ab..c0879e6f9 100644
--- a/pkg/seccomp/seccomp.go
+++ b/pkg/seccomp/seccomp.go
@@ -152,7 +152,15 @@ func BuildProgram(rules []RuleSet, defaultAction, badArchAction linux.BPFAction)
 	}
 	program.AddStmt(bpf.Ret|bpf.K, uint32(defaultAction))
 
-	return program.Instructions()
+	insns, err := program.Instructions()
+	if err != nil {
+		return insns, err
+	}
+	beforeOpt := len(insns)
+	insns = bpf.Optimize(insns)
+	afterOpt := len(insns)
+	log.Debugf("Seccomp program optimized from %d to %d instructions", beforeOpt, afterOpt)
+	return insns, nil
 }
 
 // buildIndex builds a BST to quickly search through all syscalls.