struct_ops: remove ad-hoc helpers and use findTargetInKernel

Signed-off-by: shun159 <[email protected]>

Author: shun159

collection.go

@@ -710,10 +710,14 @@ func (cl *collectionLoader) populateDeferredMaps() error {
 				return fmt.Errorf("load vmlinux BTF: %w", err)
 			}
 
-			vType, _, _, err := findStructTypeByName(s, valueType.Name)
+			target := btf.Type((*btf.Struct)(nil))
+			s, module, err := findTargetInKernel(s, valueType.Name, &target)
 			if err != nil {
-				return fmt.Errorf("find value type: %w", err)
+				return fmt.Errorf("lookup value type %q: %w", valueType.Name, err)
 			}
+			defer module.Close()
+
+			vType, _ := btf.As[*btf.Struct](target)
 
 			kernVData, err := translateStructData(valueType, userData, vType)
 			if err != nil {
@@ -858,9 +862,17 @@ func (cl *collectionLoader) initKernStructOps() error {
 		}
 
 		kTypeName := strings.TrimPrefix(vType.Name, structOpsValuePrefix)
-		kType, _, _, err := findStructTypeByName(s, kTypeName)
+
+		target := btf.Type((*btf.Struct)(nil))
+		_, module, err := findTargetInKernel(s, kTypeName, &target)
 		if err != nil {
-			return fmt.Errorf("struct type: %s %w", kTypeName, err)
+			return fmt.Errorf("lookup kern type %q: %w", kTypeName, err)
+		}
+		defer module.Close()
+
+		kType, ok := btf.As[*btf.Struct](target)
+		if !ok {
+			return fmt.Errorf("conversion target %s -> Struct in initKernStructOps", kTypeName)
 		}
 
 		cl.setStructOpsProgAttachTo(kType)

collection_test.go

@@ -809,13 +809,15 @@ func TestStructOpsMapSpecSimpleLoadAndAssign(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	_, _, _, err = findStructTypeByName(s, "bpf_testmod_ops")
+	target := btf.Type((*btf.Struct)(nil))
+	_, module, err := findTargetInKernel(s, "bpf_struct_ops_bpf_testmod_ops", &target)
 	if errors.Is(err, btf.ErrNotFound) {
 		t.Skip("bpf_testmod_ops not loaded")
 	}
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer module.Close()
 
 	coll := mustNewCollection(t, spec, nil)
 	for name := range spec.Maps {

map.go

@@ -576,13 +576,19 @@ func (spec *MapSpec) createMap(inner *sys.FD) (_ *Map, err error) {
 
 			// struct_ops: resolve value type ("bpf_struct_ops_<name>") and
 			// record kernel-specific BTF IDs / FDs needed for map creation.
-			vType, s, modBtfObjId, err := findStructTypeByName(s, valueType.Name)
+			target := btf.Type((*btf.Struct)(nil))
+			s, module, err := findTargetInKernel(s, valueType.Name, &target)
 			if err != nil {
 				return nil, fmt.Errorf("lookup value type %q: %w", valueType.Name, err)
 			}
+			defer module.Close()
+
+			vType, ok := btf.As[*btf.Struct](target)
+			if !ok {
+				return nil, fmt.Errorf("conversion target %s -> Struct for map", valueType.Name)
+			}
 
 			btfValueTypeId, err := s.TypeID(vType)
-			fmt.Println(vType, btfValueTypeId)
 			if err != nil {
 				return nil, fmt.Errorf("lookup type_id: %w", err)
 			}
@@ -591,17 +597,11 @@ func (spec *MapSpec) createMap(inner *sys.FD) (_ *Map, err error) {
 			attr.BtfVmlinuxValueTypeId = btfValueTypeId
 			attr.BtfFd = uint32(h.FD())
 
-			if modBtfObjId != 0 {
+			if module != nil {
 				// BPF_F_VTYPE_BTF_OBJ_FD is required if the type comes from a module
 				attr.MapFlags |= sys.BPF_F_VTYPE_BTF_OBJ_FD
-				// resolve a module handler for the kernel model
-				modH, err := btf.NewHandleFromID(btf.ID(modBtfObjId))
-				if err != nil {
-					return nil, fmt.Errorf("open module BTF (id=%d): %w", modBtfObjId, err)
-				}
 				// set FD for the kernel module
-				attr.ValueTypeBtfObjFd = int32(modH.FD())
-				defer modH.Close()
+				attr.ValueTypeBtfObjFd = int32(module.FD())
 			}
 		} else {
 			handle, keyTypeID, valueTypeID, err := btf.MarshalMapKV(spec.Key, spec.Value)

prog.go

@@ -408,31 +408,33 @@ func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache)
 				return nil, fmt.Errorf("load vmlinux BTF: %w", err)
 			}
 
-			st, spec2, modBtfObjID, err := findStructTypeByName(s, attachTo)
+			target := btf.Type((*btf.Struct)(nil))
+			s, module, err := findTargetInKernel(s, attachTo, &target)
 			if err != nil {
-				return nil, err
+				return nil, fmt.Errorf("lookup struct_ops kern type %q: %w", attachTo, err)
+			}
+			defer module.Close()
+
+			kType, ok := btf.As[*btf.Struct](target)
+			if !ok {
+				return nil, fmt.Errorf("conversion target %s -> Struct for prog", attachTo)
 			}
 
-			tid, err := spec2.TypeID(st)
+			tid, err := s.TypeID(kType)
 			if err != nil {
-				return nil, fmt.Errorf("type id for %s: %w", st.TypeName(), err)
+				return nil, fmt.Errorf("type id for %s: %w", kType.TypeName(), err)
 			}
 			attr.AttachBtfId = sys.TypeID(tid)
 
-			idx := getStructMemberIndexByName(st, targetMember)
+			idx := getStructMemberIndexByName(kType, targetMember)
 			if idx < 0 {
-				return nil, fmt.Errorf("member %q not found in %s", targetMember, st.Name)
+				return nil, fmt.Errorf("member %q not found in %s", targetMember, kType.Name)
 			}
 
 			attr.ExpectedAttachType = sys.AttachType(idx)
 
-			if modBtfObjID != 0 {
-				h, err := btf.NewHandleFromID(btf.ID(modBtfObjID))
-				if err != nil {
-					return nil, fmt.Errorf("open module BTF handle (id=%d): %w", modBtfObjID, err)
-				}
-				defer h.Close()
-				attr.AttachBtfObjFd = uint32(h.FD())
+			if module != nil {
+				attr.AttachBtfObjFd = uint32(module.FD())
 			}
 		}
 

struct_ops.go

@@ -1,87 +1,9 @@
 package ebpf
 
-import (
-	"errors"
-	"fmt"
-
-	"github.com/cilium/ebpf/btf"
-)
+import "github.com/cilium/ebpf/btf"
 
 const structOpsValuePrefix = "bpf_struct_ops_"
 
-// findStructTypeByName looks up a struct type with the exact name in the
-// provided base BTF spec, and falls back to scanning all loaded module BTFs
-// if it is not present in vmlinux.
-//
-// Search order and behavior:
-//  1. vmlinux (base spec): try AnyTypeByName(name). If it exists and is a
-//     *btf.Struct, return it immediately with moduleID=0.
-//     - If AnyTypeByName returns a non-notfound error, the error is propagated.
-//     - If a type is found but is not a *btf.Struct, we fall back to modules.
-//  2. modules: see findStructTypeByNameFromModule.
-//
-// Returns (*btf.Struct, *btf.Spec, moduleID, nil) on success, or btf.ErrNotFound
-// if no matching struct is present in vmlinux or any module.
-func findStructTypeByName(s *btf.Spec, name string) (*btf.Struct, *btf.Spec, uint32, error) {
-	if s == nil {
-		return nil, nil, 0, fmt.Errorf("nil BTF: %w", btf.ErrNotFound)
-	}
-
-	t, err := s.AnyTypeByName(name)
-	if err == nil {
-		if typ, ok := btf.As[*btf.Struct](t); ok {
-			return typ, s, 0, nil
-		}
-	} else if !errors.Is(err, btf.ErrNotFound) {
-		return nil, nil, 0, err
-	}
-	return findStructTypeByNameFromModule(s, name)
-}
-
-// findStructTypeByNameFromModule scans all loaded kernel modules and tries
-// to resolve a struct type with the exact name. The iteration uses the base
-// vmlinux spec for string/ID interning as required by btf.Handle.Spec(base).
-//
-// For the first module where AnyTypeByName(name) returns a *btf.Struct,
-// the function returns that struct, the module's *btf.Spec, and its BTF ID.
-// If the type is not found in any module, btf.ErrNotFound is returned.
-func findStructTypeByNameFromModule(base *btf.Spec, name string) (*btf.Struct, *btf.Spec, uint32, error) {
-	it := new(btf.HandleIterator)
-
-	for it.Next() {
-		defer it.Handle.Close()
-
-		info, err := it.Handle.Info()
-		if err != nil {
-			return nil, nil, 0, fmt.Errorf("get info for BTF ID %d: %w", it.ID, err)
-		}
-
-		if !info.IsModule() {
-			continue
-		}
-
-		spec, err := it.Handle.Spec(base)
-		if err != nil {
-			return nil, nil, 0, fmt.Errorf("parse types for module %s: %w", info.Name, err)
-		}
-
-		t, err := spec.AnyTypeByName(name)
-		if errors.Is(err, btf.ErrNotFound) {
-			continue
-		}
-
-		if err != nil {
-			return nil, nil, 0, fmt.Errorf("lookup type in module %s: %w", info.Name, err)
-		}
-
-		if typ, ok := btf.As[*btf.Struct](t); ok {
-			return typ, spec, uint32(it.ID), nil
-		}
-	}
-
-	return nil, nil, 0, btf.ErrNotFound
-}
-
 // getStructMemberIndexByName returns the index of `member` within struct `s` by
 // comparing the member name.
 func getStructMemberIndexByName(s *btf.Struct, name string) int {

struct_ops_test.go

@@ -33,13 +33,15 @@ func TestCreateStructOpsMapSpecSimple(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	_, _, _, err = findStructTypeByName(s, "bpf_struct_ops_bpf_testmod_ops")
+	target := btf.Type((*btf.Struct)(nil))
+	_, module, err := findTargetInKernel(s, "bpf_struct_ops_bpf_testmod_ops", &target)
 	if errors.Is(err, btf.ErrNotFound) {
 		t.Skip("bpf_testmod_ops not loaded")
 	}
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer module.Close()
 
 	m, err := NewMap(ms)
 	testutils.SkipIfNotSupported(t, err)