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

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

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

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