Support multiple instruction patchers (#168)

* Allow multiple instruction patchers

* Add tests for instruction patchers

examples/instruction_patching/main.go

@@ -25,7 +25,7 @@ var m1 = &manager.Manager{
 			},
 		},
 	},
-	InstructionPatcher: patchBPFTelemetry,
+	InstructionPatchers: []manager.InstructionPatcherFunc{patchBPFTelemetry},
 }
 
 const BPFTelemetryPatchCall = -1

manager.go

@@ -124,6 +124,9 @@ type Options struct {
 	SkipRingbufferReaderStartup map[string]bool
 }
 
+// InstructionPatcherFunc - A function that patches the instructions of a program
+type InstructionPatcherFunc func(m *Manager) error
+
 // Manager - Helper structure that manages multiple eBPF programs and maps
 type Manager struct {
 	collectionSpec     *ebpf.CollectionSpec
@@ -150,9 +153,9 @@ type Manager struct {
 	// and dump the current state (human-readable)
 	DumpHandler func(w io.Writer, manager *Manager, mapName string, currentMap *ebpf.Map)
 
-	// InstructionPatcher - Callback function called before loading probes, to
+	// InstructionPatchers - Callback functions called before loading probes, to
 	// provide user the ability to perform last minute instruction patching.
-	InstructionPatcher func(m *Manager) error
+	InstructionPatchers []InstructionPatcherFunc
 }
 
 // DumpMaps - Write in the w writer argument human-readable info about eBPF maps
@@ -558,8 +561,8 @@ func (m *Manager) InitWithOptions(elf io.ReaderAt, options Options) error {
 	}
 
 	// Patch instructions
-	if m.InstructionPatcher != nil {
-		if err := m.InstructionPatcher(m); err != nil {
+	for _, patcher := range m.InstructionPatchers {
+		if err := patcher(m); err != nil {
 			resetManager(m)
 			return err
 		}

manager_test.go

@@ -201,3 +201,63 @@ func TestDumpMaps(t *testing.T) {
 		t.Errorf("expected %s, got %s", dumpContents, output.String())
 	}
 }
+
+func TestInstructionPatching(t *testing.T) {
+	err := rlimit.RemoveMemlock()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// We want to test multiple patchers, so we'll use a generic one
+	// and call it twice with different constants.
+	// The patching.c program contains two invalid calls, with constants
+	// -1 and -2. We just replace them with a movimm instruction.
+	genericPatcher := func(m *Manager, constant int64) error {
+		specs, err := m.GetProgramSpecs()
+		if err != nil {
+			return err
+		}
+		for _, spec := range specs {
+			if spec == nil {
+				continue
+			}
+			iter := spec.Instructions.Iterate()
+			for iter.Next() {
+				ins := iter.Ins
+
+				if !ins.IsBuiltinCall() {
+					continue
+				}
+
+				if ins.Constant == constant {
+					*ins = asm.Mov.Imm(asm.R1, int32(0xff)).WithMetadata(ins.Metadata)
+				}
+			}
+		}
+		return nil
+	}
+
+	m := &Manager{
+		Probes: []*Probe{
+			{ProbeIdentificationPair: ProbeIdentificationPair{EBPFFuncName: "patching_test"}},
+		},
+		InstructionPatchers: []InstructionPatcherFunc{
+			func(m *Manager) error { return genericPatcher(m, -1) },
+			func(m *Manager) error { return genericPatcher(m, -2) },
+		},
+	}
+
+	f, err := os.Open("testdata/patching.elf")
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = f.Close() })
+
+	// If any of the patchers fail, they will leave an invalid call instruction
+	// in the program, which will cause the verifier to fail. This allows us
+	// to not do any extra validation.
+	err = m.InitWithOptions(f, Options{})
+	if err != nil {
+		t.Fatal(err)
+	}
+}

testdata/Makefile

@@ -1,7 +1,7 @@
 LLVM_PREFIX ?= /usr/bin
 CLANG ?= $(LLVM_PREFIX)/clang
 
-all: rewrite.elf exclude.elf
+all: rewrite.elf exclude.elf patching.elf
 
 clean:
 	-$(RM) *.elf

testdata/patching.c

@@ -0,0 +1,14 @@
+#include "common.h"
+
+char _license[] __section("license") = "MIT";
+
+static void *(*bpf_patch_1)(unsigned long, ...) = (void *)-1;
+static void *(*bpf_patch_2)(unsigned long, ...) = (void *)-2;
+
+__section("socket")
+int patching_test() {
+	int ret = 0;
+	bpf_patch_1(ret);
+	bpf_patch_2(ret);
+	return 1;
+}