collection: make RewriteConstants operate on custom .rodata* sections

Instead of only looking for constants in .rodata, iterate over all maps with
an .rodata prefix.

Moved the datasec bytes and BTF info extraction into a method on MapSpec.
Lifted patchValue back into RewriteConstants.

Signed-off-by: Timo Beckers <timo@isovalent.com>

collection.go

@@ -122,34 +122,67 @@ func (cs *CollectionSpec) RewriteMaps(maps map[string]*Map) error {
 //
 // Returns an error if a constant doesn't exist.
 func (cs *CollectionSpec) RewriteConstants(consts map[string]interface{}) error {
-	rodata := cs.Maps[".rodata"]
-	if rodata == nil {
-		return errors.New("missing .rodata section")
-	}
+	replaced := make(map[string]bool)
 
-	if rodata.BTF == nil {
-		return errors.New(".rodata section has no BTF")
-	}
+	for name, spec := range cs.Maps {
+		if !strings.HasPrefix(name, ".rodata") {
+			continue
+		}
 
-	if n := len(rodata.Contents); n != 1 {
-		return fmt.Errorf("expected one key in .rodata, found %d", n)
-	}
+		b, ds, err := spec.dataSection()
+		if errors.Is(err, errMapNoBTFValue) {
+			// Data sections without a BTF Datasec are valid, but don't support
+			// constant replacements.
+			continue
+		}
+		if err != nil {
+			return fmt.Errorf("map %s: %w", name, err)
+		}
+
+		// MapSpec.Copy() performs a shallow copy. Fully copy the byte slice
+		// to avoid any changes affecting other copies of the MapSpec.
+		cpy := make([]byte, len(b))
+		copy(cpy, b)
+
+		for _, v := range ds.Vars {
+			vname := v.Type.TypeName()
+			replacement, ok := consts[vname]
+			if !ok {
+				continue
+			}
 
-	kv := rodata.Contents[0]
-	value, ok := kv.Value.([]byte)
-	if !ok {
-		return fmt.Errorf("first value in .rodata is %T not []byte", kv.Value)
+			if replaced[vname] {
+				return fmt.Errorf("section %s: duplicate variable %s", name, vname)
+			}
+
+			if int(v.Offset+v.Size) > len(cpy) {
+				return fmt.Errorf("section %s: offset %d(+%d) for variable %s is out of bounds", name, v.Offset, v.Size, vname)
+			}
+
+			b, err := marshalBytes(replacement, int(v.Size))
+			if err != nil {
+				return fmt.Errorf("marshaling constant replacement %s: %w", vname, err)
+			}
+
+			copy(cpy[v.Offset:v.Offset+v.Size], b)
+
+			replaced[vname] = true
+		}
+
+		spec.Contents[0] = MapKV{Key: uint32(0), Value: cpy}
 	}
 
-	buf := make([]byte, len(value))
-	copy(buf, value)
+	var missing []string
+	for c := range consts {
+		if !replaced[c] {
+			missing = append(missing, c)
+		}
+	}
 
-	err := patchValue(buf, rodata.Value, consts)
-	if err != nil {
-		return err
+	if len(missing) != 0 {
+		return fmt.Errorf("spec is missing one or more constants: %s", strings.Join(missing, ","))
 	}
 
-	rodata.Contents[0] = MapKV{kv.Key, buf}
 	return nil
 }
 

map.go

@@ -8,7 +8,6 @@ import (
 	"math/rand"
 	"path/filepath"
 	"reflect"
-	"strings"
 	"time"
 	"unsafe"
 
@@ -25,6 +24,7 @@ var (
 	ErrKeyExist         = errors.New("key already exists")
 	ErrIterationAborted = errors.New("iteration aborted")
 	ErrMapIncompatible  = errors.New("map spec is incompatible with existing map")
+	errMapNoBTFValue    = errors.New("map spec does not contain a BTF Value")
 )
 
 // MapOptions control loading a map into the kernel.
@@ -128,6 +128,31 @@ func (ms *MapSpec) clampPerfEventArraySize() error {
 	return nil
 }
 
+// dataSection returns the contents and BTF Datasec descriptor of the spec.
+func (ms *MapSpec) dataSection() ([]byte, *btf.Datasec, error) {
+
+	if ms.Value == nil {
+		return nil, nil, errMapNoBTFValue
+	}
+
+	ds, ok := ms.Value.(*btf.Datasec)
+	if !ok {
+		return nil, nil, fmt.Errorf("map value BTF is a %T, not a *btf.Datasec", ms.Value)
+	}
+
+	if n := len(ms.Contents); n != 1 {
+		return nil, nil, fmt.Errorf("expected one key, found %d", n)
+	}
+
+	kv := ms.Contents[0]
+	value, ok := kv.Value.([]byte)
+	if !ok {
+		return nil, nil, fmt.Errorf("value at first map key is %T, not []byte", kv.Value)
+	}
+
+	return value, ds, nil
+}
+
 // MapKV is used to initialize the contents of a Map.
 type MapKV struct {
 	Key   interface{}
@@ -1282,60 +1307,6 @@ func marshalMap(m *Map, length int) ([]byte, error) {
 	return buf, nil
 }
 
-func patchValue(value []byte, typ btf.Type, replacements map[string]interface{}) error {
-	replaced := make(map[string]bool)
-	replace := func(name string, offset, size int, replacement interface{}) error {
-		if offset+size > len(value) {
-			return fmt.Errorf("%s: offset %d(+%d) is out of bounds", name, offset, size)
-		}
-
-		buf, err := marshalBytes(replacement, size)
-		if err != nil {
-			return fmt.Errorf("marshal %s: %w", name, err)
-		}
-
-		copy(value[offset:offset+size], buf)
-		replaced[name] = true
-		return nil
-	}
-
-	switch parent := typ.(type) {
-	case *btf.Datasec:
-		for _, secinfo := range parent.Vars {
-			name := string(secinfo.Type.(*btf.Var).Name)
-			replacement, ok := replacements[name]
-			if !ok {
-				continue
-			}
-
-			err := replace(name, int(secinfo.Offset), int(secinfo.Size), replacement)
-			if err != nil {
-				return err
-			}
-		}
-
-	default:
-		return fmt.Errorf("patching %T is not supported", typ)
-	}
-
-	if len(replaced) == len(replacements) {
-		return nil
-	}
-
-	var missing []string
-	for name := range replacements {
-		if !replaced[name] {
-			missing = append(missing, name)
-		}
-	}
-
-	if len(missing) == 1 {
-		return fmt.Errorf("unknown field: %s", missing[0])
-	}
-
-	return fmt.Errorf("unknown fields: %s", strings.Join(missing, ","))
-}
-
 // MapIterator iterates a Map.
 //
 // See Map.Iterate.