From b22985712f481ca889ad090a5f60728b7384d28c Mon Sep 17 00:00:00 2001
From: Raphael Assenat <raph@raphnet.net>
Date: Sun, 8 Jan 2017 15:14:55 -0500
Subject: [PATCH] Add stack overgrow detection

If the stack ever grows too large (and starts overwriting variables
in .bss) the firmware jumps into the bootloader. This is better than
just continuing to run with strange side effects.
---
 Makefile        |  2 +-
 Makefile.inc    |  2 +-
 Makefile.stk525 |  2 +-
 main.c          |  6 +++++
 stkchk.c        | 60 +++++++++++++++++++++++++++++++++++++++++++++++++
 stkchk.h        | 16 +++++++++++++
 6 files changed, 85 insertions(+), 3 deletions(-)
 create mode 100644 stkchk.c
 create mode 100644 stkchk.h

diff --git a/Makefile b/Makefile
index 661ec02..65fe5d2 100644
--- a/Makefile
+++ b/Makefile
@@ -7,7 +7,7 @@ include Makefile.inc
 PROGNAME=gcn64usb
 OBJDIR=objs-$(PROGNAME)
 CPU=atmega32u2
-CFLAGS=-Wall -mmcu=$(CPU) -DF_CPU=16000000L -Os -DUART1_STDOUT -DVERSIONSTR=$(VERSIONSTR) -DVERSIONSTR_SHORT=$(VERSIONSTR_SHORT)
+CFLAGS=-Wall -mmcu=$(CPU) -DF_CPU=16000000L -Os -DUART1_STDOUT -DVERSIONSTR=$(VERSIONSTR) -DVERSIONSTR_SHORT=$(VERSIONSTR_SHORT) -std=gnu99
 LDFLAGS=-mmcu=$(CPU) -Wl,-Map=$(PROGNAME).map
 HEXFILE=$(PROGNAME).hex
 
diff --git a/Makefile.inc b/Makefile.inc
index ce7668a..f4613e0 100644
--- a/Makefile.inc
+++ b/Makefile.inc
@@ -1,3 +1,3 @@
-OBJS=main.o usb.o usbpad.o mappings.o gcn64_protocol.o n64.o gamecube.o usart1.o bootloader.o eeprom.o config.o hiddata.o usbstrings.o intervaltimer.o version.o gcn64txrx0.o gcn64txrx1.o gcn64txrx2.o gcn64txrx3.o gamepads.o
+OBJS=main.o usb.o usbpad.o mappings.o gcn64_protocol.o n64.o gamecube.o usart1.o bootloader.o eeprom.o config.o hiddata.o usbstrings.o intervaltimer.o version.o gcn64txrx0.o gcn64txrx1.o gcn64txrx2.o gcn64txrx3.o gamepads.o stkchk.o
 VERSIONSTR=\"3.4.0\"
 VERSIONSTR_SHORT=\"3.4\"
diff --git a/Makefile.stk525 b/Makefile.stk525
index e071fcc..283344a 100644
--- a/Makefile.stk525
+++ b/Makefile.stk525
@@ -7,7 +7,7 @@ include Makefile.inc
 PROGNAME=gcn64usb-stk500
 OBJDIR=objs-$(PROGNAME)
 CPU=at90usb1287
-CFLAGS=-Wall -mmcu=$(CPU) -DF_CPU=16000000L -Os -DUART1_STDOUT -DSTK525 -DVERSIONSTR=$(VERSIONSTR) -DVERSIONSTR_SHORT=$(VERSIONSTR_SHORT)
+CFLAGS=-Wall -mmcu=$(CPU) -DF_CPU=16000000L -Os -DUART1_STDOUT -DSTK525 -DVERSIONSTR=$(VERSIONSTR) -DVERSIONSTR_SHORT=$(VERSIONSTR_SHORT) -std=gnu99
 LDFLAGS=-mmcu=$(CPU) -Wl,-Map=$(PROGNAME).map
 HEXFILE=$(PROGNAME).hex
 
diff --git a/main.c b/main.c
index acb19a2..f5d4634 100644
--- a/main.c
+++ b/main.c
@@ -37,6 +37,7 @@
 #include "usbstrings.h"
 #include "intervaltimer.h"
 #include "requests.h"
+#include "stkchk.h"
 
 #define MAX_PLAYERS		2
 
@@ -444,6 +445,7 @@ int main(void)
 	usart1_init();
 	eeprom_init();
 	intervaltimer_init();
+	stkchk_init();
 
 	switch (g_eeprom_data.cfg.mode)
 	{
@@ -511,6 +513,10 @@ int main(void)
 	{
 		static char last_v[MAX_PLAYERS] = { };
 
+		if (stkchk_verify()) {
+			enterBootLoader();
+		}
+
 		usb_doTasks();
 		hiddata_doTask(&hiddata_ops);
 
diff --git a/stkchk.c b/stkchk.c
new file mode 100644
index 0000000..9809442
--- /dev/null
+++ b/stkchk.c
@@ -0,0 +1,60 @@
+#include <stdio.h>
+#include <stdint.h>
+#include <avr/io.h>
+#include <util/atomic.h>
+#include "stkchk.h"
+
+extern uint16_t __stack;
+extern uint16_t _end;
+
+/** Write a canary at the end of the stack. */
+void stkchk_init(void)
+{
+	*((&_end)-1) = 0xDEAD;
+}
+
+/** Check if the canary is still alive.
+ *
+ * Call this perdiocally to check if the
+ * stack grew too large.
+ **/
+char stkchk_verify(void)
+{
+	if (*((&_end)-1) != 0xDEAD) {
+		return -1;
+	}
+
+	return 0;
+}
+
+/* In order to get an approximate idea of how
+ * much stack you are using, this can be called
+ * at strategic places where you know the
+ * call stack is deep or where large automatic
+ * buffers are used.
+ */
+#ifdef STKCHK_WITH_STATUS_CHECK
+void stkchk(const char *fname)
+{
+	static int max_usage = 0;
+	uint16_t end = ((uint16_t)&_end);
+	uint16_t s_top = ((uint16_t)&__stack);
+	uint16_t s_cur = SPL | SPH<<8;
+	uint16_t used, s_size;
+	uint8_t grew = 0;
+
+	ATOMIC_BLOCK(ATOMIC_RESTORESTATE)
+	{
+		used = s_top - s_cur;
+		s_size = s_top - end;
+		if (used > max_usage) {
+			max_usage = used;
+			grew = 1;
+		}
+	}
+
+	if (grew) {
+		printf("[stkchk] %s: %d/%d\r\n", fname, used, s_size);
+	}
+}
+#endif
diff --git a/stkchk.h b/stkchk.h
new file mode 100644
index 0000000..79ad431
--- /dev/null
+++ b/stkchk.h
@@ -0,0 +1,16 @@
+#ifndef _stkchk_h__
+#define _stkchk_h__
+
+#undef STKCHK_WITH_STATUS_CHECK
+
+#ifdef STKCHK_WITH_STATUS_CHECK
+#define stkcheck()	stkchk(__FUNCTION__);
+void stkchk(const char *label);
+#else
+#define stkcheck()
+#endif
+
+void stkchk_init(void);
+char stkchk_verify(void);
+
+#endif // _stkchk_h__