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 <dxu@dxuuu.xyz>

Author: danobi

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.

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{