btf: add FindHandle

Add a helper that iterates all BTF objects in the kernel and returns
the first one for which a user supplied callback returns true.

btf/handle.go

@@ -130,3 +130,30 @@ func (it *HandleIterator) Take() *Handle {
 func (it *HandleIterator) Err() error {
 	return it.err
 }
+
+// FindHandle returns the first handle for which predicate returns true.
+//
+// Requires CAP_SYS_ADMIN.
+//
+// Returns an error wrapping ErrNotFound if predicate never returns true or if
+// there is no BTF loaded into the kernel.
+func FindHandle(predicate func(info *HandleInfo) bool) (*Handle, error) {
+	it := new(HandleIterator)
+	defer it.Handle.Close()
+
+	for it.Next() {
+		info, err := it.Handle.Info()
+		if err != nil {
+			return nil, fmt.Errorf("info for ID %d: %w", it.ID, err)
+		}
+
+		if predicate(info) {
+			return it.Take(), nil
+		}
+	}
+	if err := it.Err(); err != nil {
+		return nil, fmt.Errorf("iterate handles: %w", err)
+	}
+
+	return nil, fmt.Errorf("find handle: %w", ErrNotFound)
+}

btf/handle_test.go

@@ -50,60 +50,25 @@ func TestParseModuleSplitSpec(t *testing.T) {
 	// See TestNewHandleFromID for reasoning.
 	testutils.SkipOnOldKernel(t, "5.11", "vmlinux BTF ID")
 
-	var module *btf.Handle
-	defer module.Close()
-
-	it := new(btf.HandleIterator)
-	defer it.Handle.Close()
-
-	for it.Next() {
-		info, err := it.Handle.Info()
-		if err != nil {
-			t.Fatal(err)
+	module, err := btf.FindHandle(func(info *btf.HandleInfo) bool {
+		if info.IsModule() {
+			t.Log("Using module", info.Name)
+			return true
 		}
-
-		if !info.IsModule() {
-			continue
-		}
-
-		t.Log("Using module", info.Name)
-		module = it.Take()
-		break
-	}
-	if err := it.Err(); err != nil {
+		return false
+	})
+	if err != nil {
 		t.Fatal(err)
 	}
+	defer module.Close()
 
-	if module == nil {
-		t.Fatal("No BTF for kernel module found")
-	}
-
-	var vmlinux *btf.Handle
-	defer vmlinux.Close()
-
-	it = new(btf.HandleIterator)
-	defer it.Handle.Close()
-
-	for it.Next() {
-		info, err := it.Handle.Info()
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		if !info.IsVmlinux() {
-			continue
-		}
-
-		vmlinux = it.Take()
-		break
-	}
-	if err := it.Err(); err != nil {
+	vmlinux, err := btf.FindHandle(func(info *btf.HandleInfo) bool {
+		return info.IsVmlinux()
+	})
+	if err != nil {
 		t.Fatal(err)
 	}
-
-	if vmlinux == nil {
-		t.Fatal("No BTF for kernel found")
-	}
+	defer vmlinux.Close()
 
 	vmlinuxSpec, err := vmlinux.Spec(nil)
 	if err != nil {