bpf: support JumpIf on RegX instead of K

Add a JumpIfX instruction which implements conditional jumps using
RegA and RegX. This is in addition to the pre-existing JumpIf
instruction which uses RegA and K.

This instruction / addressing mode is not mentionned in the original BPF
paper, but is supported by tools like bpf_asm, and has recently been
added to the kernel's filter.txt.

Simplify some of the parsing logic, and add a separate helper for
checking for "fake" JumpIfs.

Add JumpIfX support to the BPF vm.

Update testdata with JumpIfX instructions, and add tests
for both the assembler/disassembler and vm.

Fixes golang/go#27814

Change-Id: I0c3f6ac7eb5b4cd4d9c5af8784ee2e8d25195a0a
GitHub-Last-Rev: 39a8816
GitHub-Pull-Request: #20
Reviewed-on: https://go-review.googlesource.com/c/136895
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

Author: arthurfabre

bpf/constants.go

@@ -38,6 +38,7 @@ const (
 type JumpTest uint16
 
 // Supported operators for conditional jumps.
+// K can be RegX for JumpIfX
 const (
 	// K == A
 	JumpEqual JumpTest = iota
@@ -134,12 +135,9 @@ const (
 	opMaskLoadDest  = 0x01
 	opMaskLoadWidth = 0x18
 	opMaskLoadMode  = 0xe0
-	// opClsALU
-	opMaskOperandSrc = 0x08
-	opMaskOperator   = 0xf0
-	// opClsJump
-	opMaskJumpConst = 0x0f
-	opMaskJumpCond  = 0xf0
+	// opClsALU & opClsJump
+	opMaskOperand  = 0x08
+	opMaskOperator = 0xf0
 )
 
 const (
@@ -192,15 +190,21 @@ const (
 	opLoadWidth1
 )
 
-// Operator defined by ALUOp*
+// Operand for ALU and Jump instructions
+type opOperand uint16
 
+// Supported operand sources.
 const (
-	opALUSrcConstant uint16 = iota << 3
-	opALUSrcX
+	opOperandConstant opOperand = iota << 3
+	opOperandX
 )
 
+// An jumpOp is a conditional jump condition.
+type jumpOp uint16
+
+// Supported jump conditions.
 const (
-	opJumpAlways = iota << 4
+	opJumpAlways jumpOp = iota << 4
 	opJumpEqual
 	opJumpGT
 	opJumpGE

bpf/instructions.go

@@ -89,74 +89,33 @@ func (ri RawInstruction) Disassemble() Instruction {
 	case opClsALU:
 		switch op := ALUOp(ri.Op & opMaskOperator); op {
 		case ALUOpAdd, ALUOpSub, ALUOpMul, ALUOpDiv, ALUOpOr, ALUOpAnd, ALUOpShiftLeft, ALUOpShiftRight, ALUOpMod, ALUOpXor:
-			if ri.Op&opMaskOperandSrc != 0 {
+			switch operand := opOperand(ri.Op & opMaskOperand); operand {
+			case opOperandX:
 				return ALUOpX{Op: op}
+			case opOperandConstant:
+				return ALUOpConstant{Op: op, Val: ri.K}
+			default:
+				return ri
 			}
-			return ALUOpConstant{Op: op, Val: ri.K}
 		case aluOpNeg:
 			return NegateA{}
 		default:
 			return ri
 		}
 
 	case opClsJump:
-		if ri.Op&opMaskJumpConst != opClsJump {
-			return ri
-		}
-		switch ri.Op & opMaskJumpCond {
+		switch op := jumpOp(ri.Op & opMaskOperator); op {
 		case opJumpAlways:
 			return Jump{Skip: ri.K}
-		case opJumpEqual:
-			if ri.Jt == 0 {
-				return JumpIf{
-					Cond:      JumpNotEqual,
-					Val:       ri.K,
-					SkipTrue:  ri.Jf,
-					SkipFalse: 0,
-				}
-			}
-			return JumpIf{
-				Cond:      JumpEqual,
-				Val:       ri.K,
-				SkipTrue:  ri.Jt,
-				SkipFalse: ri.Jf,
-			}
-		case opJumpGT:
-			if ri.Jt == 0 {
-				return JumpIf{
-					Cond:      JumpLessOrEqual,
-					Val:       ri.K,
-					SkipTrue:  ri.Jf,
-					SkipFalse: 0,
-				}
-			}
-			return JumpIf{
-				Cond:      JumpGreaterThan,
-				Val:       ri.K,
-				SkipTrue:  ri.Jt,
-				SkipFalse: ri.Jf,
-			}
-		case opJumpGE:
-			if ri.Jt == 0 {
-				return JumpIf{
-					Cond:      JumpLessThan,
-					Val:       ri.K,
-					SkipTrue:  ri.Jf,
-					SkipFalse: 0,
-				}
-			}
-			return JumpIf{
-				Cond:      JumpGreaterOrEqual,
-				Val:       ri.K,
-				SkipTrue:  ri.Jt,
-				SkipFalse: ri.Jf,
-			}
-		case opJumpSet:
-			return JumpIf{
-				Cond:      JumpBitsSet,
-				Val:       ri.K,
-				SkipTrue:  ri.Jt,
-				SkipFalse: ri.Jf,
+		case opJumpEqual, opJumpGT, opJumpGE, opJumpSet:
+			cond, skipTrue, skipFalse := jumpOpToTest(op, ri.Jt, ri.Jf)
+			switch operand := opOperand(ri.Op & opMaskOperand); operand {
+			case opOperandX:
+				return JumpIfX{Cond: cond, SkipTrue: skipTrue, SkipFalse: skipFalse}
+			case opOperandConstant:
+				return JumpIf{Cond: cond, Val: ri.K, SkipTrue: skipTrue, SkipFalse: skipFalse}
+			default:
+				return ri
 			}
 		default:
 			return ri
@@ -187,6 +146,41 @@ func (ri RawInstruction) Disassemble() Instruction {
 	}
 }
 
+func jumpOpToTest(op jumpOp, skipTrue uint8, skipFalse uint8) (JumpTest, uint8, uint8) {
+	var test JumpTest
+
+	// Decode "fake" jump conditions that don't appear in machine code
+	// Ensures the Assemble -> Disassemble stage recreates the same instructions
+	// See https://github.com/golang/go/issues/18470
+	if skipTrue == 0 {
+		switch op {
+		case opJumpEqual:
+			test = JumpNotEqual
+		case opJumpGT:
+			test = JumpLessOrEqual
+		case opJumpGE:
+			test = JumpLessThan
+		case opJumpSet:
+			test = JumpBitsNotSet
+		}
+
+		return test, skipFalse, 0
+	}
+
+	switch op {
+	case opJumpEqual:
+		test = JumpEqual
+	case opJumpGT:
+		test = JumpGreaterThan
+	case opJumpGE:
+		test = JumpGreaterOrEqual
+	case opJumpSet:
+		test = JumpBitsSet
+	}
+
+	return test, skipTrue, skipFalse
+}
+
 // LoadConstant loads Val into register Dst.
 type LoadConstant struct {
 	Dst Register
@@ -413,7 +407,7 @@ type ALUOpConstant struct {
 // Assemble implements the Instruction Assemble method.
 func (a ALUOpConstant) Assemble() (RawInstruction, error) {
 	return RawInstruction{
-		Op: opClsALU | opALUSrcConstant | uint16(a.Op),
+		Op: opClsALU | uint16(opOperandConstant) | uint16(a.Op),
 		K:  a.Val,
 	}, nil
 }
@@ -454,7 +448,7 @@ type ALUOpX struct {
 // Assemble implements the Instruction Assemble method.
 func (a ALUOpX) Assemble() (RawInstruction, error) {
 	return RawInstruction{
-		Op: opClsALU | opALUSrcX | uint16(a.Op),
+		Op: opClsALU | uint16(opOperandX) | uint16(a.Op),
 	}, nil
 }
 
@@ -509,7 +503,7 @@ type Jump struct {
 // Assemble implements the Instruction Assemble method.
 func (a Jump) Assemble() (RawInstruction, error) {
 	return RawInstruction{
-		Op: opClsJump | opJumpAlways,
+		Op: opClsJump | uint16(opJumpAlways),
 		K:  a.Skip,
 	}, nil
 }
@@ -530,11 +524,39 @@ type JumpIf struct {
 
 // Assemble implements the Instruction Assemble method.
 func (a JumpIf) Assemble() (RawInstruction, error) {
+	return jumpToRaw(a.Cond, opOperandConstant, a.Val, a.SkipTrue, a.SkipFalse)
+}
+
+// String returns the instruction in assembler notation.
+func (a JumpIf) String() string {
+	return jumpToString(a.Cond, fmt.Sprintf("#%d", a.Val), a.SkipTrue, a.SkipFalse)
+}
+
+// JumpIfX skips the following Skip instructions in the program if A
+// <Cond> X is true.
+type JumpIfX struct {
+	Cond      JumpTest
+	SkipTrue  uint8
+	SkipFalse uint8
+}
+
+// Assemble implements the Instruction Assemble method.
+func (a JumpIfX) Assemble() (RawInstruction, error) {
+	return jumpToRaw(a.Cond, opOperandX, 0, a.SkipTrue, a.SkipFalse)
+}
+
+// String returns the instruction in assembler notation.
+func (a JumpIfX) String() string {
+	return jumpToString(a.Cond, "x", a.SkipTrue, a.SkipFalse)
+}
+
+// jumpToRaw assembles a jump instruction into a RawInstruction
+func jumpToRaw(test JumpTest, operand opOperand, k uint32, skipTrue, skipFalse uint8) (RawInstruction, error) {
 	var (
-		cond uint16
+		cond jumpOp
 		flip bool
 	)
-	switch a.Cond {
+	switch test {
 	case JumpEqual:
 		cond = opJumpEqual
 	case JumpNotEqual:
@@ -552,63 +574,63 @@ func (a JumpIf) Assemble() (RawInstruction, error) {
 	case JumpBitsNotSet:
 		cond, flip = opJumpSet, true
 	default:
-		return RawInstruction{}, fmt.Errorf("unknown JumpTest %v", a.Cond)
+		return RawInstruction{}, fmt.Errorf("unknown JumpTest %v", test)
 	}
-	jt, jf := a.SkipTrue, a.SkipFalse
+	jt, jf := skipTrue, skipFalse
 	if flip {
 		jt, jf = jf, jt
 	}
 	return RawInstruction{
-		Op: opClsJump | cond,
+		Op: opClsJump | uint16(cond) | uint16(operand),
 		Jt: jt,
 		Jf: jf,
-		K:  a.Val,
+		K:  k,
 	}, nil
 }
 
-// String returns the instruction in assembler notation.
-func (a JumpIf) String() string {
-	switch a.Cond {
+// jumpToString converts a jump instruction to assembler notation
+func jumpToString(cond JumpTest, operand string, skipTrue, skipFalse uint8) string {
+	switch cond {
 	// K == A
 	case JumpEqual:
-		return conditionalJump(a, "jeq", "jneq")
+		return conditionalJump(operand, skipTrue, skipFalse, "jeq", "jneq")
 	// K != A
 	case JumpNotEqual:
-		return fmt.Sprintf("jneq #%d,%d", a.Val, a.SkipTrue)
+		return fmt.Sprintf("jneq %s,%d", operand, skipTrue)
 	// K > A
 	case JumpGreaterThan:
-		return conditionalJump(a, "jgt", "jle")
+		return conditionalJump(operand, skipTrue, skipFalse, "jgt", "jle")
 	// K < A
 	case JumpLessThan:
-		return fmt.Sprintf("jlt #%d,%d", a.Val, a.SkipTrue)
+		return fmt.Sprintf("jlt %s,%d", operand, skipTrue)
 	// K >= A
 	case JumpGreaterOrEqual:
-		return conditionalJump(a, "jge", "jlt")
+		return conditionalJump(operand, skipTrue, skipFalse, "jge", "jlt")
 	// K <= A
 	case JumpLessOrEqual:
-		return fmt.Sprintf("jle #%d,%d", a.Val, a.SkipTrue)
+		return fmt.Sprintf("jle %s,%d", operand, skipTrue)
 	// K & A != 0
 	case JumpBitsSet:
-		if a.SkipFalse > 0 {
-			return fmt.Sprintf("jset #%d,%d,%d", a.Val, a.SkipTrue, a.SkipFalse)
+		if skipFalse > 0 {
+			return fmt.Sprintf("jset %s,%d,%d", operand, skipTrue, skipFalse)
 		}
-		return fmt.Sprintf("jset #%d,%d", a.Val, a.SkipTrue)
+		return fmt.Sprintf("jset %s,%d", operand, skipTrue)
 	// K & A == 0, there is no assembler instruction for JumpBitNotSet, use JumpBitSet and invert skips
 	case JumpBitsNotSet:
-		return JumpIf{Cond: JumpBitsSet, SkipTrue: a.SkipFalse, SkipFalse: a.SkipTrue, Val: a.Val}.String()
+		return jumpToString(JumpBitsSet, operand, skipFalse, skipTrue)
 	default:
-		return fmt.Sprintf("unknown instruction: %#v", a)
+		return fmt.Sprintf("unknown JumpTest %#v", cond)
 	}
 }
 
-func conditionalJump(inst JumpIf, positiveJump, negativeJump string) string {
-	if inst.SkipTrue > 0 {
-		if inst.SkipFalse > 0 {
-			return fmt.Sprintf("%s #%d,%d,%d", positiveJump, inst.Val, inst.SkipTrue, inst.SkipFalse)
+func conditionalJump(operand string, skipTrue, skipFalse uint8, positiveJump, negativeJump string) string {
+	if skipTrue > 0 {
+		if skipFalse > 0 {
+			return fmt.Sprintf("%s %s,%d,%d", positiveJump, operand, skipTrue, skipFalse)
 		}
-		return fmt.Sprintf("%s #%d,%d", positiveJump, inst.Val, inst.SkipTrue)
+		return fmt.Sprintf("%s %s,%d", positiveJump, operand, skipTrue)
 	}
-	return fmt.Sprintf("%s #%d,%d", negativeJump, inst.Val, inst.SkipFalse)
+	return fmt.Sprintf("%s %s,%d", negativeJump, operand, skipFalse)
 }
 
 // RetA exits the BPF program, returning the value of register A.