program: move jump and call fixup out of asm package

Jumps and calls may have to be fixed up, if their target isn't known
at compile time. This happens when a user uses the asm package to write
"inline" BPF programs, and when using function calls in C derived BPF.

The jump and call targets have to be expressed in units of raw BPF
instructions. Since some asm.Instructions have to be encoded in two
BPF instructions there is no 1:1 mapping. Figuring out the raw
instruction offset therefore involves iterating all instructions
and keeping a running count.

Due to this the fixup currently happens during instruction marshalling.
This is a bit of a hack, but so far allowed to keep the details of raw
instruction offsets contained to the asm package.

With BPF CO-RE this doesn't work too well anymore: the fixup we have to do
is also based on raw BPF instruction offsets. Cramming this into
instruction marshalling seems like the wrong thing to do, and is
probably not possible due to import cycles between the asm and btf
package.

Introduce a RawInstructionOffset type and an InstructionIterator to
encapsulate the raw offset handling, and move jump and call fixup to
a function in the main package.

Author: lmb

asm/instruction.go

@@ -12,6 +12,14 @@ import (
 // InstructionSize is the size of a BPF instruction in bytes
 const InstructionSize = 8
 
+// RawInstructionOffset is an offset in units of raw BPF instructions.
+type RawInstructionOffset uint64
+
+// Bytes returns the offset of an instruction in bytes.
+func (rio RawInstructionOffset) Bytes() uint64 {
+	return uint64(rio) * InstructionSize
+}
+
 // Instruction is a single eBPF instruction.
 type Instruction struct {
 	OpCode    OpCode
@@ -155,6 +163,13 @@ func (ins *Instruction) isLoadFromMap() bool {
 	return ins.OpCode == LoadImmOp(DWord) && (ins.Src == PseudoMapFD || ins.Src == PseudoMapValue)
 }
 
+// IsFunctionCall returns true if the instruction calls another BPF function.
+//
+// This is not the same thing as a BPF helper call.
+func (ins *Instruction) IsFunctionCall() bool {
+	return ins.OpCode.JumpOp() == Call && ins.Src == PseudoCall
+}
+
 // Format implements fmt.Formatter.
 func (ins Instruction) Format(f fmt.State, c rune) {
 	if c != 'v' {
@@ -310,28 +325,6 @@ func (insns Instructions) ReferenceOffsets() map[string][]int {
 	return offsets
 }
 
-func (insns Instructions) marshalledOffsets() (map[string]int, error) {
-	symbols := make(map[string]int)
-
-	marshalledPos := 0
-	for _, ins := range insns {
-		currentPos := marshalledPos
-		marshalledPos += ins.OpCode.marshalledInstructions()
-
-		if ins.Symbol == "" {
-			continue
-		}
-
-		if _, ok := symbols[ins.Symbol]; ok {
-			return nil, fmt.Errorf("duplicate symbol %s", ins.Symbol)
-		}
-
-		symbols[ins.Symbol] = currentPos
-	}
-
-	return symbols, nil
-}
-
 // Format implements fmt.Formatter.
 //
 // You can control indentation of symbols by
@@ -370,65 +363,67 @@ func (insns Instructions) Format(f fmt.State, c rune) {
 		symIndent = strings.Repeat(" ", symPadding)
 	}
 
-	// Figure out how many digits we need to represent the highest
-	// offset.
-	highestOffset := 0
-	for _, ins := range insns {
-		highestOffset += ins.OpCode.marshalledInstructions()
-	}
+	// Guess how many digits we need at most, by assuming that all instructions
+	// are double wide.
+	highestOffset := len(insns) * 2
 	offsetWidth := int(math.Ceil(math.Log10(float64(highestOffset))))
 
-	offset := 0
-	for _, ins := range insns {
-		if ins.Symbol != "" {
-			fmt.Fprintf(f, "%s%s:\n", symIndent, ins.Symbol)
+	iter := insns.Iterate()
+	for iter.Next() {
+		if iter.Ins.Symbol != "" {
+			fmt.Fprintf(f, "%s%s:\n", symIndent, iter.Ins.Symbol)
 		}
-		fmt.Fprintf(f, "%s%*d: %v\n", indent, offsetWidth, offset, ins)
-		offset += ins.OpCode.marshalledInstructions()
+		fmt.Fprintf(f, "%s%*d: %v\n", indent, offsetWidth, iter.Offset, iter.Ins)
 	}
 
 	return
 }
 
 // Marshal encodes a BPF program into the kernel format.
 func (insns Instructions) Marshal(w io.Writer, bo binary.ByteOrder) error {
-	absoluteOffsets, err := insns.marshalledOffsets()
-	if err != nil {
-		return err
-	}
-
-	num := 0
 	for i, ins := range insns {
-		switch {
-		case ins.OpCode.JumpOp() == Call && ins.Src == PseudoCall && ins.Constant == -1:
-			// Rewrite bpf to bpf call
-			offset, ok := absoluteOffsets[ins.Reference]
-			if !ok {
-				return fmt.Errorf("instruction %d: reference to missing symbol %s", i, ins.Reference)
-			}
-
-			ins.Constant = int64(offset - num - 1)
-
-		case ins.OpCode.Class() == JumpClass && ins.Offset == -1:
-			// Rewrite jump to label
-			offset, ok := absoluteOffsets[ins.Reference]
-			if !ok {
-				return fmt.Errorf("instruction %d: reference to missing symbol %s", i, ins.Reference)
-			}
-
-			ins.Offset = int16(offset - num - 1)
-		}
-
-		n, err := ins.Marshal(w, bo)
+		_, err := ins.Marshal(w, bo)
 		if err != nil {
 			return fmt.Errorf("instruction %d: %w", i, err)
 		}
-
-		num += int(n / InstructionSize)
 	}
 	return nil
 }
 
+// Iterate allows iterating a BPF program while keeping track of
+// various offsets.
+//
+// Modifying the instruction slice will lead to undefined behaviour.
+func (insns Instructions) Iterate() *InstructionIterator {
+	return &InstructionIterator{insns: insns}
+}
+
+// InstructionIterator iterates over a BPF program.
+type InstructionIterator struct {
+	insns Instructions
+	// The instruction in question.
+	Ins *Instruction
+	// The index of the instruction in the original instruction slice.
+	Index int
+	// The offset of the instruction in raw BPF instructions. This accounts
+	// for double-wide instructions.
+	Offset RawInstructionOffset
+}
+
+// Next returns true as long as there are any instructions remaining.
+func (iter *InstructionIterator) Next() bool {
+	if len(iter.insns) == 0 {
+		return false
+	}
+
+	if iter.Ins != nil {
+		iter.Offset += RawInstructionOffset(iter.Ins.OpCode.rawInstructions())
+	}
+	iter.Ins = &iter.insns[0]
+	iter.insns = iter.insns[1:]
+	return true
+}
+
 type bpfInstruction struct {
 	OpCode    OpCode
 	Registers bpfRegisters

asm/opcode.go

@@ -66,10 +66,10 @@ type OpCode uint8
 // InvalidOpCode is returned by setters on OpCode
 const InvalidOpCode OpCode = 0xff
 
-// marshalledInstructions returns the number of BPF instructions required
+// rawInstructions returns the number of BPF instructions required
 // to encode this opcode.
-func (op OpCode) marshalledInstructions() int {
-	if op == LoadImmOp(DWord) {
+func (op OpCode) rawInstructions() int {
+	if op.isDWordLoad() {
 		return 2
 	}
 	return 1

linker.go

@@ -84,3 +84,50 @@ func needSection(insns, section asm.Instructions) (bool, error) {
 	// None of the functions in the section are called.
 	return false, nil
 }
+
+func fixupJumpsAndCalls(insns asm.Instructions) error {
+	symbolOffsets := make(map[string]asm.RawInstructionOffset)
+	iter := insns.Iterate()
+	for iter.Next() {
+		ins := iter.Ins
+
+		if ins.Symbol == "" {
+			continue
+		}
+
+		if _, ok := symbolOffsets[ins.Symbol]; ok {
+			return fmt.Errorf("duplicate symbol %s", ins.Symbol)
+		}
+
+		symbolOffsets[ins.Symbol] = iter.Offset
+	}
+
+	iter = insns.Iterate()
+	for iter.Next() {
+		i := iter.Index
+		offset := iter.Offset
+		ins := iter.Ins
+
+		switch {
+		case ins.IsFunctionCall() && ins.Constant == -1:
+			// Rewrite bpf to bpf call
+			callOffset, ok := symbolOffsets[ins.Reference]
+			if !ok {
+				return fmt.Errorf("instruction %d: reference to missing symbol %s", i, ins.Reference)
+			}
+
+			ins.Constant = int64(callOffset - offset - 1)
+
+		case ins.OpCode.Class() == asm.JumpClass && ins.Offset == -1:
+			// Rewrite jump to label
+			jumpOffset, ok := symbolOffsets[ins.Reference]
+			if !ok {
+				return fmt.Errorf("instruction %d: reference to missing symbol %s", i, ins.Reference)
+			}
+
+			ins.Offset = int16(jumpOffset - offset - 1)
+		}
+	}
+
+	return nil
+}

prog.go

@@ -208,8 +208,15 @@ func convertProgramSpec(spec *ProgramSpec, handle *btf.Handle) (*bpfProgLoadAttr
 		return nil, fmt.Errorf("can't load %s program on %s", spec.ByteOrder, internal.NativeEndian)
 	}
 
+	insns := make(asm.Instructions, len(spec.Instructions))
+	copy(insns, spec.Instructions)
+
+	if err := fixupJumpsAndCalls(insns); err != nil {
+		return nil, err
+	}
+
 	buf := bytes.NewBuffer(make([]byte, 0, len(spec.Instructions)*asm.InstructionSize))
-	err := spec.Instructions.Marshal(buf, internal.NativeEndian)
+	err := insns.Marshal(buf, internal.NativeEndian)
 	if err != nil {
 		return nil, err
 	}