diff --git a/btf/handle.go b/btf/handle.go index 128e9b35c..7128171c8 100644 --- a/btf/handle.go +++ b/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) +} diff --git a/btf/handle_test.go b/btf/handle_test.go index fd6db85b8..2ab21e022 100644 --- a/btf/handle_test.go +++ b/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 {