From 6a7073e5f262db9a4578ff00b5b28e34335564ce Mon Sep 17 00:00:00 2001
From: Kimplul <kimi.h.kuparinen@gmail.com>
Date: Thu, 23 May 2024 19:54:55 +0300
Subject: improve id handling

+ The number of allowed threads running at the same time is limited to
  num_tids, but each thread's ID can be any larger than that. This
  should make ID reuse a lot more rare, and probably makes certain kinds
  of time-of-check-to-time-of-use attacks more difficult
---
 arch/riscv64/conf/init.c |  7 +++----
 common/tcb.c             | 17 +++++++++--------
 common/uapi/ipc.c        |  2 +-
 include/kmi/bits.h       |  2 ++
 4 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/arch/riscv64/conf/init.c b/arch/riscv64/conf/init.c
index 6783283..293d5e8 100644
--- a/arch/riscv64/conf/init.c
+++ b/arch/riscv64/conf/init.c
@@ -255,13 +255,12 @@ static void *sys_req_sharedmem(long tid, unsigned long size, void **cbuf)
 static char *rw_buf = 0;
 static size_t rw_buf_size = 4096;
 
-void callback(long status, long tid, long d0, long d1, long d2, long d3)
+void callback(long pid, long tid, long d0, long d1, long d2, long d3)
 {
-	(void)status;
-
+	(void)tid;
 	if (d0 == 1) {
 		void *cbuf = 0;
-		rw_buf = sys_req_sharedmem(tid, rw_buf_size, &cbuf);
+		rw_buf = sys_req_sharedmem(pid, rw_buf_size, &cbuf);
 		sys_ipc_resp((long)cbuf, rw_buf_size, 0, 0);
 		__builtin_unreachable();
 
diff --git a/common/tcb.c b/common/tcb.c
index 4b57c06..c0f82cc 100644
--- a/common/tcb.c
+++ b/common/tcb.c
@@ -23,7 +23,7 @@
 
 /* arguably exessively many globals... */
 /** Thread ID to start looking from when allocating new ID. */
-static id_t start_tid;
+static id_t start_tid = 0;
 
 /** Total number of possible thread IDs. */
 static id_t num_tids;
@@ -46,6 +46,7 @@ void init_tcbs()
 	 * something smaller but this is fine for now. */
 	tcbs = (struct tcb **)alloc_page(MM_O1);
 	num_tids = order_size(MM_O1) / sizeof(struct tcb *);
+	catastrophic_assert(is_powerof2(num_tids));
 	memset(tcbs, 0, order_size(MM_O1));
 }
 
@@ -64,18 +65,18 @@ static id_t __alloc_tid(struct tcb *t)
 {
 	id_t stop_tid = start_tid - 1;
 	/** \todo this would need some locking or something... */
-	for (id_t i = start_tid; 1; ++i) {
-		if (i == ID_MAX)
-			i = 0;
+	for (id_t i = start_tid;; ++i) {
+		if (i <= 0)
+			i = 1;
 
 		/* we're completely full */
 		if (i == stop_tid)
 			return ERR_NF;
 
-		if (tcbs[i] || i == 0)
+		if (get_tcb(i) || i == 0)
 			continue;
 
-		tcbs[i] = t;
+		tcbs[i & (num_tids - 1)] = t;
 		start_tid = i + 1;
 		return i;
 	}
@@ -292,10 +293,10 @@ struct tcb *get_tcb(id_t tid)
 {
 	hard_assert(tcbs, 0);
 
-	if (tid <= 0 || tid >= num_tids)
+	if (tid <= 0)
 		return NULL;
 
-	return tcbs[tid];
+	return tcbs[tid & (num_tids - 1)];
 }
 
 void set_return(struct tcb *t, vm_t v)
diff --git a/common/uapi/ipc.c b/common/uapi/ipc.c
index 18c3b1a..cb2ce28 100644
--- a/common/uapi/ipc.c
+++ b/common/uapi/ipc.c
@@ -216,7 +216,7 @@ static void do_ipc(struct tcb *t,
 	if (unlikely(!enough_rpc_stack(t)))
 		return_args1(t, ERR_OOMEM);
 
-	vm_t s = enter_rpc(t, SYS_RET6(OK, t->eid, d0, d1, d2, d3), kind);
+	vm_t s = enter_rpc(t, SYS_RET6(t->eid, t->tid, d0, d1, d2, d3), kind);
 
 	struct tcb *r = get_tcb(pid);
 	if (unlikely(!r)) {
diff --git a/include/kmi/bits.h b/include/kmi/bits.h
index 738f1d5..5fe438c 100644
--- a/include/kmi/bits.h
+++ b/include/kmi/bits.h
@@ -379,4 +379,6 @@ uint64_t __bswap64(uint64_t u);
 
 #endif
 
+#define is_powerof2(x) (((x) & ((x) - 1)) == 0)
+
 #endif /* KMI_BITS_H */
-- 
cgit v1.3