pin: make WalkDir harder to misuse and add Windows support

The current WalkDir API requires the user to close each object that
is iterated over. This isn't very intuitive, the two uses of the
function in the tetragon code base end up leaking objects.
It's also not straight forward to implement WalkDir on Windows, since
BPFFS isn't really a thing on that platform.

Instead, turn WalkDir into an iterator which returns a new Pin object.
By default, objects are closed once they have been iterated over. The
user may call Pin.Take if they wish to retain the object. This is
similar to the existing link.Iterator.Take.

One downside is that directories are not returned by the iteration
anymore. The existing code in tetragon doesn't use them and is
therefore not affected. There is also no more fine grained control
over whether to descent into a directory. Aborting an iteration is
still supported and as efficient as before.

Fixes: #1652
Signed-off-by: Lorenz Bauer <[email protected]>

Author: lmb

.github/workflows/ci.yml

@@ -346,23 +346,7 @@ jobs:
       - name: Test
         run: >
           gotestsum --raw-command --ignore-non-json-output-lines --junitfile junit.xml --
-          go test -short -count 1 -json
-          ./
-          ./asm
-          ./btf
-          ./features
-          ./internal
-          ./internal/efw
-          ./internal/kallsyms
-          ./internal/kconfig
-          ./internal/linux
-          ./internal/sys
-          ./internal/sysenc
-          ./internal/testutils
-          ./internal/testutils/testmain
-          ./internal/tracefs
-          ./internal/unix
-          ./link
+          go test -short -count 1 -json ./...
 
       - name: Upload Test Results
         if: always()

pin/load_test.go

@@ -8,16 +8,22 @@ import (
 
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/asm"
+	"github.com/cilium/ebpf/internal/platform"
 	"github.com/cilium/ebpf/internal/testutils"
 	"github.com/cilium/ebpf/internal/testutils/testmain"
 )
 
 func mustPinnedProgram(t *testing.T, path string) *ebpf.Program {
 	t.Helper()
 
+	typ := ebpf.SocketFilter
+	if platform.IsWindows {
+		typ = ebpf.WindowsXDPTest
+	}
+
 	spec := &ebpf.ProgramSpec{
 		Name: "test",
-		Type: ebpf.SocketFilter,
+		Type: typ,
 		Instructions: asm.Instructions{
 			asm.LoadImm(asm.R0, 2, asm.DWord),
 			asm.Return(),
@@ -41,9 +47,14 @@ func mustPinnedProgram(t *testing.T, path string) *ebpf.Program {
 func mustPinnedMap(t *testing.T, path string) *ebpf.Map {
 	t.Helper()
 
+	typ := ebpf.Array
+	if platform.IsWindows {
+		typ = ebpf.WindowsArray
+	}
+
 	spec := &ebpf.MapSpec{
 		Name:       "test",
-		Type:       ebpf.Array,
+		Type:       typ,
 		KeySize:    4,
 		ValueSize:  4,
 		MaxEntries: 1,

pin/pin.go

@@ -0,0 +1,22 @@
+package pin
+
+import "io"
+
+// Pin represents an object and its pinned path.
+type Pin struct {
+	Path   string
+	Object io.Closer
+}
+
+func (p *Pin) close() {
+	if p.Object != nil {
+		p.Object.Close()
+	}
+}
+
+// Take ownership of Pin.Object.
+//
+// The caller is responsible for calling close on [Pin.Object].
+func (p *Pin) Take() {
+	p.Object = nil
+}

pin/walk.go

@@ -0,0 +1,64 @@
+//go:build !windows
+
+package pin
+
+import (
+	"fmt"
+	"io/fs"
+	"iter"
+	"os"
+	"path/filepath"
+
+	"github.com/cilium/ebpf"
+	"github.com/cilium/ebpf/internal/linux"
+	"github.com/cilium/ebpf/internal/unix"
+)
+
+// WalkDir walks the file tree rooted at path and yields a [Pin] for each
+// BPF object below the path.
+//
+// Callers must invoke [Pin.Take] if they wish to hold on to the object.
+func WalkDir(root string, opts *ebpf.LoadPinOptions) iter.Seq2[*Pin, error] {
+	return func(yield func(*Pin, error) bool) {
+		fsType, err := linux.FSType(root)
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+		if fsType != unix.BPF_FS_MAGIC {
+			yield(nil, fmt.Errorf("%s is not on a bpf filesystem", root))
+			return
+		}
+
+		fn := func(path string, d fs.DirEntry, err error) error {
+			if err != nil {
+				return err
+			}
+
+			if d.IsDir() {
+				return nil
+			}
+
+			path = filepath.Join(root, path)
+			obj, err := Load(path, opts)
+			if err != nil {
+				return err
+			}
+
+			pin := &Pin{Path: path, Object: obj}
+			defer pin.close()
+
+			if !yield(pin, nil) {
+				return fs.SkipAll
+			}
+
+			return nil
+		}
+
+		err = fs.WalkDir(os.DirFS(root), ".", fn)
+		if err != nil {
+			yield(nil, fmt.Errorf("walk: %w", err))
+			return
+		}
+	}
+}

pin/walk_other.go

@@ -1,58 +1,67 @@
 package pin
 
 import (
-	"io/fs"
+	"iter"
 	"os"
 	"path/filepath"
+	"reflect"
 	"testing"
 
 	"github.com/go-quicktest/qt"
 
 	"github.com/cilium/ebpf"
+	"github.com/cilium/ebpf/internal/platform"
 	"github.com/cilium/ebpf/internal/testutils"
 )
 
 func TestWalkDir(t *testing.T) {
 	testutils.SkipOnOldKernel(t, "4.10", "reading program fdinfo")
 
 	tmp := testutils.TempBPFFS(t)
+
 	dir := filepath.Join(tmp, "dir")
-	qt.Assert(t, qt.IsNil(os.Mkdir(dir, 0755)))
+	if !platform.IsWindows {
+		// Windows doesn't have a BPF file system, so mkdir below fails.
+		qt.Assert(t, qt.IsNil(os.Mkdir(dir, 0755)))
+	}
 
-	mustPinnedProgram(t, filepath.Join(tmp, "pinned_prog"))
-	mustPinnedMap(t, filepath.Join(dir, "pinned_map"))
+	progPath := filepath.Join(tmp, "pinned_prog")
+	mustPinnedProgram(t, progPath)
+	mapPath := filepath.Join(dir, "pinned_map")
+	mustPinnedMap(t, mapPath)
 
-	entries := make(map[string]string)
+	next, stop := iter.Pull2(WalkDir(tmp, nil))
+	defer stop()
 
-	bpffn := func(path string, d fs.DirEntry, obj Pinner, err error) error {
-		qt.Assert(t, qt.IsNil(err))
+	pin, err, ok := next()
+	qt.Assert(t, qt.IsTrue(ok))
+	qt.Assert(t, qt.IsNil(err))
+	qt.Assert(t, qt.Equals(reflect.TypeOf(pin.Object), reflect.TypeFor[*ebpf.Map]()))
+	qt.Assert(t, qt.Equals(pin.Path, mapPath))
 
-		if obj != nil {
-			defer obj.Close()
-		}
+	pin, err, ok = next()
+	qt.Assert(t, qt.IsTrue(ok))
+	qt.Assert(t, qt.IsNil(err))
+	qt.Assert(t, qt.Equals(reflect.TypeOf(pin.Object), reflect.TypeFor[*ebpf.Program]()))
+	qt.Assert(t, qt.Equals(pin.Path, progPath))
 
-		if path == "." {
-			return nil
-		}
+	_, _, ok = next()
+	qt.Assert(t, qt.IsFalse(ok))
 
-		switch obj.(type) {
-		case *ebpf.Program:
-			entries[path] = "prog"
-		case *ebpf.Map:
-			entries[path] = "map"
-		default:
-			entries[path] = ""
+	t.Run("Not BPFFS", func(t *testing.T) {
+		if platform.IsWindows {
+			t.Skip("Windows does not have BPFFS")
 		}
 
-		return nil
-	}
-	qt.Assert(t, qt.IsNil(WalkDir(tmp, bpffn)))
+		next, stop := iter.Pull2(WalkDir("/", nil))
+		defer stop()
+
+		_, err, ok = next()
+		qt.Assert(t, qt.IsTrue(ok))
+		qt.Assert(t, qt.IsNotNil(err))
 
-	qt.Assert(t, qt.DeepEquals(entries, map[string]string{
-		"pinned_prog":    "prog",
-		"dir":            "",
-		"dir/pinned_map": "map",
-	}))
+		_, _, ok = next()
+		qt.Assert(t, qt.IsFalse(ok))
+	})
 
-	qt.Assert(t, qt.IsNotNil(WalkDir("/", nil)))
 }

pin/walk_test.go

@@ -0,0 +1,48 @@
+package pin
+
+import (
+	"errors"
+	"iter"
+	"strings"
+
+	"github.com/cilium/ebpf"
+	"github.com/cilium/ebpf/internal/efw"
+)
+
+// WalkDir walks the file tree rooted at path and yields a [Pin] for each
+// BPF object below the path.
+//
+// Callers must invoke [Pin.Take] if they wish to hold on to the object.
+func WalkDir(root string, opts *ebpf.LoadPinOptions) iter.Seq2[*Pin, error] {
+	return func(yield func(*Pin, error) bool) {
+		cursor := root
+		for {
+			next, _, err := efw.EbpfGetNextPinnedObjectPath(cursor, efw.EBPF_OBJECT_UNKNOWN)
+			if errors.Is(err, efw.EBPF_NO_MORE_KEYS) {
+				break
+			} else if err != nil {
+				yield(nil, err)
+				return
+			}
+
+			if !strings.HasPrefix(next, root) {
+				break
+			}
+
+			obj, err := Load(next, opts)
+			if err != nil {
+				yield(nil, err)
+				return
+			}
+
+			pin := &Pin{next, obj}
+			ok := yield(pin, nil)
+			pin.close()
+			if !ok {
+				return
+			}
+
+			cursor = next
+		}
+	}
+}