coroutine.c: Implement stack protection

1 files changed, 88 insertions, 31 deletions

diff --git a/coroutine.c b/coroutine.c
index fbcb355..97d5be6 100644
--- a/coroutine.c
+++ b/coroutine.c
@@ -6,7 +6,7 @@
 
 #include <stdint.h> /* for uint8_t */
 #include <stdio.h>  /* for printf(), fprintf(), stderr */
-#include <stdlib.h> /* for calloc(), free() */
+#include <stdlib.h> /* for malloc(), free() */
 #include <assert.h>
 #include <setjmp.h>
 
@@ -16,7 +16,8 @@
 
 #define COROUTINE_NUM           5
 #define COROUTINE_MEASURE_STACK 1
-#define COROUTINE_DEBUG         1
+#define COROUTINE_PROTECT_STACK 1
+#define COROUTINE_DEBUG         0
 
 /* Implementation *************************************************************/
 
@@ -67,12 +68,13 @@
 /*
  * Design decisions and notes:
  *
- * - Coroutines are launched with a zeroed-out stack.  Because all
- *   stack variables should be initialized before they are read, this
- *   "shouldn't" make a difference, but: (1) Initializing it to a
- *   known value allows us to measure how much of the stack was
- *   written to, which is helpful to tune stack sizes.  (2) Leaving it
- *   uninitialized just gives me the willies.
+ * - Coroutines are launched with a stack stack that is filled with
+ *   known values (zero or a pre-determined pattern, depending on the
+ *   config above).  Because all stack variables should be initialized
+ *   before they are read, this "shouldn't" make a difference, but:
+ *   (1) Initializing it to a known value allows us to measure how
+ *   much of the stack was written to, which is helpful to tune stack
+ *   sizes.  (2) Leaving it uninitialized just gives me the willies.
  *
  * - Because embedded programs should be adverse to using the heap,
  *   COROUTINE_NUM is fixed, instead of having coroutine_add()
@@ -82,13 +84,11 @@
  *   instead of having them be statically-allocated along with
  *   coroutine_table.  (1) This reduced the blast-area of damage for a
  *   stack-overflow; and indeed if the end of the stack alignes with a
- *   page-boundary memory-protection can even detect the overflow for
- *   us.  (2) Having different-looking addresses for stack-area vs
+ *   page-boundary then memory-protection can even detect the overflow
+ *   for us.  (2) Having different-looking addresses for stack-area vs
  *   static-area is handy for making things jump out at you when
- *   debugging.  (3) Given the above about wanting a zeroed-out stack,
- *   this allows us to take advantage of optimizations in calloc()
- *   instead of using memset, and this can likely also improve things
- *   with being page-aligned.
+ *   debugging.  (3) This can likely also improve things with being
+ *   page-aligned.
  *
  * - Coroutines must use cr_exit() instead of returning because if
  *   they return then they will return to call_with_stack() in
@@ -189,14 +189,45 @@ static void call_with_stack(void *stack, cr_fn_t fn, void *args) {
 	              : "r0"
 	              );
 #else
-#error unsupported architecture
+#  error unsupported architecture
 #endif
 }
 
+#if COROUTINE_MEASURE_STACK || COROUTINE_PROTECT_STACK
+/* We just need a pattern that is unlikely to occur naturaly; this is
+ * just a few bytes that I read from /dev/random.  */
+static const uint8_t const stack_pattern[] = {0x1e, 0x15, 0x16, 0x0a, 0xcc, 0x52, 0x7e, 0xb7};
+#endif
+
+#if COROUTINE_PROTECT_STACK
+void assert_stack_protection(cid_t cid) {
+	assert(coroutine_table[cid-1].stack_size);
+	assert(coroutine_table[cid-1].stack);
+	for (size_t i = 0; i < sizeof(stack_pattern); i++) {
+		size_t j = coroutine_table[cid-1].stack_size - (i+1);
+		assert(((uint8_t*)coroutine_table[cid-1].stack)[i] == stack_pattern[i]);
+		assert(((uint8_t*)coroutine_table[cid-1].stack)[j] == stack_pattern[j%sizeof(stack_pattern)]);
+	}
+}
+#else
+#  define assert_stack_protection(cid) ((void)0)
+#endif
+
+#define assert_cid_state(cid, opstate) do {                     \
+		assert((cid) > 0);                              \
+		assert((cid) <= COROUTINE_NUM);                 \
+		assert(coroutine_table[(cid)-1].state opstate); \
+		assert_stack_protection(cid);                   \
+	} while (0)
+
 cid_t coroutine_add_with_stack_size(size_t stack_size, cr_fn_t fn, void *args) {
 	static cid_t last_created = 0;
+	cid_t parent = coroutine_running;
 
-	assert(coroutine_running == 0 || coroutine_table[coroutine_running-1].state == CR_RUNNING);
+	if (parent)
+		assert_cid_state(parent, == CR_RUNNING);
+	assert(stack_size);
+	assert(fn);
 	debugf("coroutine_add_with_stack_size(%zu, %#p, %#p)...\n", stack_size, fn, args);
 
 	cid_t child;
@@ -210,26 +241,39 @@ cid_t coroutine_add_with_stack_size(size_t stack_size, cr_fn_t fn, void *args) {
 		return 0;
 	found:
 	}
+	debugf("...child=%zu\n", child);
 
 	last_created = child;
 
 	coroutine_table[child-1].stack_size = stack_size;
-	coroutine_table[child-1].stack = calloc(1, coroutine_table[child-1].stack_size);
+	coroutine_table[child-1].stack = malloc(stack_size);
+#if COROUTINE_MEASURE_STACK || COROUTINE_PROTECT_STACK
+	for (size_t i = 0; i < stack_size; i++)
+		((uint8_t*)coroutine_table[child-1].stack)[i] = stack_pattern[i%sizeof(stack_pattern)];
+#endif
 
-	cid_t parent = coroutine_running;
-	assert(parent == 0 || coroutine_table[parent-1].state == CR_RUNNING);
 	coroutine_running = child; 
 	coroutine_table[child-1].state = CR_INITIALIZING;
 	if (!setjmp(coroutine_add_env)) { /* point=a */
+		void *stack_base = coroutine_table[child-1].stack + (STACK_GROWS_DOWNWARD ? stack_size : 0);
+#if COROUTINE_PROTECT_STACK
+#  if STACK_GROWS_DOWNWARD
+		stack_base -= sizeof(stack_pattern);
+#  else
+		stack_base += sizeof(stack_pattern);
+#  endif
+#endif
+		debugf("...stack     =%#p\n", coroutine_table[child-1].stack);
+		debugf("...stack_base=%#p\n", stack_base);
 		/* run until cr_begin() */
-		call_with_stack(coroutine_table[child-1].stack + (STACK_GROWS_DOWNWARD ? coroutine_table[child-1].stack_size : 0), fn, args);
+		call_with_stack(stack_base, fn, args);
 		assert(false); /* should cr_begin() instead of returning */
 	}
-	assert(coroutine_table[child-1].state == CR_RUNNABLE);
-	assert(parent == 0 || coroutine_table[parent-1].state == CR_RUNNING);
+	assert_cid_state(child, == CR_RUNNABLE);
+	if (parent)
+		assert_cid_state(parent, == CR_RUNNING);
 	coroutine_running = parent;
 
-	debugf("coroutine_add_with_stack_size => %zu\n", child);
 	return child;
 }
 
@@ -248,12 +292,20 @@ void coroutine_main(void) {
 				longjmp(cr->env, 1); /* jump to point=c */
 				assert(false); /* should cr_exit() instead of returning */
 			}
+			assert_cid_state(coroutine_running, != CR_RUNNING);
 			if (cr->state == CR_NONE) {
 #if COROUTINE_MEASURE_STACK
-				size_t stack_used = cr->stack_size;
-				while (stack_used > 0 && ((uint8_t*)cr->stack)[STACK_GROWS_DOWNWARD ? cr->stack_size - stack_used : stack_used - 1] == 0)
+				size_t stack_size = cr->stack_size - (COROUTINE_PROTECT_STACK ? 2*sizeof(stack_pattern) : 0);
+				size_t stack_used = stack_size;
+				for (;;) {
+					size_t i = STACK_GROWS_DOWNWARD
+						? (COROUTINE_PROTECT_STACK ? sizeof(stack_pattern) : 0) + stack_size - stack_used
+						: stack_used - 1 - (COROUTINE_PROTECT_STACK ? sizeof(stack_pattern) : 0);
+					if (stack_used == 0 || ((uint8_t*)cr->stack)[i] != stack_pattern[i%sizeof(stack_pattern)])
+						break;
 					stack_used--;
-				printf("info: coroutine %zu exited having used %zu B stack space\n", coroutine_running, stack_used);
+				}
+				printf("info: cid=%zu: exited having used %zu B stack space\n", coroutine_running, stack_used);
 #endif
 				free(cr->stack);
 				coroutine_table[coroutine_running-1] = (struct coroutine){0};
@@ -267,15 +319,17 @@ void coroutine_main(void) {
 }
 
 bool cr_begin(void) {
-	assert(coroutine_table[coroutine_running-1].state == CR_INITIALIZING);
+	assert_cid_state(coroutine_running, == CR_INITIALIZING);
+
 	coroutine_table[coroutine_running-1].state = CR_RUNNABLE;
 	if (!setjmp(coroutine_table[coroutine_running-1].env)) /* point=c1 */
 		longjmp(coroutine_add_env, 1); /* jump to point=a */
 }
 
 static inline __attribute__ ((no_split_stack)) void _cr_transition(enum coroutine_state state) {
-	assert(coroutine_running && coroutine_table[coroutine_running-1].state == CR_RUNNING);
-	debugf("_cr_transition cid=%zu %i->%i\n", coroutine_running, coroutine_table[coroutine_running-1].state, state);
+	assert_cid_state(coroutine_running, == CR_RUNNING);
+	debugf("cid=%zu: transition %i->%i\n", coroutine_running, coroutine_table[coroutine_running-1].state, state);
+
 	coroutine_table[coroutine_running-1].state = state;
 	if (!setjmp(coroutine_table[coroutine_running-1].env)) /* point=c2 */
 		longjmp(coroutine_main_env, 1); /* jump to point=b */
@@ -285,14 +339,17 @@ void cr_yield(void)           { _cr_transition(CR_RUNNABLE); }
 void cr_pause_and_yield(void) { _cr_transition(CR_PAUSED); }
 
 void cr_exit(void) {
-	assert(coroutine_running && coroutine_table[coroutine_running-1].state == CR_RUNNING);
+	assert_cid_state(coroutine_running, == CR_RUNNING);
+	debugf("cid=%zu: exit\n", coroutine_running);
+
 	coroutine_table[coroutine_running-1].state = CR_NONE;
 	longjmp(coroutine_main_env, 1); /* jump to point=b */
 }
 
 void cr_unpause(cid_t cid) {
-	assert(cid && coroutine_table[cid-1].state == CR_PAUSED);
+	assert_cid_state(cid, == CR_PAUSED);
 	debugf("cr_unpause(%zu)\n", cid);
+
 	coroutine_table[cid-1].state = CR_RUNNABLE;
 }