program: use cached BTF for fixups and target search

btf/kernel.go

@@ -225,9 +225,9 @@ func findVMLinux() (*os.File, error) {
 //
 // It is not safe for concurrent use.
 type Cache struct {
-	KernelTypes   *Spec
-	ModuleTypes   map[string]*Spec
-	LoadedModules []string
+	kernelTypes   *Spec
+	moduleTypes   map[string]*Spec
+	loadedModules []string
 }
 
 // NewCache creates a new Cache.
@@ -262,26 +262,26 @@ func NewCache() *Cache {
 // Kernel is equivalent to [LoadKernelSpec], except that repeated calls do
 // not copy the Spec.
 func (c *Cache) Kernel() (*Spec, error) {
-	if c.KernelTypes != nil {
-		return c.KernelTypes, nil
+	if c.kernelTypes != nil {
+		return c.kernelTypes, nil
 	}
 
 	var err error
-	c.KernelTypes, err = LoadKernelSpec()
-	return c.KernelTypes, err
+	c.kernelTypes, err = LoadKernelSpec()
+	return c.kernelTypes, err
 }
 
 // Module is equivalent to [LoadKernelModuleSpec], except that repeated calls do
 // not copy the spec.
 //
 // All modules also share the return value of [Kernel] as their base.
 func (c *Cache) Module(name string) (*Spec, error) {
-	if spec := c.ModuleTypes[name]; spec != nil {
+	if spec := c.moduleTypes[name]; spec != nil {
 		return spec, nil
 	}
 
-	if c.ModuleTypes == nil {
-		c.ModuleTypes = make(map[string]*Spec)
+	if c.moduleTypes == nil {
+		c.moduleTypes = make(map[string]*Spec)
 	}
 
 	base, err := c.Kernel()
@@ -302,14 +302,14 @@ func (c *Cache) Module(name string) (*Spec, error) {
 	}
 
 	spec = &Spec{decoder: decoder}
-	c.ModuleTypes[name] = spec
+	c.moduleTypes[name] = spec
 	return spec, err
 }
 
 // Modules returns a sorted list of all loaded modules.
 func (c *Cache) Modules() ([]string, error) {
-	if c.LoadedModules != nil {
-		return c.LoadedModules, nil
+	if c.loadedModules != nil {
+		return c.loadedModules, nil
 	}
 
 	btfDir, err := os.Open("/sys/kernel/btf")
@@ -328,6 +328,6 @@ func (c *Cache) Modules() ([]string, error) {
 	})
 
 	sort.Strings(entries)
-	c.LoadedModules = entries
+	c.loadedModules = entries
 	return entries, nil
 }

btf/kernel_test.go

@@ -43,9 +43,9 @@ func TestCache(t *testing.T) {
 	FlushKernelSpec()
 	c := NewCache()
 
-	qt.Assert(t, qt.IsNil(c.KernelTypes))
-	qt.Assert(t, qt.HasLen(c.ModuleTypes, 0))
-	qt.Assert(t, qt.IsNil(c.LoadedModules))
+	qt.Assert(t, qt.IsNil(c.kernelTypes))
+	qt.Assert(t, qt.HasLen(c.moduleTypes, 0))
+	qt.Assert(t, qt.IsNil(c.loadedModules))
 
 	// Test that Kernel() creates only one copy
 	spec1, err := c.Kernel()
@@ -83,15 +83,15 @@ func TestCache(t *testing.T) {
 
 	// Test that NewCache populates from global cache
 	c = NewCache()
-	qt.Assert(t, qt.IsNotNil(c.KernelTypes))
-	qt.Assert(t, qt.Not(qt.Equals(c.KernelTypes, vmlinux)))
+	qt.Assert(t, qt.IsNotNil(c.kernelTypes))
+	qt.Assert(t, qt.Not(qt.Equals(c.kernelTypes, vmlinux)))
 	if testmod != nil {
-		qt.Assert(t, qt.IsNotNil(c.ModuleTypes["bpf_testmod"]))
-		qt.Assert(t, qt.Not(qt.Equals(c.ModuleTypes["bpf_testmod"], testmod)))
+		qt.Assert(t, qt.IsNotNil(c.moduleTypes["bpf_testmod"]))
+		qt.Assert(t, qt.Not(qt.Equals(c.moduleTypes["bpf_testmod"], testmod)))
 	}
 
 	// Test that Modules only reads modules once.
 	_, err = c.Modules()
 	qt.Assert(t, qt.IsNil(err))
-	qt.Assert(t, qt.IsNotNil(c.LoadedModules))
+	qt.Assert(t, qt.IsNotNil(c.loadedModules))
 }

collection.go

@@ -437,7 +437,7 @@ func newCollectionLoader(coll *CollectionSpec, opts *CollectionOptions) (*collec
 		make(map[string]*Map),
 		make(map[string]*Program),
 		make(map[string]*Variable),
-		newBTFCache(&opts.Programs),
+		btf.NewCache(),
 	}, nil
 }
 

linker.go

@@ -122,9 +122,8 @@ func hasFunctionReferences(insns asm.Instructions) bool {
 
 // applyRelocations collects and applies any CO-RE relocations in insns.
 //
-// Passing a nil target will relocate against the running kernel. insns are
-// modified in place.
-func applyRelocations(insns asm.Instructions, bo binary.ByteOrder, b *btf.Builder, c *btf.Cache, extraTargets []*btf.Spec) error {
+// insns are modified in place.
+func applyRelocations(insns asm.Instructions, bo binary.ByteOrder, b *btf.Builder, c *btf.Cache, kernelOverride *btf.Spec, extraTargets []*btf.Spec) error {
 	var relos []*btf.CORERelocation
 	var reloInsns []*asm.Instruction
 	iter := insns.Iterate()
@@ -143,28 +142,35 @@ func applyRelocations(insns asm.Instructions, bo binary.ByteOrder, b *btf.Builde
 		bo = internal.NativeEndian
 	}
 
-	kernelTarget, err := c.Kernel()
-	if err != nil {
-		return fmt.Errorf("load kernel spec: %w", err)
-	}
+	var targets []*btf.Spec
+	if kernelOverride == nil {
+		kernel, err := c.Kernel()
+		if err != nil {
+			return fmt.Errorf("load kernel spec: %w", err)
+		}
 
-	modules, err := c.Modules()
-	// Ignore ErrNotExists to cater to kernels which have CONFIG_DEBUG_INFO_BTF_MODULES
-	// or CONFIG_DEBUG_INFO_BTF disabled.
-	if err != nil && !errors.Is(err, fs.ErrNotExist) {
-		return err
-	}
+		modules, err := c.Modules()
+		// Ignore ErrNotExists to cater to kernels which have CONFIG_DEBUG_INFO_BTF_MODULES
+		// or CONFIG_DEBUG_INFO_BTF disabled.
+		if err != nil && !errors.Is(err, fs.ErrNotExist) {
+			return err
+		}
 
-	targets := make([]*btf.Spec, 0, 1+len(modules)+len(extraTargets))
-	targets = append(targets, kernelTarget)
+		targets = make([]*btf.Spec, 0, 1+len(modules)+len(extraTargets))
+		targets = append(targets, kernel)
 
-	for _, kmod := range modules {
-		spec, err := c.Module(kmod)
-		if err != nil {
-			return fmt.Errorf("load BTF for kmod %s: %w", kmod, err)
-		}
+		for _, kmod := range modules {
+			spec, err := c.Module(kmod)
+			if err != nil {
+				return fmt.Errorf("load BTF for kmod %s: %w", kmod, err)
+			}
 
-		targets = append(targets, spec)
+			targets = append(targets, spec)
+		}
+	} else {
+		// We expect kernelOverride to contain the merged types
+		// of vmlinux and kernel modules, as distributed by btfhub.
+		targets = []*btf.Spec{kernelOverride}
 	}
 
 	targets = append(targets, extraTargets...)
@@ -289,7 +295,7 @@ const kfuncCallPoisonBase = 0xdedc0de
 // fixupKfuncs loops over all instructions in search for kfunc calls.
 // If at least one is found, the current kernels BTF and module BTFis are searched to set Instruction.Constant
 // and Instruction.Offset to the correct values.
-func fixupKfuncs(insns asm.Instructions) (_ handles, err error) {
+func fixupKfuncs(insns asm.Instructions, cache *btf.Cache) (_ handles, err error) {
 	closeOnError := func(c io.Closer) {
 		if err != nil {
 			c.Close()
@@ -310,7 +316,7 @@ fixups:
 	// Only load kernel BTF if we found at least one kfunc call. kernelSpec can be
 	// nil if the kernel does not have BTF, in which case we poison all kfunc
 	// calls.
-	kernelSpec, err := btf.LoadKernelSpec()
+	_, err = cache.Kernel()
 	// ErrNotSupportedOnOS wraps ErrNotSupported, check for it first.
 	if errors.Is(err, internal.ErrNotSupportedOnOS) {
 		return nil, fmt.Errorf("kfuncs are not supported on this platform: %w", err)
@@ -342,7 +348,7 @@ fixups:
 
 		// findTargetInKernel returns btf.ErrNotFound if the input btf.Spec is nil.
 		target := btf.Type((*btf.Func)(nil))
-		spec, module, err := findTargetInKernel(kernelSpec, kfm.Func.Name, &target)
+		spec, module, err := findTargetInKernel(kfm.Func.Name, &target, cache)
 		if errors.Is(err, btf.ErrNotFound) {
 			if kfm.Binding == elf.STB_WEAK {
 				if ins.IsKfuncCall() {

prog.go

@@ -236,7 +236,7 @@ func NewProgramWithOptions(spec *ProgramSpec, opts ProgramOptions) (*Program, er
 		return nil, errors.New("can't load a program from a nil spec")
 	}
 
-	prog, err := newProgramWithOptions(spec, opts, newBTFCache(&opts))
+	prog, err := newProgramWithOptions(spec, opts, btf.NewCache())
 	if errors.Is(err, asm.ErrUnsatisfiedMapReference) {
 		return nil, fmt.Errorf("cannot load program without loading its whole collection: %w", err)
 	}
@@ -297,7 +297,7 @@ func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache)
 	copy(insns, spec.Instructions)
 
 	var b btf.Builder
-	if err := applyRelocations(insns, spec.ByteOrder, &b, c, opts.ExtraRelocationTargets); err != nil {
+	if err := applyRelocations(insns, spec.ByteOrder, &b, c, opts.KernelTypes, opts.ExtraRelocationTargets); err != nil {
 		return nil, fmt.Errorf("apply CO-RE relocations: %w", err)
 	}
 
@@ -350,7 +350,7 @@ func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache)
 		return nil, err
 	}
 
-	handles, err := fixupKfuncs(insns)
+	handles, err := fixupKfuncs(insns, c)
 	if err != nil {
 		return nil, fmt.Errorf("fixing up kfuncs: %w", err)
 	}
@@ -381,7 +381,7 @@ func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache)
 		attr.AttachBtfObjFd = uint32(spec.AttachTarget.FD())
 		defer runtime.KeepAlive(spec.AttachTarget)
 	} else if spec.AttachTo != "" {
-		module, targetID, err := findProgramTargetInKernel(spec.AttachTo, spec.Type, spec.AttachType)
+		module, targetID, err := findProgramTargetInKernel(spec.AttachTo, spec.Type, spec.AttachType, c)
 		if err != nil && !errors.Is(err, errUnrecognizedAttachType) {
 			// We ignore errUnrecognizedAttachType since AttachTo may be non-empty
 			// for programs that don't attach anywhere.
@@ -1027,7 +1027,7 @@ var errUnrecognizedAttachType = errors.New("unrecognized attach type")
 //
 // Returns errUnrecognizedAttachType if the combination of progType and attachType
 // is not recognised.
-func findProgramTargetInKernel(name string, progType ProgramType, attachType AttachType) (*btf.Handle, btf.TypeID, error) {
+func findProgramTargetInKernel(name string, progType ProgramType, attachType AttachType, cache *btf.Cache) (*btf.Handle, btf.TypeID, error) {
 	type match struct {
 		p ProgramType
 		a AttachType
@@ -1067,12 +1067,7 @@ func findProgramTargetInKernel(name string, progType ProgramType, attachType Att
 		return nil, 0, errUnrecognizedAttachType
 	}
 
-	spec, err := btf.LoadKernelSpec()
-	if err != nil {
-		return nil, 0, fmt.Errorf("load kernel spec: %w", err)
-	}
-
-	spec, module, err := findTargetInKernel(spec, typeName, &target)
+	spec, module, err := findTargetInKernel(typeName, &target, cache)
 	if errors.Is(err, btf.ErrNotFound) {
 		return nil, 0, &internal.UnsupportedFeatureError{Name: featureName}
 	}
@@ -1102,14 +1097,15 @@ func findProgramTargetInKernel(name string, progType ProgramType, attachType Att
 //
 // Returns a non-nil handle if the type was found in a module, or btf.ErrNotFound
 // if the type wasn't found at all.
-func findTargetInKernel(kernelSpec *btf.Spec, typeName string, target *btf.Type) (*btf.Spec, *btf.Handle, error) {
-	if kernelSpec == nil {
-		return nil, nil, fmt.Errorf("nil kernelSpec: %w", btf.ErrNotFound)
+func findTargetInKernel(typeName string, target *btf.Type, cache *btf.Cache) (*btf.Spec, *btf.Handle, error) {
+	kernelSpec, err := cache.Kernel()
+	if err != nil {
+		return nil, nil, err
 	}
 
-	err := kernelSpec.TypeByName(typeName, target)
+	err = kernelSpec.TypeByName(typeName, target)
 	if errors.Is(err, btf.ErrNotFound) {
-		spec, module, err := findTargetInModule(kernelSpec, typeName, target)
+		spec, module, err := findTargetInModule(typeName, target, cache)
 		if err != nil {
 			return nil, nil, fmt.Errorf("find target in modules: %w", err)
 		}
@@ -1127,7 +1123,7 @@ func findTargetInKernel(kernelSpec *btf.Spec, typeName string, target *btf.Type)
 // are searched in the order they were loaded.
 //
 // Returns btf.ErrNotFound if the target can't be found in any module.
-func findTargetInModule(base *btf.Spec, typeName string, target *btf.Type) (*btf.Spec, *btf.Handle, error) {
+func findTargetInModule(typeName string, target *btf.Type, cache *btf.Cache) (*btf.Spec, *btf.Handle, error) {
 	it := new(btf.HandleIterator)
 	defer it.Handle.Close()
 
@@ -1141,7 +1137,7 @@ func findTargetInModule(base *btf.Spec, typeName string, target *btf.Type) (*btf
 			continue
 		}
 
-		spec, err := it.Handle.Spec(base)
+		spec, err := cache.Module(info.Name)
 		if err != nil {
 			return nil, nil, fmt.Errorf("parse types for module %s: %w", info.Name, err)
 		}
@@ -1201,12 +1197,3 @@ func findTargetInProgram(prog *Program, name string, progType ProgramType, attac
 
 	return spec.TypeID(targetFunc)
 }
-
-func newBTFCache(opts *ProgramOptions) *btf.Cache {
-	c := btf.NewCache()
-	if opts.KernelTypes != nil {
-		c.KernelTypes = opts.KernelTypes
-		c.LoadedModules = []string{}
-	}
-	return c
-}