[PATCH] kcov: add AFL-style tracing

[Vegard Nossum]

From: Quentin Casasnovas <quentin.casasnovas@xxxxxxxxxx>

AFL uses a fixed-size buffer (typically 64 KiB) where each byte is
a counter representing how many times an A -> B branch was taken.
Of course, since the buffer is fixed size, it's a little imprecise
in that e.g. two different branches could map to the same counter,
but in practice it works well.

See afl:docs/technical_details.txt for more information.

Here is a small test program that demonstrates the new capability:

	#include <sys/ioctl.h>
	#include <sys/mman.h>
	#include <sys/types.h>

	#include <errno.h>
	#include <error.h>
	#include <fcntl.h>
	#include <stdio.h>
	#include <stdlib.h>
	#include <unistd.h>

	#include <linux/types.h>

	#define KCOV_INIT_TRACE                 _IOR('c', 1, unsigned long)
	#define KCOV_INIT_AFL                   _IOR('c', 2, unsigned long)
	#define KCOV_ENABLE                     _IO('c', 100)
	#define KCOV_DISABLE                    _IO('c', 101)

	int main(int argc, char *argv[])
	{
		int fd = open("/sys/kernel/debug/kcov", O_RDWR);
		if (fd == -1)
			error(1, errno, "open()");

		unsigned long size = 1 << 10;
		if (ioctl(fd, KCOV_INIT_AFL, size) != 0)
			error(1, errno, "ioctl(KCOV_INIT_AFL)");

		void *mem = mmap(NULL, size * sizeof(long), PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
		if (mem == MAP_FAILED)
			error(1, errno, "mmap()");

		/* Start kernel instrumentation */
		if (ioctl(fd, KCOV_ENABLE, 0) != 0)
			error(1, errno, "ioctl(KCOV_ENABLE)");

		printf("Hello world!\n");

		/* End kernel instrumentation*/
		if (ioctl(fd, KCOV_DISABLE, 0) != 0)
			error(1, errno, "ioctl(KCOV_DISABLE)");

		/* Hex dump of memory area */
		unsigned char *mem2 = mem;
		for (unsigned int i = 0; i < size; ++i) {
			printf("%02x ", mem2[i]);
			if (i % 32 == 31)
				printf("\n");
		}

		close(fd);
		return 0;
	}

This patch is a collaboration between Quentin Casasnovas and Vegard Nossum.

Cc: Dmitry Vyukov <dvyukov@xxxxxxxxxx>
Cc: Michal Zalewski <lcamtuf@xxxxxxxxx>
Cc: Kees Cook <keescook@xxxxxxxxxxxx>
Signed-off-by: Quentin Casasnovas <quentin.casasnovas@xxxxxxxxxx>
Signed-off-by: Vegard Nossum <vegard.nossum@xxxxxxxxxx>
---
 include/linux/kcov.h      |  6 ++++++
 include/uapi/linux/kcov.h |  1 +
 kernel/kcov.c             | 44 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 51 insertions(+)

diff --git a/include/linux/kcov.h b/include/linux/kcov.h
index 2883ac9..cb8eb3e 100644
--- a/include/linux/kcov.h
+++ b/include/linux/kcov.h
@@ -18,6 +18,12 @@ enum kcov_mode {
 	 * Covered PCs are collected in a per-task buffer.
 	 */
 	KCOV_MODE_TRACE = 1,
+	/*
+	 * AFL-style collection.
+	 * Covered branches are hashed and collected in a fixed size buffer
+	 * (see AFL documentation for more information).
+	 */
+	KCOV_MODE_AFL = 2,
 };
 
 #else
diff --git a/include/uapi/linux/kcov.h b/include/uapi/linux/kcov.h
index 574e22e..0dc1cf9 100644
--- a/include/uapi/linux/kcov.h
+++ b/include/uapi/linux/kcov.h
@@ -4,6 +4,7 @@
 #include <linux/types.h>
 
 #define KCOV_INIT_TRACE			_IOR('c', 1, unsigned long)
+#define KCOV_INIT_AFL			_IOR('c', 2, unsigned long)
 #define KCOV_ENABLE			_IO('c', 100)
 #define KCOV_DISABLE			_IO('c', 101)
 
diff --git a/kernel/kcov.c b/kernel/kcov.c
index a02f2dd..e6eb785 100644
--- a/kernel/kcov.c
+++ b/kernel/kcov.c
@@ -5,6 +5,7 @@
 #include <linux/types.h>
 #include <linux/file.h>
 #include <linux/fs.h>
+#include <linux/hash.h>
 #include <linux/mm.h>
 #include <linux/printk.h>
 #include <linux/slab.h>
@@ -34,6 +35,19 @@ struct kcov {
 	enum kcov_mode		mode;
 	/* Size of arena (in long's for KCOV_MODE_TRACE). */
 	unsigned		size;
+
+	union {
+		/* For KCOV_MODE_AFL */
+		struct {
+			/* (number of bytes) - 1, where number of
+			 * bytes must be a power of 2. */
+			unsigned int		mask;
+
+			/* Previous PC (for KCOV_MODE_AFL). */
+			unsigned int		prev_location;
+		};
+	};
+
 	/* Coverage buffer shared with user space. */
 	void			*area;
 	/* Task for which we collect coverage, or NULL. */
@@ -76,6 +90,19 @@ void notrace __sanitizer_cov_trace_pc(void)
 			area[pos] = _RET_IP_;
 			WRITE_ONCE(area[0], pos);
 		}
+	} else if (mode == KCOV_MODE_AFL) {
+		struct kcov *kcov;
+		unsigned char *area;
+		unsigned long location = _RET_IP_;
+
+		/* See above */
+		barrier();
+
+		kcov = t->kcov;
+		area = kcov->area;
+
+		++area[(kcov->prev_location ^ location) & kcov->mask];
+		kcov->prev_location = hash_long(location, BITS_PER_LONG);
 	}
 }
 EXPORT_SYMBOL(__sanitizer_cov_trace_pc);
@@ -200,6 +227,23 @@ static int kcov_ioctl_locked(struct kcov *kcov, unsigned int cmd,
 		kcov->size = size;
 		kcov->mode = KCOV_MODE_TRACE;
 		return 0;
+	case KCOV_INIT_AFL:
+		if (kcov->mode != KCOV_MODE_DISABLED)
+			return -EBUSY;
+
+		/*
+		 * AFL mode needs a fixed-size area that is a power of 2.
+		 */
+		size = arg;
+		if (size < 2 || size > INT_MAX / sizeof(unsigned long))
+			return -EINVAL;
+		if (!is_power_of_2(size))
+			return -EINVAL;
+		kcov->size = size;
+		kcov->mask = (size * sizeof(unsigned long)) - 1;
+		kcov->prev_location = hash_long(0, BITS_PER_LONG);
+		kcov->mode = KCOV_MODE_AFL;
+		return 0;
 	case KCOV_ENABLE:
 		/*
 		 * Enable coverage for the current task.
-- 
1.9.1