Merge pull request #411 from ISSOtm/recursion_limit

Add a recursion limit

include/asm/fstack.h

@@ -35,6 +35,8 @@ struct sContext {
 	uint32_t nREPTBlockSize;
 };
 
+extern unsigned int nMaxRecursionDepth;
+
 void fstk_RunInclude(char *tzFileName);
 void fstk_RunMacroArg(int32_t s);
 void fstk_Init(char *s);

include/asm/lexer.h

@@ -40,6 +40,13 @@ enum eLexerState {
 	LEX_STATE_MACROARGS
 };
 
+struct sStringExpansionPos {
+	char *tzName;
+	char *pBuffer;
+	char *pBufferPos;
+	struct sStringExpansionPos *pParent;
+};
+
 #define INITIAL		0
 #define macroarg	3
 
@@ -62,14 +69,16 @@ void lex_FloatDeleteSecondRange(uint32_t id, uint16_t start, uint16_t end);
 void lex_Init(void);
 void lex_AddStrings(const struct sLexInitString *lex);
 void lex_SetBuffer(char *buffer, uint32_t len);
+void lex_BeginStringExpansion(const char *tzName);
 int yywrap(void);
 int yylex(void);
 void yyunput(char c);
-void yyunputstr(char *s);
+void yyunputstr(const char *s);
 void yyskipbytes(uint32_t count);
 void yyunputbytes(uint32_t count);
 
 extern YY_BUFFER_STATE pCurrentBuffer;
+extern struct sStringExpansionPos *pCurrentStringExpansion;
 
 void upperstring(char *s);
 void lowerstring(char *s);

src/asm/fstack.c

@@ -28,6 +28,8 @@
 #include "types.h"
 
 static struct sContext *pFileStack;
+static unsigned int nFileStackDepth;
+unsigned int nMaxRecursionDepth;
 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 > nMaxRecursionDepth)
+		fatalerror("Recursion limit (%d) exceeded", nMaxRecursionDepth);
+
 	ppFileStack = &pFileStack;
 	while (*ppFileStack)
 		ppFileStack = &((*ppFileStack)->pNext);
@@ -154,6 +161,8 @@ static int32_t popcontext(void)
 		fatalerror("%s: Internal error.", __func__);
 	}
 
+	nFileStackDepth--;
+
 	free(*ppLastFile);
 	*ppLastFile = NULL;
 	yy_switch_to_buffer(CurrentFlexHandle);
@@ -413,6 +422,7 @@ void fstk_Init(char *pFileName)
 		if (pCurrentFile == NULL)
 			err(1, "Unable to open file '%s'", pFileName);
 	}
+	nFileStackDepth = 0;
 
 	nMacroCount = 0;
 	nCurrentStatus = STAT_isInclude;

src/asm/globlex.c

@@ -188,11 +188,11 @@ uint32_t ParseNumber(char *s, uint32_t size)
 }
 
 /*
- * If the symbol name ends before the end of the macro arg, return true
- * and point "rest" to the rest of the macro arg.
- * Otherwise, return false.
+ * If the symbol name ends before the end of the macro arg,
+ * return a pointer to the rest of the macro arg.
+ * Otherwise, return NULL.
  */
-bool AppendMacroArg(char whichArg, char *dest, size_t *destIndex, char **rest)
+char *AppendMacroArg(char whichArg, char *dest, size_t *destIndex)
 {
 	char *marg;
 
@@ -222,14 +222,13 @@ bool AppendMacroArg(char whichArg, char *dest, size_t *destIndex, char **rest)
 			dest[*destIndex] = ch;
 			(*destIndex)++;
 		} else {
-			*rest = marg;
-			return true;
+			return marg;
 		}
 
 		marg++;
 	}
 
-	return false;
+	return NULL;
 }
 
 uint32_t ParseSymbol(char *src, uint32_t size)
@@ -251,7 +250,9 @@ uint32_t ParseSymbol(char *src, uint32_t size)
 			 */
 			ch = src[srcIndex++];
 
-			if (AppendMacroArg(ch, dest, &destIndex, &rest))
+			rest = AppendMacroArg(ch, dest, &destIndex);
+			/* If the symbol's end was in the middle of the token */
+			if (rest)
 				break;
 		} else {
 			if (destIndex >= MAXSYMLEN)
@@ -262,28 +263,35 @@ uint32_t ParseSymbol(char *src, uint32_t size)
 
 	dest[destIndex] = 0;
 
+	/* Tell the lexer we read all bytes that we did */
+	yyskipbytes(srcIndex);
+
+	/*
+	 * If an escape's expansion left some chars after the symbol's end,
+	 * such as the `::` in a `Backup\1` expanded to `BackupCamX::`,
+	 * put those into the buffer.
+	 * Note that this NEEDS to be done after the `yyskipbytes` above.
+	 */
+	if (rest)
+		yyunputstr(rest);
+
+	/* If the symbol is an EQUS, expand it */
 	if (!oDontExpandStrings && sym_isString(dest)) {
 		char *s;
 
-		yyskipbytes(srcIndex);
-
-		if (rest)
-			yyunputstr(rest);
+		lex_BeginStringExpansion(dest);
 
+		/* Feed the symbol's contents into the buffer */
 		yyunputstr(s = sym_GetStringValue(dest));
 
+		/* Lines inserted this way shall not increase nLineNo */
 		while (*s) {
 			if (*s++ == '\n')
-				nLineNo -= 1;
+				nLineNo--;
 		}
 		return 0;
 	}
 
-	yyskipbytes(srcIndex);
-
-	if (rest)
-		yyunputstr(rest);
-
 	strcpy(yylval.tzSym, dest);
 	return 1;
 }

src/asm/lexer.c

@@ -51,6 +51,9 @@ uint32_t tFloatingChars[256];
 uint32_t nFloating;
 enum eLexerState lexerstate = LEX_STATE_NORMAL;
 
+struct sStringExpansionPos *pCurrentStringExpansion;
+static unsigned int nNbStringExpansions;
+
 /* UTF-8 byte order mark */
 static const unsigned char bom[BOM_SIZE] = { 0xEF, 0xBB, 0xBF };
 
@@ -88,17 +91,49 @@ void yyunput(char c)
 	*(--pLexBuffer) = c;
 }
 
-void yyunputstr(char *s)
+void yyunputstr(const char *s)
 {
-	int32_t i, len;
+	int32_t len;
 
 	len = strlen(s);
 
-	if (pLexBuffer - len < pLexBufferRealStart)
+	/*
+	 * It would be undefined behavior to subtract `len` from pLexBuffer and
+	 * potentially have it point outside of pLexBufferRealStart's buffer,
+	 * this is why the check is done this way.
+	 * Refer to https://github.com/rednex/rgbds/pull/411#discussion_r319779797
+	 */
+	if (pLexBuffer - pLexBufferRealStart < len)
 		fatalerror("Buffer safety margin exceeded");
 
-	for (i = len - 1; i >= 0; i--)
-		*(--pLexBuffer) = s[i];
+	pLexBuffer -= len;
+
+	memcpy(pLexBuffer, s, len);
+}
+
+/*
+ * Marks that a new string expansion with name `tzName` ends here
+ * Enforces recursion depth
+ */
+void lex_BeginStringExpansion(const char *tzName)
+{
+	if (++nNbStringExpansions > nMaxRecursionDepth)
+		fatalerror("Recursion limit (%d) exceeded", nMaxRecursionDepth);
+
+	struct sStringExpansionPos *pNewStringExpansion =
+		malloc(sizeof(*pNewStringExpansion));
+	char *tzNewExpansionName = strdup(tzName);
+
+	if (!pNewStringExpansion || !tzNewExpansionName)
+		fatalerror("Could not allocate memory to expand '%s'",
+			   tzName);
+
+	pNewStringExpansion->tzName = tzNewExpansionName;
+	pNewStringExpansion->pBuffer = pLexBufferRealStart;
+	pNewStringExpansion->pBufferPos = pLexBuffer;
+	pNewStringExpansion->pParent = pCurrentStringExpansion;
+
+	pCurrentStringExpansion = pNewStringExpansion;
 }
 
 void yy_switch_to_buffer(YY_BUFFER_STATE buf)
@@ -423,6 +458,9 @@ void lex_Init(void)
 
 	nLexMaxLength = 0;
 	nFloating = 0;
+
+	pCurrentStringExpansion = NULL;
+	nNbStringExpansions = 0;
 }
 
 void lex_AddStrings(const struct sLexInitString *lex)
@@ -967,12 +1005,30 @@ static uint32_t yylex_MACROARGS(void)
 
 int yylex(void)
 {
+	int returnedChar;
 	switch (lexerstate) {
 	case LEX_STATE_NORMAL:
-		return yylex_NORMAL();
+		returnedChar = yylex_NORMAL();
+		break;
 	case LEX_STATE_MACROARGS:
-		return yylex_MACROARGS();
+		returnedChar = yylex_MACROARGS();
+		break;
 	default:
 		fatalerror("%s: Internal error.", __func__);
 	}
+
+	/* Check if string expansions were fully read */
+	while (pCurrentStringExpansion
+	    && pCurrentStringExpansion->pBuffer == pLexBufferRealStart
+	    && pCurrentStringExpansion->pBufferPos <= pLexBuffer) {
+		struct sStringExpansionPos *pParent =
+			pCurrentStringExpansion->pParent;
+		free(pCurrentStringExpansion->tzName);
+		free(pCurrentStringExpansion);
+
+		pCurrentStringExpansion = pParent;
+		nNbStringExpansions--;
+	}
+
+	return returnedChar;
 }

src/asm/main.c

@@ -289,7 +289,8 @@ static void print_usage(void)
 {
 	printf(
 "usage: rgbasm [-EhLVvw] [-b chars] [-Dname[=value]] [-g chars] [-i path]\n"
-"              [-M dependfile] [-o outfile] [-p pad_value] file.asm\n");
+"              [-M dependfile] [-o outfile] [-p pad_value]\n"
+"              [-r recursion_depth] file.asm\n");
 	exit(1);
 }
 
@@ -316,6 +317,8 @@ int main(int argc, char *argv[])
 
 	/* yydebug=1; */
 
+	nMaxRecursionDepth = 64;
+
 	DefaultOptions.gbgfx[0] = '0';
 	DefaultOptions.gbgfx[1] = '1';
 	DefaultOptions.gbgfx[2] = '2';
@@ -333,7 +336,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) {
@@ -387,6 +390,12 @@ int main(int argc, char *argv[])
 				errx(1, "Argument for option 'p' must be between 0 and 0xFF");
 
 			break;
+		case 'r':
+			nMaxRecursionDepth = 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);

src/asm/rgbasm.1

@@ -21,6 +21,7 @@
 .Op Fl M Ar dependfile
 .Op Fl o Ar outfile
 .Op Fl p Ar pad_value
+.Op Fl r Ar recursion_depth
 .Ar file
 .Sh DESCRIPTION
 The
@@ -77,6 +78,8 @@ Write an object file to the given filename.
 .It Fl p Ar pad_value
 When padding an image, pad with this value.
 The default is 0x00.
+.It Fl r Ar recursion_depth
+Specifies the recursion depth at which RGBASM will assume being in an infinite loop.
 .It Fl V
 Print the version of the program and exit.
 .It Fl v

src/asm/rgbasm.5

@@ -455,8 +455,7 @@ String equates can't be exported or imported.
 .Sy Important note :
 An EQUS can be expanded to a string that contains another EQUS
 and it will be expanded as well.
-This means that, if you aren't careful, you may trap the assembler into an
-infinite loop if there's a circular dependency in the expansions.
+If this creates an infinite loop, RGBASM will error out once a certain depth is reached. See the -r command-line option.
 Also, a MACRO can have inside an EQUS which references the same MACRO, which has
 the same problem.
 .Pp

test/asm/equs-recursion.asm

@@ -0,0 +1,2 @@
+recurse EQUS "recurse"
+recurse

test/asm/equs-recursion.out

@@ -0,0 +1,2 @@
+ERROR: equs-recursion.asm(2):
+    Recursion limit (64) exceeded

test/asm/equs-recursion.out.pipe

@@ -0,0 +1,2 @@
+ERROR: -(2):
+    Recursion limit (64) exceeded

test/asm/include-recursion.asm

@@ -0,0 +1 @@
+INCLUDE "include-recursion.asm"

test/asm/include-recursion.out

@@ -0,0 +1,2 @@
+ERROR: include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1):
+    Recursion limit (64) exceeded

test/asm/include-recursion.out.pipe

@@ -0,0 +1,2 @@
+ERROR: -(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1) -> include-recursion.asm(1):
+    Recursion limit (64) exceeded

test/asm/macro-recursion.asm

@@ -0,0 +1,4 @@
+recurse: MACRO
+	recurse
+ENDM
+	recurse

test/asm/macro-recursion.out

@@ -0,0 +1,2 @@
+ERROR: macro-recursion.asm(4) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1):
+    Recursion limit (64) exceeded

test/asm/macro-recursion.out.pipe

@@ -0,0 +1,2 @@
+ERROR: -(4) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1) -> recurse(1):
+    Recursion limit (64) exceeded