From d5dcf5e6fc71e437ec318e0d2e4b65e1cf5bb9c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mauricio=20V=C3=A1squez?= Date: Fri, 22 Apr 2022 15:23:25 -0500 Subject: [PATCH] collection: Add MapReplacements to CollectionOptions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a new way to replace references to specific maps when creating a new collection by introducing a MapReplacements 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 https://github.com/cilium/ebpf/issues/517 is fixed. Signed-off-by: Mauricio Vásquez --- collection.go | 44 ++++++++++++-- collection_test.go | 139 +++++++++++++++++++++++++++++++++++++++++++++ map.go | 2 +- 3 files changed, 180 insertions(+), 5 deletions(-) diff --git a/collection.go b/collection.go index 4d6b35922..406cf22c9 100644 --- a/collection.go +++ b/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) diff --git a/collection_test.go b/collection_test.go index 00f3a9103..9b7bf22eb 100644 --- a/collection_test.go +++ b/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{ diff --git a/map.go b/map.go index b49b40187..21c163313 100644 --- a/map.go +++ b/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.