From 6068b565f5f487ad725cacc1ff79ed4451cd89b7 Mon Sep 17 00:00:00 2001
From: ISSOtm <eldredhabert0@gmail.com>
Date: Fri, 30 Aug 2019 04:42:56 +0200
Subject: [PATCH] Add recursion limit for INCLUDE and macros (And REPT.) Not
 exactly a *recursion* limit, more like a *stack depth* limit, but calling it
 "recursion" conveys its purpose better. The default of 64 is super overkill:
 even in a a project with what I believe to be above-average levels of
 nesting, the level only peaked at 6. Keeping in mind the purpose of this is
 to catch infinite recursion, which is still caught quickly (in usual cases,
 anyways), this default seems sensible. And it passes tests. What more do you
 need?

---
 include/asm/fstack.h |  2 ++
 src/asm/fstack.c     | 10 ++++++++++
 src/asm/main.c       | 10 +++++++++-
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/include/asm/fstack.h b/include/asm/fstack.h
index 0c0c74c2..f3e8a878 100644
--- a/include/asm/fstack.h
+++ b/include/asm/fstack.h
@@ -35,6 +35,8 @@ struct sContext {
 	uint32_t nREPTBlockSize;
 };
 
+extern unsigned int nMaxFileStackDepth;
+
 void fstk_RunInclude(char *tzFileName);
 void fstk_RunMacroArg(int32_t s);
 void fstk_Init(char *s);
diff --git a/src/asm/fstack.c b/src/asm/fstack.c
index 3a84877e..b471cdd6 100644
--- a/src/asm/fstack.c
+++ b/src/asm/fstack.c
@@ -28,6 +28,8 @@
 #include "types.h"
 
 static struct sContext *pFileStack;
+static unsigned int nFileStackDepth;
+unsigned int nMaxFileStackDepth;
 static struct sSymbol *pCurrentMacro;
 static YY_BUFFER_STATE CurrentFlexHandle;
 static FILE *pCurrentFile;
@@ -51,6 +53,8 @@ uint32_t ulMacroReturnValue;
 #define STAT_isMacroArg		2
 #define STAT_isREPTBlock	3
 
+/* Max context stack size */
+
 /*
  * Context push and pop
  */
@@ -58,6 +62,9 @@ static void pushcontext(void)
 {
 	struct sContext **ppFileStack;
 
+	if (++nFileStackDepth > nMaxFileStackDepth)
+		fatalerror("Recursion limit (%d) exceeded", nMaxFileStackDepth);
+
 	ppFileStack = &pFileStack;
 	while (*ppFileStack)
 		ppFileStack = &((*ppFileStack)->pNext);
@@ -158,6 +165,8 @@ static int32_t popcontext(void)
 		fatalerror("%s: Internal error.", __func__);
 	}
 
+	nFileStackDepth--;
+
 	free(*ppLastFile);
 	*ppLastFile = NULL;
 	yy_switch_to_buffer(CurrentFlexHandle);
@@ -417,6 +426,7 @@ void fstk_Init(char *pFileName)
 		if (pCurrentFile == NULL)
 			err(1, "Unable to open file '%s'", pFileName);
 	}
+	nFileStackDepth = 0;
 
 	nMacroCount = 0;
 	nCurrentStatus = STAT_isInclude;
diff --git a/src/asm/main.c b/src/asm/main.c
index c2c79e3e..0d928e92 100644
--- a/src/asm/main.c
+++ b/src/asm/main.c
@@ -315,6 +315,8 @@ int main(int argc, char *argv[])
 
 	/* yydebug=1; */
 
+	nMaxFileStackDepth = 64;
+
 	DefaultOptions.gbgfx[0] = '0';
 	DefaultOptions.gbgfx[1] = '1';
 	DefaultOptions.gbgfx[2] = '2';
@@ -332,7 +334,7 @@ int main(int argc, char *argv[])
 
 	newopt = CurrentOptions;
 
-	while ((ch = getopt(argc, argv, "b:D:Eg:hi:LM:o:p:Vvw")) != -1) {
+	while ((ch = getopt(argc, argv, "b:D:Eg:hi:LM:o:p:r:Vvw")) != -1) {
 		switch (ch) {
 		case 'b':
 			if (strlen(optarg) == 2) {
@@ -386,6 +388,12 @@ int main(int argc, char *argv[])
 				errx(1, "Argument for option 'p' must be between 0 and 0xFF");
 
 			break;
+		case 'r':
+			nMaxFileStackDepth = strtoul(optarg, &ep, 0);
+
+			if (optarg[0] == '\0' || *ep != '\0')
+				errx(1, "Invalid argument for option 'r'");
+			break;
 		case 'V':
 			printf("rgbasm %s\n", get_package_version_string());
 			exit(0);