variable: add accessors for values represented by a Variable

Signed-off-by: Timo Beckers <[email protected]>

testdata/variables.c

@@ -38,4 +38,13 @@ __section("socket") int check_struct() {
 
 // Variable aligned on page boundary to ensure all bytes in the mapping can be
 // accessed through the Variable API.
-volatile char var_array[8192] __section(".data.array");
+volatile uint8_t var_array[8192] __section(".data.array");
+__section("socket") int check_array() {
+	return var_array[sizeof(var_array) - 1] == 0xff;
+}
+
+volatile uint32_t var_atomic __section(".data.atomic");
+__section("socket") int add_atomic() {
+	__sync_fetch_and_add(&var_atomic, 1);
+	return 0;
+}

variable.go

@@ -1,6 +1,7 @@
 package ebpf
 
 import (
+	"encoding/binary"
 	"fmt"
 	"io"
 
@@ -232,3 +233,37 @@ func (v *Variable) Get(out any) error {
 
 	return nil
 }
+
+func checkVariable[T any](v *Variable) error {
+	if v.ReadOnly() {
+		return ErrReadOnly
+	}
+
+	var t T
+	size := binary.Size(t)
+	if size < 0 {
+		return fmt.Errorf("can't determine size of type %T: %w", t, ErrInvalidType)
+	}
+
+	if v.size != uint64(size) {
+		return fmt.Errorf("can't create %d-byte accessor to %d-byte variable", size, v.size)
+	}
+	return nil
+}
+
+// VariablePointer returns a pointer to a variable of type T backed by memory
+// shared with the BPF program. Requires
+// [CollectionOptions.UnsafeVariableExperiment] to be true.
+//
+// Taking a pointer to a read-only Variable is not supported. T must be a
+// fixed-size type according to [binary.Size]. Types containing Go pointers are
+// not valid.
+//
+// When accessing structs, embedding [structs.HostLayout] may help ensure the
+// layout of the Go struct matches the one in the BPF C program.
+func VariablePointer[T any](v *Variable) (*T, error) {
+	if err := checkVariable[T](v); err != nil {
+		return nil, fmt.Errorf("variable pointer %s: %w", v.name, err)
+	}
+	return memoryPointer[T](v.mm, v.offset)
+}

variable_test.go

@@ -1,6 +1,9 @@
 package ebpf
 
 import (
+	"runtime"
+	"structs"
+	"sync/atomic"
 	"testing"
 
 	"github.com/go-quicktest/qt"
@@ -182,3 +185,111 @@ func TestVariableFallback(t *testing.T) {
 		qt.Assert(t, qt.ErrorIs(err, ErrNotSupported))
 	}
 }
+
+func TestVariablePointer(t *testing.T) {
+	testutils.SkipIfNotSupported(t, haveMmapableMaps())
+
+	file := testutils.NativeFile(t, "testdata/variables-%s.elf")
+	spec, err := LoadCollectionSpec(file)
+	qt.Assert(t, qt.IsNil(err))
+
+	obj := struct {
+		AddAtomic   *Program `ebpf:"add_atomic"`
+		CheckStruct *Program `ebpf:"check_struct"`
+		CheckArray  *Program `ebpf:"check_array"`
+
+		Atomic *Variable `ebpf:"var_atomic"`
+		Struct *Variable `ebpf:"var_struct"`
+		Array  *Variable `ebpf:"var_array"`
+	}{}
+
+	qt.Assert(t, qt.IsNil(spec.LoadAndAssign(&obj, &CollectionOptions{UnsafeVariableExperiment: true})))
+	t.Cleanup(func() {
+		obj.AddAtomic.Close()
+		obj.CheckStruct.Close()
+		obj.CheckArray.Close()
+	})
+
+	// Bump the value by 1 using a bpf program.
+	want := uint32(1338)
+	a32, err := VariablePointer[atomic.Uint32](obj.Atomic)
+	qt.Assert(t, qt.IsNil(err))
+	a32.Store(want - 1)
+
+	mustReturn(t, obj.AddAtomic, 0)
+	qt.Assert(t, qt.Equals(a32.Load(), want))
+
+	_, err = VariablePointer[*uint32](obj.Atomic)
+	qt.Assert(t, qt.ErrorIs(err, ErrInvalidType))
+
+	_, err = VariablePointer[struct{ A, B *uint64 }](obj.Struct)
+	qt.Assert(t, qt.ErrorIs(err, ErrInvalidType))
+
+	type S struct {
+		_    structs.HostLayout
+		A, B uint64
+	}
+
+	s, err := VariablePointer[S](obj.Struct)
+	qt.Assert(t, qt.IsNil(err))
+	*s = S{A: 0xa, B: 0xb}
+	mustReturn(t, obj.CheckStruct, 1)
+
+	a, err := VariablePointer[[8192]byte](obj.Array)
+	qt.Assert(t, qt.IsNil(err))
+	a[len(a)-1] = 0xff
+	mustReturn(t, obj.CheckArray, 1)
+}
+
+func TestVariablePointerError(t *testing.T) {
+	file := testutils.NativeFile(t, "testdata/variables-%s.elf")
+	spec, err := LoadCollectionSpec(file)
+	qt.Assert(t, qt.IsNil(err))
+
+	obj := struct {
+		Atomic *Variable `ebpf:"var_atomic"`
+	}{}
+
+	qt.Assert(t, qt.IsNil(spec.LoadAndAssign(&obj, nil)))
+
+	_, err = VariablePointer[atomic.Uint32](obj.Atomic)
+	qt.Assert(t, qt.ErrorIs(err, ErrNotSupported))
+}
+
+func TestVariablePointerGC(t *testing.T) {
+	testutils.SkipIfNotSupported(t, haveMmapableMaps())
+
+	file := testutils.NativeFile(t, "testdata/variables-%s.elf")
+	spec, err := LoadCollectionSpec(file)
+	qt.Assert(t, qt.IsNil(err))
+
+	obj := struct {
+		AddAtomic *Program  `ebpf:"add_atomic"`
+		Atomic    *Variable `ebpf:"var_atomic"`
+	}{}
+	qt.Assert(t, qt.IsNil(spec.LoadAndAssign(&obj, &CollectionOptions{UnsafeVariableExperiment: true})))
+
+	// Pull out Program handle and Variable pointer so obj reference is dropped.
+	prog := obj.AddAtomic
+	t.Cleanup(func() {
+		prog.Close()
+	})
+
+	a32, err := VariablePointer[atomic.Uint32](obj.Atomic)
+	qt.Assert(t, qt.IsNil(err))
+
+	// No obj references past this point. Trigger multiple GC cycles to ensure
+	// obj is collected.
+	runtime.GC()
+	runtime.GC()
+	runtime.GC()
+
+	// Trigger prog and read memory multiple times to ensure reference is still
+	// valid.
+	mustReturn(t, prog, 0)
+	qt.Assert(t, qt.Equals(a32.Load(), 1))
+	mustReturn(t, prog, 0)
+	qt.Assert(t, qt.Equals(a32.Load(), 2))
+	mustReturn(t, prog, 0)
+	qt.Assert(t, qt.Equals(a32.Load(), 3))
+}