pkg/kcov: cleanup and support for 32-bit arches

Use uintptr instead of uint64 to allow for non-64-bit arch support, and
do some small cleanup and suggested in code review.

Author: Ethan Graham

pkg/kcov/cdefs.go

@@ -3,11 +3,8 @@ package kcov
 import "unsafe"
 
 const (
-	// sizeof(uintptr).
-	sizeofUintPtr = 8
-)
+	sizeofUintPtr = int(unsafe.Sizeof((*int)(nil)))
 
-const (
 	_IOC_NRBITS   = 8
 	_IOC_TYPEBITS = 8
 	_IOC_SIZEBITS = 14

pkg/kcov/kcov.go

@@ -21,26 +21,27 @@ type KCOVState struct {
 	cover []byte
 }
 
-// Trace invokes `f` and returns its result, as well as collected kernel
-// coverage during its invocation.
-func (st *KCOVState) Trace(f func() error) ([]uintptr, error) {
+type KCOVTraceResult struct {
+	Result   error     // Result of the call.
+	Coverage []uintptr // Collected program counters.
+}
+
+// Trace invokes `f` and returns a KCOVTraceResult.
+func (st *KCOVState) Trace(f func() error) KCOVTraceResult {
 	// First 8 bytes holds the number of collected PCs since last poll.
-	countPtr := (*uint64)(unsafe.Pointer(&st.cover[0]))
+	countPtr := (*uintptr)(unsafe.Pointer(&st.cover[0]))
 	// Reset coverage for this run.
-	atomic.StoreUint64(countPtr, 0)
+	atomic.StoreUintptr(countPtr, 0)
 	// Trigger call.
 	err := f()
 	// Load the number of PCs that were hit during trigger.
-	n := atomic.LoadUint64(countPtr)
-	if n == 0 {
-		return nil, nil
-	}
+	n := atomic.LoadUintptr(countPtr)
 
 	pcDataPtr := (*uintptr)(unsafe.Pointer(&st.cover[sizeofUintPtr]))
 	pcs := unsafe.Slice(pcDataPtr, n)
 	pcsCopy := make([]uintptr, n)
 	copy(pcsCopy, pcs)
-	return pcsCopy, err
+	return KCOVTraceResult{Result: err, Coverage: pcsCopy}
 }
 
 // EnableTracingForCurrentGoroutine prepares the current goroutine for kcov tracing.
@@ -53,47 +54,42 @@ func EnableTracingForCurrentGoroutine() (*KCOVState, error) {
 	if err != nil {
 		return nil, err
 	}
-
-	cleanupOnError := func() {
-		file.Close()
-		runtime.UnlockOSThread()
+	st := KCOVState{
+		file: file,
 	}
-
 	fd := file.Fd()
 
 	// Setup trace mode and size.
 	if err := unix.IoctlSetInt(int(fd), uint(KCOV_INIT_TRACE), kcovCoverSize); err != nil {
-		cleanupOnError()
+		st.DisableTracing()
 		return nil, err
 	}
 
-	// Mmap buffer shared between kernel- and user-space.
-	coverageBuffer, err := unix.Mmap(
+	// Mmap buffer shared between kernel- and user-space. For more information,
+	// see the Linux KCOV documentation: https://docs.kernel.org/dev-tools/kcov.html.
+	st.cover, err = unix.Mmap(
 		int(fd),
-		0, // offset
+		0, // Offset.
 		kcovCoverSize*sizeofUintPtr,
-		unix.PROT_READ|unix.PROT_WRITE, // a read/write mapping
-		unix.MAP_SHARED,                // changes are shared with the kernel
+		unix.PROT_READ|unix.PROT_WRITE,
+		unix.MAP_SHARED,
 	)
 	if err != nil {
-		cleanupOnError()
+		st.DisableTracing()
 		return nil, err
 	}
 
 	// Enable coverage collection on the current thread.
 	if err := unix.IoctlSetInt(int(fd), uint(KCOV_ENABLE), KCOV_TRACE_PC); err != nil {
-		cleanupOnError()
-		unix.Munmap(coverageBuffer)
+		st.DisableTracing()
 		return nil, err
 	}
 
-	return &KCOVState{
-		file:  file,
-		cover: coverageBuffer,
-	}, nil
+	return &st, nil
 }
 
 func (st *KCOVState) DisableTracing() {
 	runtime.UnlockOSThread()
 	st.file.Close()
+	unix.Munmap(st.cover)
 }

pkg/kfuzztest/kfuzztest.go

@@ -186,5 +186,6 @@ func ExecKFuzzTestCallLocal(st *kcov.KCOVState, call *prog.Call) ([]uintptr, err
 	finalBlob := prog.MarshallKFuzztestArg(dataArg.Res)
 	inputPath := path.Join("/sys/kernel/debug/kfuzztest/", testName, "input")
 
-	return st.Trace(func() error { return os.WriteFile(inputPath, finalBlob, 0644) })
+	res := st.Trace(func() error { return os.WriteFile(inputPath, finalBlob, 0644) })
+	return res.Coverage, res.Result
 }