From 786a3be46f8d1112ebc61542c54ad9164983c934 Mon Sep 17 00:00:00 2001 From: Daniel Xu Date: Wed, 2 Aug 2023 16:51:18 -0600 Subject: [PATCH 1/2] collection: Add Assign method Assign() bridges bpf2go generated code and functionality implemented in Collection. A few things are easier to do in Collection, like iterating through maps and programs. It is harder to guarantee invariants like map and program initialization using reflection as well as being more complicated. Signed-off-by: Daniel Xu --- collection.go | 65 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) 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. From d391bc3379eef6a2dc29f80a61ad580a8b6588b9 Mon Sep 17 00:00:00 2001 From: Daniel Xu Date: Wed, 2 Aug 2023 17:05:29 -0600 Subject: [PATCH 2/2] test: collection: Add test for Assign Signed-off-by: Daniel Xu --- collection_test.go | 85 +++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 84 insertions(+), 1 deletion(-) 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{