libfuzzer: Address review

Author: kyakdan

src/cmd/compile/internal/typecheck/builtin.go

@@ -265,8 +265,11 @@ func libfuzzerTraceConstCmp1(uint8, uint8, int)
 func libfuzzerTraceConstCmp2(uint16, uint16, int)
 func libfuzzerTraceConstCmp4(uint32, uint32, int)
 func libfuzzerTraceConstCmp8(uint64, uint64, int)
-func libfuzzerHookStrCmp(string, string, int, int)
+func libfuzzerHookStrCmp(string, string, bool, int)
+func libfuzzerHookEqualFold(string, string, int)
 func libfuzzerIncrementCounter(*uint8)
+
+// This function should be called by the fuzz target on start to register the 8bit counters with libfuzzer
 func LibfuzzerInitializeCounters()
 
 // architecture variants

src/cmd/compile/internal/typecheck/builtin/runtime.go

@@ -5,8 +5,9 @@
 package walk
 
 import (
+	"crypto/sha1"
 	"go/constant"
-	"math/rand"
+	"strconv"
 
 	"cmd/compile/internal/base"
 	"cmd/compile/internal/ir"
@@ -16,8 +17,15 @@ import (
 	"cmd/compile/internal/types"
 )
 
-func fakePC() ir.Node {
-	return ir.NewInt(int64(rand.Uint32()))
+func fakePC(n ir.Node) ir.Node {
+	// In order to get deterministic IDs, we include the package path, file index, line number, column number
+	// in the calculation of the fakePC for the IR node
+	data := base.Ctxt.Pkgpath
+	data += strconv.FormatInt(int64(n.Pos().FileIndex()), 10)
+	data += n.Pos().LineNumber()
+	data += strconv.FormatUint(uint64(n.Pos().Col()), 10)
+	hash := sha1.Sum([]byte(data))
+	return ir.NewInt(int64(uint32(hash[0]) | uint32(hash[1])<<8 | uint32(hash[2])<<16 | uint32(hash[3])<<24))
 }
 
 // The result of walkCompare MUST be assigned back to n, e.g.
@@ -136,7 +144,8 @@ func walkCompare(n *ir.BinaryExpr, init *ir.Nodes) ir.Node {
 			default:
 				base.Fatalf("unexpected integer size %d for %v", t.Size(), t)
 			}
-			init.Append(mkcall(fn, nil, init, tracecmpArg(l, paramType, init), tracecmpArg(r, paramType, init), fakePC()))
+
+			init.Append(mkcall(fn, nil, init, tracecmpArg(l, paramType, init), tracecmpArg(r, paramType, init), fakePC(n)))
 		}
 		return n
 	case types.TARRAY:
@@ -286,14 +295,6 @@ func walkCompareInterface(n *ir.BinaryExpr, init *ir.Nodes) ir.Node {
 }
 
 func walkCompareString(n *ir.BinaryExpr, init *ir.Nodes) ir.Node {
-	if base.Debug.Libfuzzer != 0 {
-		fn := "libfuzzerHookStrCmp"
-		l := cheapExpr(n.X, init)
-		r := cheapExpr(n.Y, init)
-		paramType := types.Types[types.TSTRING]
-		init.Append(mkcall(fn, nil, init, tracecmpArg(l, paramType, init), tracecmpArg(r, paramType, init), ir.NewInt(1), fakePC()))
-	}
-
 	// Rewrite comparisons to short constant strings as length+byte-wise comparisons.
 	var cs, ncs ir.Node // const string, non-const string
 	switch {
@@ -412,8 +413,16 @@ func walkCompareString(n *ir.BinaryExpr, init *ir.Nodes) ir.Node {
 		r = mkcall("cmpstring", types.Types[types.TINT], init, typecheck.Conv(n.X, types.Types[types.TSTRING]), typecheck.Conv(n.Y, types.Types[types.TSTRING]))
 		r = ir.NewBinaryExpr(base.Pos, n.Op(), r, ir.NewInt(0))
 	}
-
-	return finishCompare(n, r, init)
+	result := finishCompare(n, r, init)
+	if base.Debug.Libfuzzer != 0 {
+		fn := "libfuzzerHookStrCmp"
+		x := cheapExpr(n.X, init)
+		y := cheapExpr(n.Y, init)
+		z := cheapExpr(result, init)
+		paramType := types.Types[types.TSTRING]
+		init.Append(mkcall(fn, nil, init, tracecmpArg(x, paramType, init), tracecmpArg(y, paramType, init), tracecmpArg(z, n.Type(), init), fakePC(n)))
+	}
+	return result
 }
 
 // The result of finishCompare MUST be assigned back to n, e.g.

src/cmd/compile/internal/walk/compare.go

@@ -495,6 +495,16 @@ func walkAddString(n *ir.AddStringExpr, init *ir.Nodes) ir.Node {
 	return r1
 }
 
+type hookInfo struct {
+	paramType   types.Kind
+	argsNum     int
+	runtimeFunc string
+}
+
+var hooks = map[string]hookInfo{
+	"strings.EqualFold": {paramType: types.TSTRING, argsNum: 2, runtimeFunc: "libfuzzerHookEqualFold"},
+}
+
 // walkCall walks an OCALLFUNC or OCALLINTER node.
 func walkCall(n *ir.CallExpr, init *ir.Nodes) ir.Node {
 	if n.Op() == ir.OCALLMETH {
@@ -590,13 +600,15 @@ func walkCall1(n *ir.CallExpr, init *ir.Nodes) {
 	}
 
 	n.Args = args
-	if base.Debug.Libfuzzer != 0 && n.X.Sym() != nil {
-		switch n.X.Sym().Pkg.Path {
-		case "strings":
-			if n.X.Sym().Name == "EqualFold" && len(args) == 2 {
-				paramType := types.Types[types.TSTRING]
-				init.Append(mkcall("libfuzzerHookStrCmp", nil, init, tracecmpArg(args[0], paramType, init), tracecmpArg(args[1], paramType, init), ir.NewInt(1), fakePC()))
+	funSym := n.X.Sym()
+	if base.Debug.Libfuzzer != 0 && funSym != nil {
+		if hook, found := hooks[funSym.Pkg.Path+"."+funSym.Name]; found && len(args) == hook.argsNum {
+			var hookArgs []ir.Node
+			for _, arg := range args {
+				hookArgs = append(hookArgs, tracecmpArg(arg, types.Types[hook.paramType], init))
 			}
+			hookArgs = append(hookArgs, fakePC(n))
+			init.Append(mkcall(hook.runtimeFunc, nil, init, hookArgs...))
 		}
 	}
 }

src/cmd/compile/internal/walk/expr.go

@@ -32,6 +32,7 @@ func libfuzzerTraceConstCmp2(arg0, arg1 uint16, fakePC int) {}
 func libfuzzerTraceConstCmp4(arg0, arg1 uint32, fakePC int) {}
 func libfuzzerTraceConstCmp8(arg0, arg1 uint64, fakePC int) {}
 
-func libfuzzerHookStrCmp(arg0, arg1 string, result, fakePC int) {}
+func libfuzzerHookStrCmp(arg0, arg1 string, result bool, fakePC int) {}
+func libfuzzerHookEqualFold(arg0, arg1 string, fakePC int)           {}
 
 func libfuzzerIncrementCounter(counter *uint8) {}

src/internal/fuzz/trace.go

@@ -8,6 +8,7 @@ package runtime
 
 import "unsafe" // for go:linkname
 
+// Keep in sync with the definition of ret_sled in src/runtime/libfuzzer_amd64.s
 const retSledSize = 512
 
 func libfuzzerCallTraceIntCmp(fn *byte, arg0, arg1, fakePC uintptr)
@@ -54,31 +55,35 @@ func libfuzzerTraceConstCmp8(arg0, arg1 uint64, fakePC int) {
 	libfuzzerCallTraceIntCmp(&__sanitizer_cov_trace_const_cmp8, uintptr(arg0), uintptr(arg1), uintptr(fakePC))
 }
 
-func libfuzzerHookStrCmp(s1, s2 string, result, fakePC int) {
-	libfuzzerCallHookStrCmp(&__sanitizer_weak_hook_strcmp, uintptr(fakePC), cstring(s1), cstring(s2), uintptr(result))
+func libfuzzerHookStrCmp(s1, s2 string, equal bool, fakePC int) {
+	if !equal {
+		libfuzzerCallHookStrCmp(&__sanitizer_weak_hook_strcmp, uintptr(fakePC), cstring(s1), cstring(s2), uintptr(1))
+	}
+}
+
+func libfuzzerHookEqualFold(s1, s2 string, fakePC int) {
+	libfuzzerCallHookStrCmp(&__sanitizer_weak_hook_strcmp, uintptr(fakePC), cstring(s1), cstring(s2), uintptr(1))
 }
 
 var pcTables []byte
 
 func LibfuzzerInitializeCounters() {
 	libfuzzerCallTraceInit(&__sanitizer_cov_8bit_counters_init, &__start___sancov_cntrs, &__stop___sancov_cntrs)
-	var offset uintptr = 0
 	start := unsafe.Pointer(&__start___sancov_cntrs)
 	end := unsafe.Pointer(&__stop___sancov_cntrs)
 
-	cur := start
-	for cur != end {
-		offset++
-		cur = unsafe.Pointer(uintptr(start) + offset)
-	}
-
-	size := (offset + 1) * unsafe.Sizeof(uintptr(0)) * 2
+	// PC tables are arrays of ptr-sized integers representing pairs [PC,PCFlags] for every instrumented block.
+	// The number of PCs and PCFlags is the same as the number of 8-bit counters. Each PC table entry has
+	// the size of two ptr-sized integers. We allocate one more byte that what we actually need so that we can
+	// get a pointer representing the end of the PC table array.
+	size := (uintptr(end)-uintptr(start))*unsafe.Sizeof(uintptr(0))*2 + 1
 	pcTables = make([]byte, size)
 	libfuzzerCallTraceInit(&__sanitizer_cov_pcs_init, &pcTables[0], &pcTables[size-1])
 }
 
-// libfuzzerIncrementCounter guarantees that the counter never becomes zero again once
-// it is incremented. It implements the NeverZero optimization presented by the paper:
+// libfuzzerIncrementCounter guarantees that the counter never becomes zero
+// again once it has been incremented once. It implements the NeverZero
+// optimization presented by the paper:
 // "AFL++: Combining Incremental Steps of Fuzzing Research"
 func libfuzzerIncrementCounter(counter *uint8) {
 	if *counter == 0xff {

src/runtime/libfuzzer.go

@@ -28,7 +28,7 @@ TEXT	runtime·libfuzzerCallHookStrCmp(SB), NOSPLIT, $0-40
 	MOVQ	fn+0(FP), AX
 	MOVQ	hookId+8(FP), RARG0
 	MOVQ	s1+16(FP), RARG1
-    MOVQ	s2+24(FP), RARG2
+	MOVQ	s2+24(FP), RARG2
 	MOVQ	result+32(FP), RARG3
 
 	get_tls(R12)
@@ -116,13 +116,24 @@ TEXT	runtime·libfuzzerCallTraceIntCmp(SB), NOSPLIT, $0-32
 	MOVQ	(g_sched+gobuf_sp)(R10), SP
 call:
 	ANDQ	$~15, SP	// alignment for gcc ABI
+	// Load the address of the end of the function and push it into the stack.
+	// This address will be jumped to after executing the return instruction
+	// from the return sled. There we reset the stack pointer and return.
 	MOVQ    $end_of_function(SB), BX
 	PUSHQ   BX
+	// Load the starting address of the return sled into BX.
 	MOVQ    $ret_sled(SB), BX
+	// Load the address of the i'th return instruction fron the return sled.
+	// The index is given in the fakePC argument.
 	ADDQ    RARG2, BX
 	PUSHQ   BX
+	// Call the original function with the fakePC return address on the stack.
+	// Function arguments arg0 and arg1 are passed unchanged in the registers
+	// RDI and RSI as specified by the x64 calling convention.
 	JMP     AX
-non_reachable:
+// This code will not be executed and is only there to statisfy assembler
+// check of a balanced stack.
+not_reachable:
 	POPQ    BX
 	POPQ    BX
 	RET