From 2124a4d7f10886e6b4b8950561486002a55ec094 Mon Sep 17 00:00:00 2001
From: Alban Crequy <albancrequy@linux.microsoft.com>
Date: Tue, 17 Jan 2023 12:50:08 +0000
Subject: [PATCH] collection: define MissingConstantsError

RewriteConstants can return the error "spec is missing one or more
constants" if the constant given as parameter does not exist. This patch
defines the error as a specific type so it can be tested with:

    var mErr *MissingConstantsError
    if !errors.As(err, &mErr) { ... }

New unit tests are added, for both the success and failure cases.

Signed-off-by: Alban Crequy <albancrequy@linux.microsoft.com>
Co-developed-by: Lorenz Bauer <oss@lmb.io>
---
 collection.go      | 14 ++++++++++++--
 collection_test.go | 47 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/collection.go b/collection.go
index 729e5e9d7..5dec14a4e 100644
--- a/collection.go
+++ b/collection.go
@@ -107,6 +107,16 @@ func (cs *CollectionSpec) RewriteMaps(maps map[string]*Map) error {
 	return nil
 }
 
+// MissingConstantsError is returned by [CollectionSpec.RewriteConstants].
+type MissingConstantsError struct {
+	// The constants missing from .rodata.
+	Constants []string
+}
+
+func (m *MissingConstantsError) Error() string {
+	return fmt.Sprintf("some constants are missing from .rodata: %s", strings.Join(m.Constants, ", "))
+}
+
 // RewriteConstants replaces the value of multiple constants.
 //
 // The constant must be defined like so in the C program:
@@ -120,7 +130,7 @@ func (cs *CollectionSpec) RewriteMaps(maps map[string]*Map) error {
 //
 // From Linux 5.5 the verifier will use constants to eliminate dead code.
 //
-// Returns an error if a constant doesn't exist.
+// Returns an error wrapping [MissingConstantsError] if a constant doesn't exist.
 func (cs *CollectionSpec) RewriteConstants(consts map[string]interface{}) error {
 	replaced := make(map[string]bool)
 
@@ -184,7 +194,7 @@ func (cs *CollectionSpec) RewriteConstants(consts map[string]interface{}) error
 	}
 
 	if len(missing) != 0 {
-		return fmt.Errorf("spec is missing one or more constants: %s", strings.Join(missing, ","))
+		return fmt.Errorf("rewrite constants: %w", &MissingConstantsError{Constants: missing})
 	}
 
 	return nil
diff --git a/collection_test.go b/collection_test.go
index 68e888ba8..49722994d 100644
--- a/collection_test.go
+++ b/collection_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/cilium/ebpf/btf"
 	"github.com/cilium/ebpf/internal"
 	"github.com/cilium/ebpf/internal/testutils"
+	qt "github.com/frankban/quicktest"
 )
 
 func TestCollectionSpecNotModified(t *testing.T) {
@@ -328,6 +329,52 @@ func TestCollectionSpecMapReplacements_SpecMismatch(t *testing.T) {
 	}
 }
 
+func TestCollectionRewriteConstants(t *testing.T) {
+	cs := &CollectionSpec{
+		Maps: map[string]*MapSpec{
+			".rodata": {
+				Type:       Array,
+				KeySize:    4,
+				ValueSize:  4,
+				MaxEntries: 1,
+				Value: &btf.Datasec{
+					Vars: []btf.VarSecinfo{
+						{
+							Type: &btf.Var{
+								Name: "the_constant",
+								Type: &btf.Int{Size: 4},
+							},
+							Offset: 0,
+							Size:   4,
+						},
+					},
+				},
+				Contents: []MapKV{
+					{Key: uint32(0), Value: []byte{1, 1, 1, 1}},
+				},
+			},
+		},
+	}
+
+	err := cs.RewriteConstants(map[string]interface{}{
+		"fake_constant_one": uint32(1),
+		"fake_constant_two": uint32(2),
+	})
+	qt.Assert(t, err, qt.IsNotNil, qt.Commentf("RewriteConstants did not fail"))
+
+	var mErr *MissingConstantsError
+	if !errors.As(err, &mErr) {
+		t.Fatal("Error doesn't wrap MissingConstantsError:", err)
+	}
+	qt.Assert(t, mErr.Constants, qt.ContentEquals, []string{"fake_constant_one", "fake_constant_two"})
+
+	err = cs.RewriteConstants(map[string]interface{}{
+		"the_constant": uint32(0x42424242),
+	})
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, cs.Maps[".rodata"].Contents[0].Value, qt.ContentEquals, []byte{0x42, 0x42, 0x42, 0x42})
+}
+
 func TestCollectionSpec_LoadAndAssign_LazyLoading(t *testing.T) {
 	spec := &CollectionSpec{
 		Maps: map[string]*MapSpec{