executor: protect kcov/output regions with pkeys

Protect KCOV regions with pkeys if they are available.
Protect output region with pkeys in snapshot mode.
Snapshot mode is especially sensitive to output buffer corruption
since its location is not randomized.

4 files changed, 84 insertions, 10 deletions

diff --git a/executor/common_linux.h b/executor/common_linux.h
index 193afcdda..3669dee0f 100644
--- a/executor/common_linux.h
+++ b/executor/common_linux.h
@@ -5712,10 +5712,16 @@ static long syz_clone3(volatile long a0, volatile long a1)
 #endif
 
 #if SYZ_EXECUTOR || __NR_syz_pkey_set
+#include <errno.h>
+#define RESERVED_PKEY 15
 // syz_pkey_set(key pkey, val flags[pkey_flags])
 static long syz_pkey_set(volatile long pkey, volatile long val)
 {
 #if GOARCH_amd64 || GOARCH_386
+	if (pkey == RESERVED_PKEY) {
+		errno = EINVAL;
+		return -1;
+	}
 	uint32 eax = 0;
 	uint32 ecx = 0;
 	asm volatile("rdpkru"
diff --git a/executor/executor.cc b/executor/executor.cc
index 02971f374..3757f2698 100644
--- a/executor/executor.cc
+++ b/executor/executor.cc
@@ -226,11 +226,12 @@ private:
 class ShmemBuilder : ShmemAllocator, public flatbuffers::FlatBufferBuilder
 {
 public:
-	ShmemBuilder(OutputData* data, size_t size)
+	ShmemBuilder(OutputData* data, size_t size, bool store_size)
 	    : ShmemAllocator(data + 1, size - sizeof(*data)),
 	      FlatBufferBuilder(size - sizeof(*data), this)
 	{
-		data->size.store(size, std::memory_order_relaxed);
+		if (store_size)
+			data->size.store(size, std::memory_order_relaxed);
 		size_t consumed = data->consumed.load(std::memory_order_relaxed);
 		if (consumed >= size - sizeof(*data))
 			failmsg("ShmemBuilder: too large output offset", "size=%zd consumed=%zd", size, consumed);
@@ -491,6 +492,35 @@ static void parse_handshake(const handshake_req& req);
 #error "unknown OS"
 #endif
 
+class CoverAccessScope final
+{
+public:
+	CoverAccessScope(cover_t* cov)
+	    : cov_(cov)
+	{
+		// CoverAccessScope must not be used recursively b/c on Linux pkeys protection is global,
+		// so cover_protect for one cov overrides previous cover_unprotect for another cov.
+		if (used_)
+			fail("recursion in CoverAccessScope");
+		used_ = true;
+		cover_unprotect(cov_);
+	}
+	~CoverAccessScope()
+	{
+		cover_protect(cov_);
+		used_ = false;
+	}
+
+private:
+	cover_t* const cov_;
+	static bool used_;
+
+	CoverAccessScope(const CoverAccessScope&) = delete;
+	CoverAccessScope& operator=(const CoverAccessScope&) = delete;
+};
+
+bool CoverAccessScope::used_;
+
 #if !SYZ_HAVE_FEATURES
 static feature_t features[] = {};
 #endif
@@ -818,7 +848,9 @@ void execute_one()
 		SnapshotStart();
 	else
 		realloc_output_data();
-	output_builder.emplace(output_data, output_size);
+	// Output buffer may be pkey-protected in snapshot mode, so don't write the output size
+	// (it's fixed and known anyway).
+	output_builder.emplace(output_data, output_size, !flag_snapshot);
 	uint64 start = current_time_ms();
 	uint8* input_pos = input_data;
 
@@ -1116,10 +1148,8 @@ uint32 write_cover(flatbuffers::FlatBufferBuilder& fbb, cover_t* cov)
 	cover_data_t* cover_data = (cover_data_t*)(cov->data + cov->data_offset);
 	if (flag_dedup_cover) {
 		cover_data_t* end = cover_data + cover_size;
-		cover_unprotect(cov);
 		std::sort(cover_data, end);
 		cover_size = std::unique(cover_data, end) - cover_data;
-		cover_protect(cov);
 	}
 	fbb.StartVector(cover_size, sizeof(uint64));
 	for (uint32 i = 0; i < cover_size; i++)
@@ -1134,7 +1164,6 @@ uint32 write_comparisons(flatbuffers::FlatBufferBuilder& fbb, cover_t* cov)
 	kcov_comparison_t* cov_start = (kcov_comparison_t*)(cov->data + sizeof(uint64));
 	if ((char*)(cov_start + ncomps) > cov->data_end)
 		failmsg("too many comparisons", "ncomps=%llu", ncomps);
-	cover_unprotect(cov);
 	rpc::ComparisonRaw* start = (rpc::ComparisonRaw*)cov_start;
 	rpc::ComparisonRaw* end = start;
 	// We will convert kcov_comparison_t to ComparisonRaw inplace.
@@ -1156,7 +1185,6 @@ uint32 write_comparisons(flatbuffers::FlatBufferBuilder& fbb, cover_t* cov)
 			 return a.pc() == b.pc() && a.op1() == b.op1() && a.op2() == b.op2();
 		 }) -
 		 start;
-	cover_protect(cov);
 	return fbb.CreateVectorOfStructs(start, ncomps).o;
 }
 
@@ -1229,6 +1257,7 @@ void copyout_call_results(thread_t* th)
 
 void write_output(int index, cover_t* cov, rpc::CallFlag flags, uint32 error, bool all_signal)
 {
+	CoverAccessScope scope(cov);
 	auto& fbb = *output_builder;
 	const uint32 start_size = output_builder->GetSize();
 	(void)start_size;
@@ -1304,11 +1333,13 @@ void write_extra_output()
 flatbuffers::span<uint8_t> finish_output(OutputData* output, int proc_id, uint64 req_id, uint32 num_calls, uint64 elapsed,
 					 uint64 freshness, uint32 status, const std::vector<uint8_t>* process_output)
 {
-	int output_size = output->size.load(std::memory_order_relaxed) ?: kMaxOutput;
+	// In snapshot mode the output size is fixed and output_size is always initialized, so use it.
+	int out_size = flag_snapshot ? output_size : output->size.load(std::memory_order_relaxed) ?
+												  : kMaxOutput;
 	uint32 completed = output->completed.load(std::memory_order_relaxed);
 	completed = std::min(completed, kMaxCalls);
-	debug("handle completion: completed=%u output_size=%u\n", completed, output_size);
-	ShmemBuilder fbb(output, output_size);
+	debug("handle completion: completed=%u output_size=%u\n", completed, out_size);
+	ShmemBuilder fbb(output, out_size, false);
 	auto empty_call = rpc::CreateCallInfoRawDirect(fbb, rpc::CallFlag::NONE, 998);
 	std::vector<flatbuffers::Offset<rpc::CallInfoRaw>> calls(num_calls, empty_call);
 	std::vector<flatbuffers::Offset<rpc::CallInfoRaw>> extra;
diff --git a/executor/executor_linux.h b/executor/executor_linux.h
index 072831816..82a1d559a 100644
--- a/executor/executor_linux.h
+++ b/executor/executor_linux.h
@@ -11,6 +11,8 @@
 #include <sys/syscall.h>
 #include <unistd.h>
 
+static bool pkeys_enabled;
+
 const unsigned long KCOV_TRACE_PC = 0;
 const unsigned long KCOV_TRACE_CMP = 1;
 
@@ -65,6 +67,22 @@ static void os_init(int argc, char** argv, char* data, size_t data_size)
 	struct sigaction act = {};
 	act.sa_handler = [](int) {};
 	sigaction(SIGCHLD, &act, nullptr);
+
+	// Use the last available pkey so that C reproducers get the the same keys from pkey_alloc.
+	int pkeys[RESERVED_PKEY + 1];
+	int npkey = 0;
+	for (; npkey <= RESERVED_PKEY; npkey++) {
+		int pk = pkey_alloc(0, 0);
+		if (pk == -1)
+			break;
+		if (pk == RESERVED_PKEY) {
+			pkeys_enabled = true;
+			break;
+		}
+		pkeys[npkey] = pk;
+	}
+	while (npkey--)
+		pkey_free(pkeys[npkey]);
 }
 
 static intptr_t execute_syscall(const call_t* c, intptr_t a[kMaxArgs])
@@ -87,14 +105,20 @@ static void cover_open(cover_t* cov, bool extra)
 	if (ioctl(cov->fd, kcov_init_trace, cover_size))
 		fail("cover init trace write failed");
 	cov->mmap_alloc_size = cover_size * (is_kernel_64_bit ? 8 : 4);
+	if (pkeys_enabled)
+		debug("pkey protection enabled\n");
 }
 
 static void cover_protect(cover_t* cov)
 {
+	if (pkeys_enabled && pkey_set(RESERVED_PKEY, PKEY_DISABLE_WRITE))
+		debug("pkey_set failed: %d\n", errno);
 }
 
 static void cover_unprotect(cover_t* cov)
 {
+	if (pkeys_enabled && pkey_set(RESERVED_PKEY, 0))
+		debug("pkey_set failed: %d\n", errno);
 }
 
 static void cover_mmap(cover_t* cov)
@@ -113,6 +137,8 @@ static void cover_mmap(cover_t* cov)
 				PROT_READ | PROT_WRITE, MAP_SHARED | MAP_FIXED, cov->fd, 0);
 	if (cov->data == MAP_FAILED)
 		exitf("cover mmap failed");
+	if (pkeys_enabled && pkey_mprotect(cov->data, cov->mmap_alloc_size, PROT_READ | PROT_WRITE, RESERVED_PKEY))
+		exitf("failed to pkey_mprotect kcov buffer");
 	cov->data_end = cov->data + cov->mmap_alloc_size;
 	cov->data_offset = is_kernel_64_bit ? sizeof(uint64_t) : sizeof(uint32_t);
 	cov->pc_offset = 0;
@@ -151,7 +177,9 @@ static void cover_reset(cover_t* cov)
 			fail("cover_reset: current_thread == 0");
 		cov = &current_thread->cov;
 	}
+	cover_unprotect(cov);
 	*(uint64*)cov->data = 0;
+	cover_protect(cov);
 }
 
 static void cover_collect(cover_t* cov)
diff --git a/executor/snapshot.h b/executor/snapshot.h
index 5479a162f..ce3e355b0 100644
--- a/executor/snapshot.h
+++ b/executor/snapshot.h
@@ -87,6 +87,11 @@ static void FindIvshmemDevices()
 			      input, static_cast<uint64>(rpc::Const::MaxInputSize));
 			debug("mapped shmem output at at %p/%llu\n",
 			      output, static_cast<uint64>(rpc::Const::MaxOutputSize));
+#if GOOS_linux
+			if (pkeys_enabled && pkey_mprotect(output, static_cast<uint64>(rpc::Const::MaxOutputSize),
+							   PROT_READ | PROT_WRITE, RESERVED_PKEY))
+				exitf("failed to pkey_mprotect output buffer");
+#endif
 		}
 		close(res2);
 	}
@@ -175,6 +180,8 @@ static void TouchMemory(void* ptr, size_t size)
 #if SYZ_EXECUTOR_USES_FORK_SERVER
 static void SnapshotPrepareParent()
 {
+	// This allows access to the output region.
+	CoverAccessScope scope(nullptr);
 	TouchMemory((char*)output_data + output_size - kOutputPopulate, kOutputPopulate);
 	// Notify SnapshotStart that we finished prefaulting memory in the parent.
 	output_data->completed = 1;
@@ -188,6 +195,7 @@ static void SnapshotPrepareParent()
 static void SnapshotStart()
 {
 	debug("SnapshotStart\n");
+	CoverAccessScope scope(nullptr);
 	// Prefault as much memory as we can before the snapshot is taken.
 	// Also pre-create some threads and let them block.
 	// This is intended to make execution after each snapshot restore faster,
@@ -241,6 +249,7 @@ static void SnapshotStart()
 NORETURN static void SnapshotDone(bool failed)
 {
 	debug("SnapshotDone\n");
+	CoverAccessScope scope(nullptr);
 	uint32 num_calls = output_data->num_calls.load(std::memory_order_relaxed);
 	auto data = finish_output(output_data, 0, 0, num_calls, 0, 0, failed ? kFailStatus : 0, nullptr);
 	ivs.hdr->output_offset = data.data() - reinterpret_cast<volatile uint8_t*>(ivs.hdr);