[PATCH v2] io_uring: annotate remote tasks for kcoverage

[Robert Femmer <robert@fmmr.tech>]

Fuzzers use coverage information to guide generation of test cases
towards new or interesting code paths. Syzkaller, specifically, makes
use kcoverage (CONFIG_KCOV). Coverage information is not collected for
kernel tasks unless annotated by kcov_remote_start and kcov_remote_stop.
This patch annotates io-uring's work queue and sqpoll tasks.

Signed-off-by: Robert Femmer <robert@fmmr.tech>
---
 include/linux/io_uring_types.h |  4 ++++
 io_uring/io-wq.c               |  4 ++++
 io_uring/io_uring.c            |  3 +++
 io_uring/io_uring.h            | 24 ++++++++++++++++++++++++
 io_uring/sqpoll.c              |  4 ++++
 5 files changed, 39 insertions(+)

diff --git a/include/linux/io_uring_types.h b/include/linux/io_uring_types.h
index 244392026c6d..b92b8e7169ea 100644
--- a/include/linux/io_uring_types.h
+++ b/include/linux/io_uring_types.h
@@ -504,6 +504,10 @@ struct io_ring_ctx {
 	struct io_mapped_region		ring_region;
 	/* used for optimised request parameter and wait argument passing  */
 	struct io_mapped_region		param_region;
+
+#ifdef CONFIG_KCOV
+	u64				kcov_handle;
+#endif
 };
 
 /*
diff --git a/io_uring/io-wq.c b/io_uring/io-wq.c
index 8cc7b47d3089..16af75b1cfe0 100644
--- a/io_uring/io-wq.c
+++ b/io_uring/io-wq.c
@@ -639,6 +639,7 @@ static void io_worker_handle_work(struct io_wq_acct *acct,
 		/* handle a whole dependent link */
 		do {
 			struct io_wq_work *next_hashed, *linked;
+			struct io_kiocb *req;
 			unsigned int work_flags = atomic_read(&work->flags);
 			unsigned int hash = __io_wq_is_hashed(work_flags)
 				? __io_get_work_hash(work_flags)
@@ -649,7 +650,10 @@ static void io_worker_handle_work(struct io_wq_acct *acct,
 			if (do_kill &&
 			    (work_flags & IO_WQ_WORK_UNBOUND))
 				atomic_or(IO_WQ_WORK_CANCEL, &work->flags);
+			req = container_of(work, struct io_kiocb, work);
+			io_kcov_remote_start(req->ctx);
 			io_wq_submit_work(work);
+			io_kcov_remote_stop(req->ctx);
 			io_assign_current_work(worker, NULL);
 
 			linked = io_wq_free_work(work);
diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
index 036145ee466c..f38b8eca6bbb 100644
--- a/io_uring/io_uring.c
+++ b/io_uring/io_uring.c
@@ -293,6 +293,9 @@ static __cold struct io_ring_ctx *io_ring_ctx_alloc(struct io_uring_params *p)
 	INIT_HLIST_HEAD(&ctx->cancelable_uring_cmd);
 	io_napi_init(ctx);
 	mutex_init(&ctx->mmap_lock);
+#ifdef CONFIG_KCOV
+	ctx->kcov_handle = current->kcov_handle;
+#endif
 
 	return ctx;
 
diff --git a/io_uring/io_uring.h b/io_uring/io_uring.h
index e612a66ee80e..881d43bd529c 100644
--- a/io_uring/io_uring.h
+++ b/io_uring/io_uring.h
@@ -7,6 +7,7 @@
 #include <linux/resume_user_mode.h>
 #include <linux/poll.h>
 #include <linux/io_uring_types.h>
+#include <linux/kcov.h>
 #include <uapi/linux/eventpoll.h>
 #include "alloc_cache.h"
 #include "io-wq.h"
@@ -581,4 +582,27 @@ static inline bool io_has_work(struct io_ring_ctx *ctx)
 	return test_bit(IO_CHECK_CQ_OVERFLOW_BIT, &ctx->check_cq) ||
 	       io_local_work_pending(ctx);
 }
+
+#ifdef CONFIG_KCOV
+static inline void io_kcov_remote_start(struct io_ring_ctx *ctx)
+{
+	if (ctx->kcov_handle)
+		kcov_remote_start(ctx->kcov_handle);
+}
+
+static inline void io_kcov_remote_stop(struct io_ring_ctx *ctx)
+{
+	if (ctx->kcov_handle)
+		kcov_remote_stop();
+}
+#else
+static inline void io_kcov_remote_start(struct io_ring_ctx *ctx)
+{
+}
+
+static inline void io_kcov_remote_stop(struct io_ring_ctx *ctx)
+{
+}
+#endif
+
 #endif
diff --git a/io_uring/sqpoll.c b/io_uring/sqpoll.c
index 46c12afec73e..8d2876e31acb 100644
--- a/io_uring/sqpoll.c
+++ b/io_uring/sqpoll.c
@@ -342,19 +342,23 @@ static int io_sq_thread(void *data)
 
 		cap_entries = !list_is_singular(&sqd->ctx_list);
 		list_for_each_entry(ctx, &sqd->ctx_list, sqd_list) {
+			io_kcov_remote_start(ctx);
 			int ret = __io_sq_thread(ctx, sqd, cap_entries, &ist);
 
 			if (!sqt_spin && (ret > 0 || !list_empty(&ctx->iopoll_list)))
 				sqt_spin = true;
+			io_kcov_remote_stop(ctx);
 		}
 		if (io_sq_tw(&retry_list, IORING_TW_CAP_ENTRIES_VALUE))
 			sqt_spin = true;
 
 		list_for_each_entry(ctx, &sqd->ctx_list, sqd_list) {
+			io_kcov_remote_start(ctx);
 			if (io_napi(ctx)) {
 				io_sq_start_worktime(&ist);
 				io_napi_sqpoll_busy_poll(ctx);
 			}
+			io_kcov_remote_stop(ctx);
 		}
 
 		io_sq_update_worktime(sqd, &ist);
-- 
2.54.0