collection: Add ReplaceMaps to CollectionOptions

Add a new way to replace references to specific maps when creating a
new collection by introducing a ReplaceMaps field to the CollectionOptions
struct.

The goal of this functionality is the same provided by RewriteMaps(),
but this new approach plays nicely with bpf2go as #517
is fixed.

Signed-off-by: Mauricio Vásquez <mauriciov@microsoft.com>

collection.go

@@ -18,6 +18,12 @@ import (
 type CollectionOptions struct {
 	Maps     MapOptions
 	Programs ProgramOptions
+
+	// OverrideMaps defines a set of maps that will be used instead
+	// of creating new ones when loading the object. The maps' specs
+	// should be compatible and there must exist a map in
+	// CollectionSpec.Maps for each map in OverrideMaps.
+	OverrideMaps map[string]*Map
 }
 
 // CollectionSpec describes a collection.
@@ -214,7 +220,10 @@ func (cs *CollectionSpec) Assign(to interface{}) error {
 // Returns an error if any of the fields can't be found, or
 // if the same Map or Program is assigned multiple times.
 func (cs *CollectionSpec) LoadAndAssign(to interface{}, opts *CollectionOptions) error {
-	loader := newCollectionLoader(cs, opts)
+	loader, err := newCollectionLoader(cs, opts)
+	if err != nil {
+		return err
+	}
 	defer loader.cleanup()
 
 	// Support assigning Programs and Maps, lazy-loading the required objects.
@@ -289,7 +298,10 @@ func NewCollection(spec *CollectionSpec) (*Collection, error) {
 // Omitting Collection.Close() during application shutdown is an error.
 // See the package documentation for details around Map and Program lifecycle.
 func NewCollectionWithOptions(spec *CollectionSpec, opts CollectionOptions) (*Collection, error) {
-	loader := newCollectionLoader(spec, &opts)
+	loader, err := newCollectionLoader(spec, &opts)
+	if err != nil {
+		return nil, err
+	}
 	defer loader.cleanup()
 
 	// Create maps first, as their fds need to be linked into programs.
@@ -379,18 +391,25 @@ type collectionLoader struct {
 	handles  *handleCache
 }
 
-func newCollectionLoader(coll *CollectionSpec, opts *CollectionOptions) *collectionLoader {
+func newCollectionLoader(coll *CollectionSpec, opts *CollectionOptions) (*collectionLoader, error) {
 	if opts == nil {
 		opts = &CollectionOptions{}
 	}
 
+	// check that there exist maps in coll.Maps for all maps in opts.OverrideMaps
+	for name := range opts.OverrideMaps {
+		if _, ok := coll.Maps[name]; !ok {
+			return nil, fmt.Errorf("override map %s doesn't exist in the spec", name)
+		}
+	}
+
 	return &collectionLoader{
 		coll,
 		opts,
 		make(map[string]*Map),
 		make(map[string]*Program),
 		newHandleCache(),
-	}
+	}, nil
 }
 
 // finalize should be called when all the collectionLoader's resources
@@ -426,9 +445,19 @@ func (cl *collectionLoader) loadMap(mapName string) (*Map, error) {
 		return nil, fmt.Errorf("map %s: BTF doesn't match collection", mapName)
 	}
 
-	m, err := newMapWithOptions(mapSpec, cl.opts.Maps, cl.handles)
-	if err != nil {
-		return nil, fmt.Errorf("map %s: %w", mapName, err)
+	var m *Map
+	var err error
+	var ok bool
+
+	if m, ok = cl.opts.OverrideMaps[mapName]; ok {
+		if err := mapSpec.checkCompatibility(m); err != nil {
+			return nil, fmt.Errorf("use override map %s: %w", mapSpec.Name, err)
+		}
+	} else {
+		m, err = newMapWithOptions(mapSpec, cl.opts.Maps, cl.handles)
+		if err != nil {
+			return nil, fmt.Errorf("map %s: %w", mapName, err)
+		}
 	}
 
 	cl.maps[mapName] = m

collection_test.go

@@ -1,6 +1,7 @@
 package ebpf
 
 import (
+	"errors"
 	"fmt"
 	"reflect"
 	"testing"
@@ -184,6 +185,106 @@ func TestCollectionSpecRewriteMaps(t *testing.T) {
 	}
 }
 
+func TestCollectionSpecOverrideMaps(t *testing.T) {
+	insns := asm.Instructions{
+		// R1 map
+		asm.LoadMapPtr(asm.R1, 0).WithReference("test-map"),
+		// R2 key
+		asm.Mov.Reg(asm.R2, asm.R10),
+		asm.Add.Imm(asm.R2, -4),
+		asm.StoreImm(asm.R2, 0, 0, asm.Word),
+		// Lookup map[0]
+		asm.FnMapLookupElem.Call(),
+		asm.JEq.Imm(asm.R0, 0, "ret"),
+		asm.LoadMem(asm.R0, asm.R0, 0, asm.Word),
+		asm.Return().WithSymbol("ret"),
+	}
+
+	cs := &CollectionSpec{
+		Maps: map[string]*MapSpec{
+			"test-map": {
+				Type:       Array,
+				KeySize:    4,
+				ValueSize:  4,
+				MaxEntries: 1,
+			},
+		},
+		Programs: map[string]*ProgramSpec{
+			"test-prog": {
+				Type:         SocketFilter,
+				Instructions: insns,
+				License:      "MIT",
+			},
+		},
+	}
+
+	// Override the map with another one
+	newMap, err := NewMap(cs.Maps["test-map"])
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer newMap.Close()
+
+	err = newMap.Put(uint32(0), uint32(2))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	coll1, err := NewCollectionWithOptions(cs, CollectionOptions{
+		OverrideMaps: map[string]*Map{
+			"test-map": newMap,
+		},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer coll1.Close()
+
+	ret, _, err := coll1.Programs["test-prog"].Test(make([]byte, 14))
+	testutils.SkipIfNotSupported(t, err)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if ret != 2 {
+		t.Fatal("new / override map not used")
+	}
+
+	// Override non-existing map
+	coll2, err := NewCollectionWithOptions(cs, CollectionOptions{
+		OverrideMaps: map[string]*Map{
+			"non-existing-map": newMap,
+		},
+	})
+	if err == nil {
+		coll2.Close()
+		t.Fatal("Overriding a non existing map did not fail")
+	}
+
+	// Override map with mismatching spec
+	newMap2, err := NewMap(&MapSpec{
+		Type:       Array,
+		KeySize:    4,
+		ValueSize:  8, // this is different
+		MaxEntries: 1,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	coll3, err := NewCollectionWithOptions(cs, CollectionOptions{
+		OverrideMaps: map[string]*Map{
+			"test-map": newMap2,
+		},
+	})
+	if err == nil {
+		coll3.Close()
+		t.Fatal("Overriding a map with a mismatching spec did not fail")
+	}
+	if !errors.Is(err, ErrMapIncompatible) {
+		t.Fatalf("Overriding a map with a mismatching spec failed with the wrong error")
+	}
+}
+
 func TestCollectionSpec_LoadAndAssign_LazyLoading(t *testing.T) {
 	spec := &CollectionSpec{
 		Maps: map[string]*MapSpec{

map.go

@@ -24,7 +24,7 @@ var (
 	ErrKeyNotExist      = errors.New("key does not exist")
 	ErrKeyExist         = errors.New("key already exists")
 	ErrIterationAborted = errors.New("iteration aborted")
-	ErrMapIncompatible  = errors.New("map's spec is incompatible with pinned map")
+	ErrMapIncompatible  = errors.New("map's spec is incompatible")
 )
 
 // MapOptions control loading a map into the kernel.