collection: Add ReplaceMaps to CollectionOptions

collection.go

@@ -18,6 +18,15 @@ import (
 type CollectionOptions struct {
 	Maps     MapOptions
 	Programs ProgramOptions
+
+	// MapReplacements 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 MapReplacements.
+	// MapReplacements are cloned before being used in the
+	// Collection, so the user can Close() them if not needed
+	// anymore.
+	MapReplacements map[string]*Map
 }
 
 // CollectionSpec describes a collection.
@@ -214,7 +223,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 +301,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 +394,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 for existing MapSpecs in the CollectionSpec for all provided replacement maps.
+	for name := range opts.MapReplacements {
+		if _, ok := coll.Maps[name]; !ok {
+			return nil, fmt.Errorf("replacement map %s not found in CollectionSpec", 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,6 +448,20 @@ func (cl *collectionLoader) loadMap(mapName string) (*Map, error) {
 		return nil, fmt.Errorf("map %s: BTF doesn't match collection", mapName)
 	}
 
+	if replaceMap, ok := cl.opts.MapReplacements[mapName]; ok {
+		if err := mapSpec.checkCompatibility(replaceMap); err != nil {
+			return nil, fmt.Errorf("use replacement map %s: %w", mapSpec.Name, err)
+		}
+		// Clone the map to avoid closing user's map later on.
+		m, err := replaceMap.Clone()
+		if err != nil {
+			return nil, err
+		}
+
+		cl.maps[mapName] = m
+		return m, nil
+	}
+
 	m, err := newMapWithOptions(mapSpec, cl.opts.Maps, cl.handles)
 	if err != nil {
 		return nil, fmt.Errorf("map %s: %w", mapName, err)

collection_test.go

@@ -1,6 +1,7 @@
 package ebpf
 
 import (
+	"errors"
 	"fmt"
 	"reflect"
 	"testing"
@@ -184,6 +185,144 @@ func TestCollectionSpecRewriteMaps(t *testing.T) {
 	}
 }
 
+func TestCollectionSpecMapReplacements(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",
+			},
+		},
+	}
+
+	// Replace 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)
+	}
+
+	coll, err := NewCollectionWithOptions(cs, CollectionOptions{
+		MapReplacements: map[string]*Map{
+			"test-map": newMap,
+		},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer coll.Close()
+
+	ret, _, err := coll.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")
+	}
+
+	// Check that newMap isn't closed when the collection is closed
+	coll.Close()
+	err = newMap.Put(uint32(0), uint32(3))
+	if err != nil {
+		t.Fatalf("failed to update replaced map: %s", err)
+	}
+}
+func TestCollectionSpecMapReplacements_NonExistingMap(t *testing.T) {
+	cs := &CollectionSpec{
+		Maps: map[string]*MapSpec{
+			"test-map": {
+				Type:       Array,
+				KeySize:    4,
+				ValueSize:  4,
+				MaxEntries: 1,
+			},
+		},
+	}
+
+	// Override non-existing map
+	newMap, err := NewMap(cs.Maps["test-map"])
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer newMap.Close()
+
+	coll, err := NewCollectionWithOptions(cs, CollectionOptions{
+		MapReplacements: map[string]*Map{
+			"non-existing-map": newMap,
+		},
+	})
+	if err == nil {
+		coll.Close()
+		t.Fatal("Overriding a non existing map did not fail")
+	}
+}
+
+func TestCollectionSpecMapReplacements_SpecMismatch(t *testing.T) {
+	cs := &CollectionSpec{
+		Maps: map[string]*MapSpec{
+			"test-map": {
+				Type:       Array,
+				KeySize:    4,
+				ValueSize:  4,
+				MaxEntries: 1,
+			},
+		},
+	}
+
+	// Override map with mismatching spec
+	newMap, err := NewMap(&MapSpec{
+		Type:       Array,
+		KeySize:    4,
+		ValueSize:  8, // this is different
+		MaxEntries: 1,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	coll, err := NewCollectionWithOptions(cs, CollectionOptions{
+		MapReplacements: map[string]*Map{
+			"test-map": newMap,
+		},
+	})
+	if err == nil {
+		coll.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 spec is incompatible with existing map")
 )
 
 // MapOptions control loading a map into the kernel.