program: support tracing of kernel modules

btf/btf.go

@@ -609,6 +609,9 @@ func (s *Spec) AnyTypeByName(name string) (Type, error) {
 // Type exists in the Spec. If multiple candidates are found,
 // an error is returned.
 func (s *Spec) TypeByName(name string, typ interface{}) error {
+	typeInterface := reflect.TypeOf((*Type)(nil)).Elem()
+
+	// typ may be **T or *Type
 	typValue := reflect.ValueOf(typ)
 	if typValue.Kind() != reflect.Ptr {
 		return fmt.Errorf("%T is not a pointer", typ)
@@ -620,7 +623,12 @@ func (s *Spec) TypeByName(name string, typ interface{}) error {
 	}
 
 	wanted := typPtr.Type()
-	if !wanted.AssignableTo(reflect.TypeOf((*Type)(nil)).Elem()) {
+	if wanted == typeInterface {
+		// This is *Type. Unwrap the value's type.
+		wanted = typPtr.Elem().Type()
+	}
+
+	if !wanted.AssignableTo(typeInterface) {
 		return fmt.Errorf("%T does not satisfy Type interface", typ)
 	}
 
@@ -643,7 +651,7 @@ func (s *Spec) TypeByName(name string, typ interface{}) error {
 	}
 
 	if candidate == nil {
-		return fmt.Errorf("type %s: %w", name, ErrNotFound)
+		return fmt.Errorf("%s %s: %w", wanted, name, ErrNotFound)
 	}
 
 	typPtr.Set(reflect.ValueOf(candidate))

btf/btf_test.go

@@ -130,6 +130,12 @@ func TestTypeByName(t *testing.T) {
 		t.Fatal("multiple TypeByName calls for `iphdr` name do not return the same addresses")
 	}
 
+	// It's valid to pass a *Type to TypeByName.
+	typ := Type(iphdr2)
+	if err := spec.TypeByName("iphdr", &typ); err != nil {
+		t.Fatal("Can't look up using *Type:", err)
+	}
+
 	// Excerpt from linux/ip.h, https://elixir.bootlin.com/linux/latest/A/ident/iphdr
 	//
 	// struct iphdr {

btf/handle.go

@@ -71,51 +71,89 @@ func (i *HandleInfo) IsModule() bool {
 
 // HandleIterator allows enumerating BTF blobs loaded into the kernel.
 type HandleIterator struct {
-	// The ID of the last retrieved handle. Only valid after a call to Next.
-	ID  ID
-	err error
+	// The ID of the current handle. Only valid after a call to Next.
+	ID ID
+	// The current Handle. Only valid until a call to Next.
+	// See Take if you want to retain the handle.
+	Handle *Handle
+	err    error
 }
 
-// Next retrieves a handle for the next BTF blob.
-//
-// [Handle.Close] is called if *handle is non-nil to avoid leaking fds.
+// Next retrieves a handle for the next BTF object.
 //
-// Returns true if another BTF blob was found. Call [HandleIterator.Err] after
+// Returns true if another BTF object was found. Call [HandleIterator.Err] after
 // the function returns false.
-func (it *HandleIterator) Next(handle **Handle) bool {
-	if *handle != nil {
-		(*handle).Close()
-		*handle = nil
-	}
-
+func (it *HandleIterator) Next() bool {
 	id := it.ID
 	for {
 		attr := &sys.BtfGetNextIdAttr{Id: id}
 		err := sys.BtfGetNextId(attr)
 		if errors.Is(err, os.ErrNotExist) {
 			// There are no more BTF objects.
-			return false
+			break
 		} else if err != nil {
 			it.err = fmt.Errorf("get next BTF ID: %w", err)
-			return false
+			break
 		}
 
 		id = attr.NextId
-		*handle, err = NewHandleFromID(id)
+		handle, err := NewHandleFromID(id)
 		if errors.Is(err, os.ErrNotExist) {
 			// Try again with the next ID.
 			continue
 		} else if err != nil {
 			it.err = fmt.Errorf("retrieve handle for ID %d: %w", id, err)
-			return false
+			break
 		}
 
-		it.ID = id
+		it.Handle.Close()
+		it.ID, it.Handle = id, handle
 		return true
 	}
+
+	// No more handles or we encountered an error.
+	it.Handle.Close()
+	it.Handle = nil
+	return false
+}
+
+// Take the ownership of the current handle.
+//
+// It's the callers responsibility to close the handle.
+func (it *HandleIterator) Take() *Handle {
+	handle := it.Handle
+	it.Handle = nil
+	return handle
 }
 
 // Err returns an error if iteration failed for some reason.
 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

@@ -14,20 +14,19 @@ func TestHandleIterator(t *testing.T) {
 	// See https://github.com/torvalds/linux/commit/5329722057d41aebc31e391907a501feaa42f7d9
 	testutils.SkipOnOldKernel(t, "5.11", "vmlinux BTF ID")
 
-	var h *btf.Handle
-	defer h.Close()
-
 	it := new(btf.HandleIterator)
-	if !it.Next(&h) {
+	defer it.Handle.Close()
+
+	if !it.Next() {
 		t.Fatalf("No BTF loaded")
 	}
-	if h == nil {
+	if it.Handle == nil {
 		t.Fatal("Next doesn't assign handle")
 	}
 	prev := it.ID
-	for it.Next(&h) {
+	for it.Next() {
 		// Iterate all loaded BTF.
-		if h == nil {
+		if it.Handle == nil {
 			t.Fatal("Next doesn't assign handle")
 		}
 		if it.ID == prev {
@@ -39,7 +38,7 @@ func TestHandleIterator(t *testing.T) {
 		t.Fatal("Iteration returned an error:", err)
 	}
 
-	if h != nil {
+	if it.Handle != nil {
 		t.Fatal("Next doesn't clean up handle on last iteration")
 	}
 	if prev != it.ID {
@@ -51,54 +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)
-	for it.Next(&module) {
-		info, err := module.Info()
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		if !info.IsModule() {
-			continue
+	module, err := btf.FindHandle(func(info *btf.HandleInfo) bool {
+		if info.IsModule() {
+			t.Log("Using module", info.Name)
+			return true
 		}
-
-		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.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 {
@@ -117,13 +87,16 @@ func TestParseModuleSplitSpec(t *testing.T) {
 }
 
 func ExampleHandleIterator() {
-	var handle *btf.Handle
-	// Ensure that handle is cleaned up. This is valid for nil handles as well.
-	defer handle.Close()
-
 	it := new(btf.HandleIterator)
-	for it.Next(&handle) {
-		fmt.Printf("Found handle with ID %d\n", it.ID)
+	defer it.Handle.Close()
+
+	for it.Next() {
+		info, err := it.Handle.Info()
+		if err != nil {
+			panic(err)
+		}
+
+		fmt.Printf("Found handle with ID %d and name %s\n", it.ID, info.Name)
 	}
 	if err := it.Err(); err != nil {
 		panic(err)

internal/cmd/gentypes/main.go

@@ -285,6 +285,7 @@ import (
 					"fd_array",
 					"core_relos",
 				),
+				choose(20, "attach_btf_obj_fd"),
 			},
 		},
 		{