collection: Add Assign method

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{