token: use coll opts for setup

Signed-off-by: Andrej Svenke <anryko@nebius.com>

Author: anryko

collection.go

@@ -19,6 +19,11 @@ import (
 	"github.com/cilium/ebpf/internal/sys"
 )
 
+type TokenOptions struct {
+	Path    string
+	Enabled bool
+}
+
 // CollectionOptions control loading a collection into the kernel.
 //
 // Maps and Programs are passed to NewMapWithOptions and NewProgramsWithOptions.
@@ -36,6 +41,8 @@ type CollectionOptions struct {
 	// The given Maps are Clone()d before being used in the Collection, so the
 	// caller can Close() them freely when they are no longer needed.
 	MapReplacements map[string]*Map
+
+	Token TokenOptions
 }
 
 // CollectionSpec describes a collection.
@@ -323,6 +330,7 @@ type collectionLoader struct {
 	programs map[string]*Program
 	vars     map[string]*Variable
 	types    *btf.Cache
+	token    *sys.FD
 }
 
 func newCollectionLoader(coll *CollectionSpec, opts *CollectionOptions) (*collectionLoader, error) {
@@ -341,13 +349,23 @@ func newCollectionLoader(coll *CollectionSpec, opts *CollectionOptions) (*collec
 		return nil, fmt.Errorf("populating kallsyms caches: %w", err)
 	}
 
+	var token *sys.FD
+	var err error
+	if opts.Token.Enabled {
+		token, err = sys.BpffsGetTokenFD(opts.Token.Path)
+		if err != nil {
+			return nil, fmt.Errorf("getting bpf token for collection: %w", err)
+		}
+	}
+
 	return &collectionLoader{
 		coll,
 		opts,
 		make(map[string]*Map),
 		make(map[string]*Program),
 		make(map[string]*Variable),
 		btf.NewCache(),
+		token,
 	}, nil
 }
 
@@ -426,7 +444,7 @@ func (cl *collectionLoader) loadMap(mapName string) (*Map, error) {
 		return nil, fmt.Errorf("assembling contents of map %s: %w", mapName, err)
 	}
 
-	m, err := newMapWithOptions(mapSpec, cl.opts.Maps, cl.types)
+	m, err := newMapWithOptions(mapSpec, cl.opts.Maps, cl.types, cl.token)
 	if err != nil {
 		return nil, fmt.Errorf("map %s: %w", mapName, err)
 	}
@@ -489,7 +507,7 @@ func (cl *collectionLoader) loadProgram(progName string) (*Program, error) {
 		}
 	}
 
-	prog, err := newProgramWithOptions(progSpec, cl.opts.Programs, cl.types)
+	prog, err := newProgramWithOptions(progSpec, cl.opts.Programs, cl.types, cl.token)
 	if err != nil {
 		return nil, fmt.Errorf("program %s: %w", progName, err)
 	}

features/link.go

@@ -27,7 +27,7 @@ var haveBPFLinkUprobeMulti = internal.NewFeatureTest("bpf_link_uprobe_multi", fu
 		},
 		AttachType: ebpf.AttachTraceUprobeMulti,
 		License:    "MIT",
-	})
+	}, nil)
 	if errors.Is(err, unix.E2BIG) {
 		// Kernel doesn't support AttachType field.
 		return ebpf.ErrNotSupported
@@ -77,7 +77,7 @@ var haveBPFLinkKprobeMulti = internal.NewFeatureTest("bpf_link_kprobe_multi", fu
 		},
 		AttachType: ebpf.AttachTraceKprobeMulti,
 		License:    "MIT",
-	})
+	}, nil)
 	if errors.Is(err, unix.E2BIG) {
 		// Kernel doesn't support AttachType field.
 		return ebpf.ErrNotSupported
@@ -125,7 +125,7 @@ var haveBPFLinkKprobeSession = internal.NewFeatureTest("bpf_link_kprobe_session"
 		},
 		AttachType: ebpf.AttachTraceKprobeSession,
 		License:    "MIT",
-	})
+	}, nil)
 	if errors.Is(err, unix.E2BIG) {
 		// Kernel doesn't support AttachType field.
 		return ebpf.ErrNotSupported

features/prog.go

@@ -30,7 +30,7 @@ func probeProgram(spec *ebpf.ProgramSpec) error {
 	}
 	prog, err := ebpf.NewProgramWithOptions(spec, ebpf.ProgramOptions{
 		LogDisabled: true,
-	})
+	}, nil)
 	if err == nil {
 		prog.Close()
 	}
@@ -141,6 +141,7 @@ var haveProgramTypeMatrix = internal.FeatureMatrix[ebpf.ProgramType]{
 				ebpf.ProgramOptions{
 					LogDisabled: true,
 				},
+				nil,
 			)
 			if err != nil {
 				return err
@@ -271,7 +272,7 @@ func haveProgramHelper(pt ebpf.ProgramType, helper asm.BuiltinFunc) error {
 
 	prog, err := ebpf.NewProgramWithOptions(spec, ebpf.ProgramOptions{
 		LogLevel: 1,
-	})
+	}, nil)
 	if err == nil {
 		prog.Close()
 	}

internal/sys/syscall_bpffs.go

@@ -86,18 +86,18 @@ func MapCreateWithToken(attr *MapCreateAttr) (*FD, error) {
 }
 
 // ProgLoadWithToken try to load prog, on permission issue try using bpf token.
-func ProgLoadWithToken(attr *ProgLoadAttr) (*FD, error) {
-	fd, err := ProgLoad(attr)
-
-	// On permission error try privilege delegation using BPF Token.
-	if errors.Is(err, unix.EPERM) {
-		if tokenFD, tokenErr := BpffsGetTokenFD(""); tokenErr == nil {
-			defer tokenFD.Close()
-			attr.ProgTokenFd = int32(tokenFD.Int())
-			attr.ProgFlags |= BPF_F_TOKEN_FD
-			fd, err = ProgLoad(attr)
-		}
-	}
-
-	return fd, err
-}
+// func ProgLoadWithToken(attr *ProgLoadAttr) (*FD, error) {
+// 	fd, err := ProgLoad(attr)
+//
+// 	// On permission error try privilege delegation using BPF Token.
+// 	if errors.Is(err, unix.EPERM) {
+// 		if tokenFD, tokenErr := BpffsGetTokenFD(""); tokenErr == nil {
+// 			defer tokenFD.Close()
+// 			attr.ProgTokenFd = int32(tokenFD.Int())
+// 			attr.ProgFlags |= BPF_F_TOKEN_FD
+// 			fd, err = ProgLoad(attr)
+// 		}
+// 	}
+//
+// 	return fd, err
+// }

link/perf_event.go

@@ -313,7 +313,7 @@ var haveBPFLinkPerfEvent = internal.NewFeatureTest("bpf_link_perf_event", func()
 			asm.Return(),
 		},
 		License: "MIT",
-	})
+	}, nil)
 	if err != nil {
 		return err
 	}

link/syscalls.go

@@ -20,7 +20,7 @@ var haveProgAttach = internal.NewFeatureTest("BPF_PROG_ATTACH", func() error {
 			asm.Mov.Imm(asm.R0, 0),
 			asm.Return(),
 		},
-	})
+	}, nil)
 	if err != nil {
 		return internal.ErrNotSupported
 	}
@@ -45,7 +45,7 @@ var haveProgAttachReplace = internal.NewFeatureTest("BPF_PROG_ATTACH atomic repl
 			asm.Mov.Imm(asm.R0, 0),
 			asm.Return(),
 		},
-	})
+	}, nil)
 
 	if err != nil {
 		return internal.ErrNotSupported
@@ -119,7 +119,7 @@ var haveTCX = internal.NewFeatureTest("tcx", func() error {
 			asm.Mov.Imm(asm.R0, 0),
 			asm.Return(),
 		},
-	})
+	}, nil)
 
 	if err != nil {
 		return internal.ErrNotSupported
@@ -154,7 +154,7 @@ var haveNetkit = internal.NewFeatureTest("netkit", func() error {
 			asm.Mov.Imm(asm.R0, 0),
 			asm.Return(),
 		},
-	})
+	}, nil)
 
 	if err != nil {
 		return internal.ErrNotSupported

linker.go

@@ -461,7 +461,7 @@ func resolveKconfigReferences(insns asm.Instructions) (_ *Map, err error) {
 		return nil, err
 	}
 
-	kconfig, err := NewMap(cpy)
+	kconfig, err := NewMap(cpy, nil)
 	if err != nil {
 		return nil, err
 	}

map.go

@@ -361,8 +361,8 @@ func newMapFromFD(fd *sys.FD) (*Map, error) {
 // NewMap creates a new Map.
 //
 // It's equivalent to calling NewMapWithOptions with default options.
-func NewMap(spec *MapSpec) (*Map, error) {
-	return NewMapWithOptions(spec, MapOptions{})
+func NewMap(spec *MapSpec, token *sys.FD) (*Map, error) {
+	return NewMapWithOptions(spec, MapOptions{}, token)
 }
 
 // NewMapWithOptions creates a new Map.
@@ -375,8 +375,8 @@ func NewMap(spec *MapSpec) (*Map, error) {
 // by calling rlimit.RemoveMemlock() prior to calling NewMapWithOptions.
 //
 // May return an error wrapping ErrMapIncompatible.
-func NewMapWithOptions(spec *MapSpec, opts MapOptions) (*Map, error) {
-	m, err := newMapWithOptions(spec, opts, btf.NewCache())
+func NewMapWithOptions(spec *MapSpec, opts MapOptions, token *sys.FD) (*Map, error) {
+	m, err := newMapWithOptions(spec, opts, btf.NewCache(), token)
 	if err != nil {
 		return nil, fmt.Errorf("creating map: %w", err)
 	}
@@ -389,7 +389,7 @@ func NewMapWithOptions(spec *MapSpec, opts MapOptions) (*Map, error) {
 	return m, nil
 }
 
-func newMapWithOptions(spec *MapSpec, opts MapOptions, c *btf.Cache) (_ *Map, err error) {
+func newMapWithOptions(spec *MapSpec, opts MapOptions, c *btf.Cache, token *sys.FD) (_ *Map, err error) {
 	closeOnError := func(c io.Closer) {
 		if err != nil {
 			c.Close()
@@ -439,7 +439,7 @@ func newMapWithOptions(spec *MapSpec, opts MapOptions, c *btf.Cache) (_ *Map, er
 			return nil, errors.New("inner maps cannot be pinned")
 		}
 
-		template, err := spec.InnerMap.createMap(nil, c)
+		template, err := spec.InnerMap.createMap(nil, c, token)
 		if err != nil {
 			return nil, fmt.Errorf("inner map: %w", err)
 		}
@@ -451,7 +451,7 @@ func newMapWithOptions(spec *MapSpec, opts MapOptions, c *btf.Cache) (_ *Map, er
 		innerFd = template.fd
 	}
 
-	m, err := spec.createMap(innerFd, c)
+	m, err := spec.createMap(innerFd, c, token)
 	if err != nil {
 		return nil, err
 	}
@@ -546,7 +546,7 @@ func (m *Map) memorySize() (int, error) {
 
 // createMap validates the spec's properties and creates the map in the kernel
 // using the given opts. It does not populate or freeze the map.
-func (spec *MapSpec) createMap(inner *sys.FD, c *btf.Cache) (_ *Map, err error) {
+func (spec *MapSpec) createMap(inner *sys.FD, c *btf.Cache, token *sys.FD) (_ *Map, err error) {
 	closeOnError := func(closer io.Closer) {
 		if err != nil {
 			closer.Close()
@@ -584,6 +584,11 @@ func (spec *MapSpec) createMap(inner *sys.FD, c *btf.Cache) (_ *Map, err error)
 		MapExtra:   spec.MapExtra,
 	}
 
+	if token != nil {
+		attr.MapTokenFd = int32(token.Int())
+		attr.MapFlags |= sys.BPF_F_TOKEN_FD
+	}
+
 	if inner != nil {
 		attr.InnerMapFd = inner.Uint()
 	}
@@ -640,7 +645,7 @@ func (spec *MapSpec) createMap(inner *sys.FD, c *btf.Cache) (_ *Map, err error)
 		}
 	}
 
-	fd, err := sys.MapCreateWithToken(&attr)
+	fd, err := sys.MapCreate(&attr)
 
 	// Some map types don't support BTF k/v in earlier kernel versions.
 	// Remove BTF metadata and retry map creation.

prog.go

@@ -222,8 +222,8 @@ type Program struct {
 //
 // Returns a [VerifierError] containing the full verifier log if the program is
 // rejected by the kernel.
-func NewProgram(spec *ProgramSpec) (*Program, error) {
-	return NewProgramWithOptions(spec, ProgramOptions{})
+func NewProgram(spec *ProgramSpec, token *sys.FD) (*Program, error) {
+	return NewProgramWithOptions(spec, ProgramOptions{}, token)
 }
 
 // NewProgramWithOptions creates a new Program.
@@ -233,12 +233,12 @@ func NewProgram(spec *ProgramSpec) (*Program, error) {
 //
 // Returns a [VerifierError] containing the full verifier log if the program is
 // rejected by the kernel.
-func NewProgramWithOptions(spec *ProgramSpec, opts ProgramOptions) (*Program, error) {
+func NewProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, token *sys.FD) (*Program, error) {
 	if spec == nil {
 		return nil, errors.New("can't load a program from a nil spec")
 	}
 
-	prog, err := newProgramWithOptions(spec, opts, btf.NewCache())
+	prog, err := newProgramWithOptions(spec, opts, btf.NewCache(), token)
 	if errors.Is(err, asm.ErrUnsatisfiedMapReference) {
 		return nil, fmt.Errorf("cannot load program without loading its whole collection: %w", err)
 	}
@@ -254,7 +254,7 @@ var (
 	kfuncBadCall = []byte(fmt.Sprintf("invalid func unknown#%d\n", kfuncCallPoisonBase))
 )
 
-func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache) (*Program, error) {
+func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache, token *sys.FD) (*Program, error) {
 	if len(spec.Instructions) == 0 {
 		return nil, errors.New("instructions cannot be empty")
 	}
@@ -295,6 +295,11 @@ func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache)
 		KernVersion:        kv,
 	}
 
+	if token != nil {
+		attr.ProgTokenFd = int32(token.Int())
+		attr.ProgFlags |= sys.BPF_F_TOKEN_FD
+	}
+
 	insns := make(asm.Instructions, len(spec.Instructions))
 	copy(insns, spec.Instructions)
 
@@ -441,7 +446,7 @@ func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache)
 	var fd *sys.FD
 	if opts.LogDisabled {
 		// Loading with logging disabled should never retry.
-		fd, err = sys.ProgLoadWithToken(attr)
+		fd, err = sys.ProgLoad(attr)
 		if err == nil {
 			return &Program{"", fd, spec.Name, "", spec.Type}, nil
 		}
@@ -461,7 +466,7 @@ func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache)
 				attr.LogBuf = sys.SlicePointer(logBuf)
 			}
 
-			fd, err = sys.ProgLoadWithToken(attr)
+			fd, err = sys.ProgLoad(attr)
 			if err == nil {
 				return &Program{unix.ByteSliceToString(logBuf), fd, spec.Name, "", spec.Type}, nil
 			}
@@ -842,7 +847,7 @@ var haveProgRun = internal.NewFeatureTest("BPF_PROG_RUN", func() error {
 			asm.Return(),
 		},
 		License: "MIT",
-	})
+	}, nil)
 	if err != nil {
 		// This may be because we lack sufficient permissions, etc.
 		return err