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

@@ -119,3 +119,33 @@ func (it *HandleIterator) Next(handle **Handle) bool {
 func (it *HandleIterator) Err() error {
 	return it.err
 }
+
+// FindHandle returns the first handle for which predicate returns true.
+//
+// Requires CAP_SYS_ADMIN.
+//
+// Returns 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) {
+	var handle *Handle
+	defer handle.Close()
+
+	it := new(HandleIterator)
+	for it.Next(&handle) {
+		info, err := handle.Info()
+		if err != nil {
+			return nil, fmt.Errorf("info for ID %d: %w", it.ID, err)
+		}
+
+		if predicate(info) {
+			tmp := handle
+			handle = nil
+			return tmp, 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

@@ -51,54 +51,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)
-	for it.Next(&module) {
-		info, err := module.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)
-		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)
-	for it.Next(&vmlinux) {
-		info, err := vmlinux.Info()
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		if !info.IsVmlinux() {
-			continue
-		}
-
-		break
-	}
-	if err := it.Err(); err != nil {
+	vmlinux, err := btf.FindHandle(func(info *btf.HandleInfo) bool {
+		return info.IsKernel
+	})
+	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 {