From 455e090eaf43923adea4cb98c35ac6c6883d4542 Mon Sep 17 00:00:00 2001
From: Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
Date: Sat, 9 Mar 2024 16:14:54 -0500
Subject: [PATCH] mempool: introduce poison attribute

Signed-off-by: Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
Change-Id: I056c0d8c4f0d5b1d2c401592a08b22ff6f239d11
---
 include/rseq/mempool.h | 13 +++++++++++++
 src/rseq-mempool.c     | 38 ++++++++++++++++++++++++++++++++++++++
 tests/mempool_test.c   |  8 ++++++++
 3 files changed, 59 insertions(+)

diff --git a/include/rseq/mempool.h b/include/rseq/mempool.h
index 0b552df..81884d4 100644
--- a/include/rseq/mempool.h
+++ b/include/rseq/mempool.h
@@ -8,6 +8,7 @@
 #include <stddef.h>
 #include <sys/types.h>
 #include <sys/mman.h>
+#include <stdint.h>
 
 /*
  * rseq/mempool.h: rseq memory pool allocator.
@@ -476,6 +477,18 @@ int rseq_mempool_attr_set_global(struct rseq_mempool_attr *attr, size_t stride);
 int rseq_mempool_attr_set_max_nr_ranges(struct rseq_mempool_attr *attr,
 		unsigned long max_nr_ranges);
 
+/*
+ * rseq_mempool_attr_set_poison: Set pool poison value.
+ *
+ * Set a poison value to be set over freed pool entries. This can be
+ * used to anonymize freed memory, and for memory corruption checks
+ * with the robust attribute.
+ *
+ * Returns 0 on success, -1 with errno=EINVAL if arguments are invalid.
+ */
+int rseq_mempool_attr_set_poison(struct rseq_mempool_attr *attr,
+		uintptr_t poison);
+
 /*
  * rseq_mempool_range_init_numa: NUMA initialization helper for memory range.
  *
diff --git a/src/rseq-mempool.c b/src/rseq-mempool.c
index bf569a6..a68d1f8 100644
--- a/src/rseq-mempool.c
+++ b/src/rseq-mempool.c
@@ -80,6 +80,9 @@ struct rseq_mempool_attr {
 	int max_nr_cpus;
 
 	unsigned long max_nr_ranges;
+
+	bool poison_set;
+	uintptr_t poison;
 };
 
 struct rseq_mempool_range;
@@ -149,6 +152,23 @@ void rseq_percpu_zero_item(struct rseq_mempool *pool,
 	}
 }
 
+static
+void rseq_percpu_poison_item(struct rseq_mempool *pool,
+		struct rseq_mempool_range *range, uintptr_t item_offset)
+{
+	uintptr_t poison = pool->attr.poison;
+	int i;
+
+	for (i = 0; i < pool->attr.max_nr_cpus; i++) {
+		char *p = __rseq_pool_range_percpu_ptr(range, i,
+				item_offset, pool->attr.stride);
+		size_t offset;
+
+		for (offset = 0; offset < pool->item_len; offset += sizeof(uintptr_t))
+			*((uintptr_t *) p) = poison;
+	}
+}
+
 #ifdef HAVE_LIBNUMA
 int rseq_mempool_range_init_numa(void *addr, size_t len, int cpu, int numa_flags)
 {
@@ -725,8 +745,14 @@ void librseq_mempool_percpu_free(void __rseq_percpu *_ptr, size_t stride)
 	clear_alloc_slot(pool, range, item_offset);
 	/* Add ptr to head of free list */
 	head = pool->free_list_head;
+	if (pool->attr.poison_set)
+		rseq_percpu_poison_item(pool, range, item_offset);
 	/* Free-list is in CPU 0 range. */
 	item = (struct free_list_node *) ptr;
+	/*
+	 * Setting the next pointer will overwrite the first uintptr_t
+	 * poison for CPU 0.
+	 */
 	item->next = head;
 	pool->free_list_head = item;
 	pthread_mutex_unlock(&pool->lock);
@@ -919,6 +945,18 @@ int rseq_mempool_attr_set_max_nr_ranges(struct rseq_mempool_attr *attr,
 	return 0;
 }
 
+int rseq_mempool_attr_set_poison(struct rseq_mempool_attr *attr,
+		uintptr_t poison)
+{
+	if (!attr) {
+		errno = EINVAL;
+		return -1;
+	}
+	attr->poison_set = true;
+	attr->poison = poison;
+	return 0;
+}
+
 int rseq_mempool_get_max_nr_cpus(struct rseq_mempool *mempool)
 {
 	if (!mempool || mempool->attr.type != MEMPOOL_TYPE_PERCPU) {
diff --git a/tests/mempool_test.c b/tests/mempool_test.c
index e02a95d..05108cb 100644
--- a/tests/mempool_test.c
+++ b/tests/mempool_test.c
@@ -24,6 +24,12 @@
 #include "list.h"
 #include "tap.h"
 
+#if RSEQ_BITS_PER_LONG == 64
+# define POISON_VALUE	0xABCDABCDABCDABCDULL
+#else
+# define POISON_VALUE	0xABCDABCDUL
+#endif
+
 struct test_data {
 	uintptr_t value;
 	struct test_data __rseq_percpu *backref;
@@ -48,6 +54,8 @@ static void test_mempool_fill(unsigned long max_nr_ranges, size_t stride)
 	ok(ret == 0, "Setting mempool percpu type");
 	ret = rseq_mempool_attr_set_max_nr_ranges(attr, max_nr_ranges);
 	ok(ret == 0, "Setting mempool max_nr_ranges=%lu", max_nr_ranges);
+	ret = rseq_mempool_attr_set_poison(attr, POISON_VALUE);
+	ok(ret == 0, "Setting mempool poison");
 	mempool = rseq_mempool_create("test_data",
 			sizeof(struct test_data), attr);
 	ok(mempool, "Create mempool of size %zu", stride);
-- 
2.34.1