hook: Simplify Hook.Patch()

cilium/ebpf now has a CollectionSpec.RewriteMaps() API that we can use
to rewrite maps instead of patching instructions ourselves.

This originally attempted to workaround an upstream issue with
cilium/ebpf that meant Hook.Patch() didn't work well with
CollectionSpec.LoadAndAssign().
cilium/ebpf@04b5c2a
fixes this, but keep the extra test to ensure it works.

For posterity's sake, the original description of the issue:

The "classic" cilium/ebpf way of doing things is to get a
CollectionSpec, and use NewCollection() to create all the
maps and load the programs into the kernel.
The user can then retrieve the maps and programs of interest from
the Collection.

xdpcap.Patch() works well with this API:

- if it's called, the program instructions are rewritten to use the
  hook map, and the map is removed from the spec to prevent
  NewCollection() from creating another one.
- if it isn't called, NewCollection() creates an empy hook map.

cilium/ebpf also supports a new API, CollectionSpec.LoadAndAssign() that
allows programs and maps to be loaded into the kernel and assigned to a
struct in one call. To avoid GC related issues I don't full understand,
it requires that ProgramArray maps are either:

- rewritten before LoadAndAssign(), so the caller already has a
  reference to them.
- stored in the LoadAndAssign() struct if they're created by
  LoadAndAssign(), so the caller also has a reference to them.

If xdpcap.Patch() isn't called, this breaks as users typically wouldn't
store the hook map in the LoadAndAssign() struct as they have no use for
it.

hook.go

@@ -7,7 +7,6 @@ import (
 	"github.com/cloudflare/xdpcap/internal"
 
 	"github.com/cilium/ebpf"
-	"github.com/cilium/ebpf/asm"
 	"github.com/pkg/errors"
 )
 
@@ -61,26 +60,7 @@ func (h *Hook) Patch(spec *ebpf.CollectionSpec, hookMapSymbol string) error {
 		return nil
 	}
 
-	if spec.Maps[hookMapSymbol] == nil {
-		return errors.Errorf("missing map %s", hookMapSymbol)
-	}
-
-	// We can't specify to use an already existing map in a spec, so:
-	// - Rewrite the hook map symbol of every program
-	// - Remove the map from spec (so it isn't created later on)
-	for progName, progSpec := range spec.Programs {
-		err := progSpec.Instructions.RewriteMapPtr(hookMapSymbol, h.hookMap.FD())
-		// Not all programs need to use the hook
-		if asm.IsUnreferencedSymbol(err) {
-			continue
-		}
-
-		if err != nil {
-			return errors.Wrapf(err, "program %s", progName)
-		}
-	}
-
-	delete(spec.Maps, hookMapSymbol)
-
-	return nil
+	return spec.RewriteMaps(map[string]*ebpf.Map{
+		hookMapSymbol: h.hookMap,
+	})
 }

hook_test.go

@@ -11,20 +11,59 @@ const (
 	hookSymbol = "xdpcap_hook"
 )
 
-// Test loading an elf with a hook map
-func TestHook(t *testing.T) {
+// Test loading an elf with a hook map using ebpf.NewCollection().
+func TestHookNewCollection(t *testing.T) {
 	hook, err := NewHook("foo")
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer hook.Close()
 
-	spec, err := ebpf.LoadCollectionSpec(elf)
+	spec := mustPatchSpec(t, hook)
+
+	coll, err := ebpf.NewCollection(spec)
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer coll.Close()
+}
 
-	err = hook.Patch(spec, hookSymbol)
+// Test loading an elf with a hook map using ebpf.CollectionSpec.LoadAndReplace(),
+// which didn't always work: https://github.com/cilium/ebpf/commit/04b5c2a901f3bcfa7d7a13c59f7c1c556f2f3d5f
+func TestHookLoadAndReplace(t *testing.T) {
+	test := func(t *testing.T, hook *Hook) {
+		spec := mustPatchSpec(t, hook)
+
+		var objs struct {
+			// Works for both programs that do and don't use the hook.
+			Hook   *ebpf.Program `ebpf:"xdp_hook"`
+			NoHook *ebpf.Program `ebpf:"xdp_nohook"`
+		}
+		if err := spec.LoadAndAssign(&objs, nil); err != nil {
+			t.Fatal(err)
+		}
+		defer objs.Hook.Close()
+		defer objs.NoHook.Close()
+	}
+
+	t.Run("nil", func(t *testing.T) {
+		test(t, nil)
+	})
+
+	t.Run("not-nil", func(t *testing.T) {
+		hook, err := NewHook("foo")
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer hook.Close()
+
+		test(t, hook)
+	})
+}
+
+// Test loading an elf that uses a hook without an explicit hook map
+func TestNoHook(t *testing.T) {
+	spec, err := ebpf.LoadCollectionSpec(elf)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -36,16 +75,16 @@ func TestHook(t *testing.T) {
 	defer coll.Close()
 }
 
-// Test loading an elf that uses a hook without an explicit hook map
-func TestNoHook(t *testing.T) {
+func mustPatchSpec(tb testing.TB, hook *Hook) *ebpf.CollectionSpec {
 	spec, err := ebpf.LoadCollectionSpec(elf)
 	if err != nil {
-		t.Fatal(err)
+		tb.Fatal(err)
 	}
 
-	coll, err := ebpf.NewCollection(spec)
+	err = hook.Patch(spec, hookSymbol)
 	if err != nil {
-		t.Fatal(err)
+		tb.Fatal(err)
 	}
-	defer coll.Close()
+
+	return spec
 }

testdata/xdp_hook.c

@@ -25,13 +25,13 @@ struct bpf_map_def xdpcap_hook = XDPCAP_HOOK();
 /**
  * Example / test XDP program that exposes packets through an xdpcap hook.
  */
-__section("test_hook") int xdp_hook(struct xdp_md *ctx) {
+__section("xdp/hook") int xdp_hook(struct xdp_md *ctx) {
 	return xdpcap_exit(ctx, &xdpcap_hook, XDP_PASS);
 }
 
 /**
  * Second entrypoint that doesn't use xdpcap exit() to test patching.
  */
-__section("test_nohook") int xdp_nohook(struct xdp_md *ctx) {
+__section("xdp/nohook") int xdp_nohook(struct xdp_md *ctx) {
     return XDP_PASS;
 }