diff --git a/collection.go b/collection.go index a581ecf44..ac7108e62 100644 --- a/collection.go +++ b/collection.go @@ -695,6 +695,71 @@ func LoadCollection(file string) (*Collection, error) { return NewCollection(spec) } +// Assign the contents of a Collection to a struct. +// +// This function bridges functionality between bpf2go generated +// code and any functionality better implemented in Collection. +// +// 'to' must be a pointer to a struct. A field of the +// struct is updated with values from Programs or Maps if it +// has an `ebpf` tag and its type is *Program or *Map. +// The tag's value specifies the name of the program or map as +// found in the CollectionSpec. +// +// struct { +// Foo *ebpf.Program `ebpf:"xdp_foo"` +// Bar *ebpf.Map `ebpf:"bar_map"` +// Ignored int +// } +// +// Returns an error if any of the eBPF objects can't be found, or +// if the same Map or Program is assigned multiple times. +// +// Ownership and Close()ing responsibility is transferred to `to` +// for any successful assigns. On error `to` is left in an undefined state. +func (coll *Collection) Assign(to interface{}) error { + assignedMaps := make(map[string]bool) + assignedProgs := make(map[string]bool) + + // Assign() only transfers already-loaded Maps and Programs. No extra + // loading is done. + getValue := func(typ reflect.Type, name string) (interface{}, error) { + switch typ { + + case reflect.TypeOf((*Program)(nil)): + if p := coll.Programs[name]; p != nil { + assignedProgs[name] = true + return p, nil + } + return nil, fmt.Errorf("missing program %q", name) + + case reflect.TypeOf((*Map)(nil)): + if m := coll.Maps[name]; m != nil { + assignedMaps[name] = true + return m, nil + } + return nil, fmt.Errorf("missing map %q", name) + + default: + return nil, fmt.Errorf("unsupported type %s", typ) + } + } + + if err := assignValues(to, getValue); err != nil { + return err + } + + // Finalize ownership transfer + for p := range assignedProgs { + delete(coll.Programs, p) + } + for m := range assignedMaps { + delete(coll.Maps, m) + } + + return nil +} + // Close frees all maps and programs associated with the collection. // // The collection mustn't be used afterwards. diff --git a/collection_test.go b/collection_test.go index f3dab3e1b..5c3d5b9e9 100644 --- a/collection_test.go +++ b/collection_test.go @@ -434,7 +434,7 @@ func TestCollectionSpec_LoadAndAssign_LazyLoading(t *testing.T) { } } -func TestCollectionAssign(t *testing.T) { +func TestCollectionSpecAssign(t *testing.T) { var specs struct { Program *ProgramSpec `ebpf:"prog1"` Map *MapSpec `ebpf:"map1"` @@ -559,6 +559,89 @@ func TestAssignValues(t *testing.T) { } +func TestCollectionAssign(t *testing.T) { + var objs struct { + Program *Program `ebpf:"prog1"` + Map *Map `ebpf:"map1"` + } + + cs := &CollectionSpec{ + Maps: map[string]*MapSpec{ + "map1": { + Type: Array, + KeySize: 4, + ValueSize: 4, + MaxEntries: 1, + }, + }, + Programs: map[string]*ProgramSpec{ + "prog1": { + Type: SocketFilter, + Instructions: asm.Instructions{ + asm.LoadImm(asm.R0, 0, asm.DWord), + asm.Return(), + }, + License: "MIT", + }, + }, + } + + coll, err := NewCollection(cs) + qt.Assert(t, err, qt.IsNil) + defer coll.Close() + + qt.Assert(t, coll.Assign(&objs), qt.IsNil) + defer objs.Program.Close() + defer objs.Map.Close() + + // Check that objs has received ownership of map and prog + qt.Assert(t, objs.Program.FD() >= 0, qt.IsTrue) + qt.Assert(t, objs.Map.FD() >= 0, qt.IsTrue) + + // Check that the collection has lost ownership + qt.Assert(t, coll.Programs["prog1"], qt.IsNil) + qt.Assert(t, coll.Maps["map1"], qt.IsNil) +} + +func TestCollectionAssignFail(t *testing.T) { + // `map2` does not exist + var objs struct { + Program *Program `ebpf:"prog1"` + Map *Map `ebpf:"map2"` + } + + cs := &CollectionSpec{ + Maps: map[string]*MapSpec{ + "map1": { + Type: Array, + KeySize: 4, + ValueSize: 4, + MaxEntries: 1, + }, + }, + Programs: map[string]*ProgramSpec{ + "prog1": { + Type: SocketFilter, + Instructions: asm.Instructions{ + asm.LoadImm(asm.R0, 0, asm.DWord), + asm.Return(), + }, + License: "MIT", + }, + }, + } + + coll, err := NewCollection(cs) + qt.Assert(t, err, qt.IsNil) + defer coll.Close() + + qt.Assert(t, coll.Assign(&objs), qt.IsNotNil) + + // Check that the collection has retained ownership + qt.Assert(t, coll.Programs["prog1"], qt.IsNotNil) + qt.Assert(t, coll.Maps["map1"], qt.IsNotNil) +} + func TestIncompleteLoadAndAssign(t *testing.T) { spec := &CollectionSpec{ Programs: map[string]*ProgramSpec{