asm: support direct value access for maps

Linux added support for directly loading a value from a map at a known offset
in d8eca5bbb2be ("bpf: implement lookup-free direct value access for maps").
This allows to skip the call to bpf_map_lookup_elem and is a building block
for data sections.

Fix RewriteMapPtr to preserve any existing offset, and add a RewriteMapOffset
function.

Author: lmb

asm/instruction.go

@@ -103,22 +103,52 @@ func (ins Instruction) Marshal(w io.Writer, bo binary.ByteOrder) (uint64, error)
 
 // RewriteMapPtr changes an instruction to use a new map fd.
 //
-// Returns an error if the fd is invalid, or the instruction
-// is incorrect.
+// Returns an error if the instruction doesn't load a map.
 func (ins *Instruction) RewriteMapPtr(fd int) error {
 	if !ins.OpCode.isDWordLoad() {
 		return xerrors.Errorf("%s is not a 64 bit load", ins.OpCode)
 	}
 
-	if fd < 0 {
-		return xerrors.New("invalid fd")
+	if ins.Src != PseudoMapFD && ins.Src != PseudoMapValue {
+		return xerrors.New("not a load from a map")
 	}
 
-	ins.Src = R1
-	ins.Constant = int64(fd)
+	// Preserve the offset value for direct map loads.
+	offset := uint64(ins.Constant) & (math.MaxUint32 << 32)
+	rawFd := uint64(uint32(fd))
+	ins.Constant = int64(offset | rawFd)
 	return nil
 }
 
+func (ins *Instruction) mapPtr() uint32 {
+	return uint32(uint64(ins.Constant) & math.MaxUint32)
+}
+
+// RewriteMapOffset changes the offset of a direct load from a map.
+//
+// Returns an error if the instruction is not a direct load.
+func (ins *Instruction) RewriteMapOffset(offset uint32) error {
+	if !ins.OpCode.isDWordLoad() {
+		return xerrors.Errorf("%s is not a 64 bit load", ins.OpCode)
+	}
+
+	if ins.Src != PseudoMapValue {
+		return xerrors.New("not a direct load from a map")
+	}
+
+	fd := uint64(ins.Constant) & math.MaxUint32
+	ins.Constant = int64(uint64(offset)<<32 | fd)
+	return nil
+}
+
+func (ins *Instruction) mapOffset() uint32 {
+	return uint32(uint64(ins.Constant) >> 32)
+}
+
+func (ins *Instruction) isLoadFromMap() bool {
+	return ins.OpCode == LoadImmOp(DWord) && (ins.Src == PseudoMapFD || ins.Src == PseudoMapValue)
+}
+
 // Format implements fmt.Formatter.
 func (ins Instruction) Format(f fmt.State, c rune) {
 	if c != 'v' {
@@ -139,6 +169,19 @@ func (ins Instruction) Format(f fmt.State, c rune) {
 		return
 	}
 
+	if ins.isLoadFromMap() {
+		fd := int32(ins.mapPtr())
+		switch ins.Src {
+		case PseudoMapFD:
+			fmt.Fprintf(f, "LoadMapPtr dst: %s fd: %d", ins.Dst, fd)
+
+		case PseudoMapValue:
+			fmt.Fprintf(f, "LoadMapValue dst: %s, fd: %d off: %d", ins.Dst, fd, ins.mapOffset())
+		}
+
+		goto ref
+	}
+
 	fmt.Fprintf(f, "%v ", op)
 	switch cls := op.Class(); cls {
 	case LdClass, LdXClass, StClass, StXClass:
@@ -183,6 +226,7 @@ func (ins Instruction) Format(f fmt.State, c rune) {
 		}
 	}
 
+ref:
 	if ins.Reference != "" {
 		fmt.Fprintf(f, " <%s>", ins.Reference)
 	}

asm/instruction_test.go

@@ -59,18 +59,50 @@ func TestSignedJump(t *testing.T) {
 	}
 }
 
-func TestInstructionRewriteMapPtr(t *testing.T) {
+func TestInstructionRewriteMapConstant(t *testing.T) {
 	ins := LoadMapPtr(R2, 0)
+	ins.Src = R2
+
 	if err := ins.RewriteMapPtr(1); err != nil {
-		t.Fatal("Can't rewrite map pointer")
+		t.Fatal("Can't rewrite map pointer:", err)
+	}
+	if ins.mapPtr() != 1 {
+		t.Fatal("Expected map ptr to be 1, got", ins.Constant)
+	}
+
+	if err := ins.RewriteMapOffset(123); err != nil {
+		t.Fatal("Can't rewrite map offset:", err)
+	}
+	if fd := ins.mapPtr(); fd != 1 {
+		t.Fatal("Expected map ptr to be 1 after changing the offset, got", fd)
+	}
+
+	if err := ins.RewriteMapPtr(2); err != nil {
+		t.Fatal("Can't rewrite map pointer:", err)
 	}
-	if ins.Constant != 1 {
-		t.Error("Expected Constant to be 1, got", ins.Constant)
+	if off := ins.mapOffset(); off != 123 {
+		t.Fatal("Expected map offset to be 123 after changin the pointer, got", off)
 	}
 
 	ins = Mov.Imm(R1, 32)
 	if err := ins.RewriteMapPtr(1); err == nil {
-		t.Error("Allows rewriting bogus instruction")
+		t.Error("RewriteMapPtr rewriting bogus instruction")
+	}
+	if err := ins.RewriteMapOffset(1); err == nil {
+		t.Error("RewriteMapOffset rewriting bogus instruction")
+	}
+}
+
+func TestInstructionLoadMapValue(t *testing.T) {
+	ins := LoadMapValue(R0, 1, 123)
+	if !ins.isLoadFromMap() {
+		t.Error("isLoadFromMap returns false")
+	}
+	if fd := ins.mapPtr(); fd != 1 {
+		t.Error("Expected map fd to be 1, got", fd)
+	}
+	if off := ins.mapOffset(); off != 123 {
+		t.Fatal("Expected map offset to be 123 after changin the pointer, got", off)
 	}
 }
 

asm/load_store.go

@@ -110,11 +110,26 @@ func LoadMapPtr(dst Register, fd int) Instruction {
 	return Instruction{
 		OpCode:   LoadImmOp(DWord),
 		Dst:      dst,
-		Src:      R1,
+		Src:      PseudoMapFD,
 		Constant: int64(fd),
 	}
 }
 
+// LoadMapValue stores a pointer to the value at a certain offset of a map.
+func LoadMapValue(dst Register, fd int, offset uint32) Instruction {
+	if fd < 0 {
+		return Instruction{OpCode: InvalidOpCode}
+	}
+
+	fdAndOffset := (uint64(offset) << 32) | uint64(uint32(fd))
+	return Instruction{
+		OpCode:   LoadImmOp(DWord),
+		Dst:      dst,
+		Src:      PseudoMapValue,
+		Constant: int64(fdAndOffset),
+	}
+}
+
 // LoadIndOp returns the OpCode for loading a value of given size from an sk_buff.
 func LoadIndOp(size Size) OpCode {
 	return OpCode(LdClass).SetMode(IndMode).SetSize(size)

asm/register.go

@@ -33,6 +33,12 @@ const (
 	RFP          = R10
 )
 
+// Pseudo registers used by 64bit loads
+const (
+	PseudoMapFD    = R1 // BPF_PSEUDO_MAP_FD
+	PseudoMapValue = R2 // BPF_PSEUDO_MAP_VALUE
+)
+
 func (r Register) String() string {
 	v := uint8(r)
 	if v == 10 {

collection.go

@@ -1,7 +1,10 @@
 package ebpf
 
 import (
+	"math"
+
 	"github.com/cilium/ebpf/asm"
+	"github.com/cilium/ebpf/internal"
 	"github.com/cilium/ebpf/internal/btf"
 	"golang.org/x/xerrors"
 )
@@ -153,21 +156,27 @@ func NewCollectionWithOptions(spec *CollectionSpec, opts CollectionOptions) (col
 
 		// Rewrite any reference to a valid map.
 		for i := range progSpec.Instructions {
-			var (
-				ins = &progSpec.Instructions[i]
-				m   = maps[ins.Reference]
-			)
+			ins := &progSpec.Instructions[i]
 
-			if ins.Reference == "" || m == nil {
+			if ins.OpCode != asm.LoadImmOp(asm.DWord) || ins.Reference == "" {
 				continue
 			}
 
-			if ins.Src == asm.R1 {
+			if uint32(ins.Constant) != math.MaxUint32 {
 				// Don't overwrite maps already rewritten, users can
 				// rewrite programs in the spec themselves
 				continue
 			}
 
+			m := maps[ins.Reference]
+			if m == nil {
+				return nil, xerrors.Errorf("program %s: missing map %s", progName, ins.Reference)
+			}
+
+			fd := m.FD()
+			if fd < 0 {
+				return nil, xerrors.Errorf("map %s: %w", ins.Reference, internal.ErrClosedFd)
+			}
 			if err := ins.RewriteMapPtr(m.FD()); err != nil {
 				return nil, xerrors.Errorf("progam %s: map %s: %w", progName, ins.Reference, err)
 			}

elf_reader.go

@@ -231,6 +231,12 @@ func (ec *elfCode) loadInstructions(section *elf.Section, symbols, relocations m
 
 		ins.Symbol = symbols[offset]
 		ins.Reference = relocations[offset]
+		if ins.OpCode == asm.LoadImmOp(asm.DWord) && ins.Reference != "" {
+			ins.Src = asm.R1
+			if err := ins.RewriteMapPtr(-1); err != nil {
+				return nil, 0, err
+			}
+		}
 
 		insns = append(insns, ins)
 		offset += n