From a27a5bc9fd7e59679c98d9140ccfdb89420434f1 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Fri, 27 Oct 2023 12:00:38 -0700
Subject: [PATCH] `bpf`: Add logic for verifying whether a program fragment
 modifies register A.

This is useful for value matching rules which look for the value of the `A`
register. If they do not modify this value, then we do not need to reload it
between sequential matchers over the same data.

PiperOrigin-RevId: 577266390
---
 pkg/bpf/bpf.go                  | 15 ++++++++++++++
 pkg/bpf/program_builder.go      | 13 ++++++++++++
 pkg/bpf/program_builder_test.go | 36 +++++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+)

diff --git a/pkg/bpf/bpf.go b/pkg/bpf/bpf.go
index 107f4935f..2abeef584 100644
--- a/pkg/bpf/bpf.go
+++ b/pkg/bpf/bpf.go
@@ -187,3 +187,18 @@ func (ins Instruction) JumpOffsets() []JumpOffset {
 	}
 	return []JumpOffset{{JumpDirect, ins.K}}
 }
+
+// ModifiesRegisterA returns true iff this instruction modifies the value
+// of the "A" register.
+func (ins Instruction) ModifiesRegisterA() bool {
+	switch ins.OpCode & instructionClassMask {
+	case Ld:
+		return true
+	case Alu:
+		return true
+	case Misc:
+		return ins.OpCode == Misc|Tax
+	default:
+		return false
+	}
+}
diff --git a/pkg/bpf/program_builder.go b/pkg/bpf/program_builder.go
index 584cbef90..644017753 100644
--- a/pkg/bpf/program_builder.go
+++ b/pkg/bpf/program_builder.go
@@ -352,3 +352,16 @@ func (f ProgramFragment) Outcomes() FragmentOutcomes {
 	}
 	return outcomes
 }
+
+// MayModifyRegisterA returns whether this fragment may modify register A.
+// A value of "true" does not necessarily mean that A *will* be modified,
+// as the control flow of this fragment may skip over instructions that
+// modify the A register.
+func (f ProgramFragment) MayModifyRegisterA() bool {
+	for pc := f.fromPC; pc < f.toPC; pc++ {
+		if f.b.instructions[pc].ModifiesRegisterA() {
+			return true
+		}
+	}
+	return false
+}
diff --git a/pkg/bpf/program_builder_test.go b/pkg/bpf/program_builder_test.go
index 77583834b..3acb13cbe 100644
--- a/pkg/bpf/program_builder_test.go
+++ b/pkg/bpf/program_builder_test.go
@@ -439,3 +439,39 @@ func TestProgramBuilderOutcomes(t *testing.T) {
 		})
 	}
 }
+
+func TestProgramBuilderMayModifyRegisterA(t *testing.T) {
+	t.Run("empty program", func(t *testing.T) {
+		if got := NewProgramBuilder().Record()().MayModifyRegisterA(); got != false {
+			t.Errorf("MayModifyRegisterA: got %v want %v", got, false)
+		}
+	})
+	t.Run("does not modify register A", func(t *testing.T) {
+		b := NewProgramBuilder()
+		stop := b.Record()
+		b.AddJump(Jmp|Ja, 0, 0, 0)
+		b.AddJump(Jmp|Jeq|K, 0, 0, 0)
+		b.AddStmt(Misc|Txa, 0)
+		b.AddStmt(Ret|K, 1337)
+		if got := stop().MayModifyRegisterA(); got != false {
+			t.Errorf("MayModifyRegisterA: got %v want %v", got, false)
+		}
+	})
+	for _, ins := range []Instruction{
+		Stmt(Ld|Abs|W, 0),
+		Stmt(Alu|Neg, 0),
+		Stmt(Misc|Tax, 0),
+	} {
+		t.Run(fmt.Sprintf("modifies register A via %v", ins), func(t *testing.T) {
+			b := NewProgramBuilder()
+			stop := b.Record()
+			b.AddJump(Jmp|Ja, 0, 0, 0)
+			b.AddJump(Jmp|Jeq|K, 0, 0, 0)
+			b.AddStmt(ins.OpCode, ins.K)
+			b.AddStmt(Ret|K, 1337)
+			if got := stop().MayModifyRegisterA(); got != true {
+				t.Errorf("MayModifyRegisterA: got %v want %v", got, true)
+			}
+		})
+	}
+}