From cd072d5e6a9031ae326e9d95c7680bde94a88d81 Mon Sep 17 00:00:00 2001
From: ISSOtm <eldredhabert0@gmail.com>
Date: Fri, 19 Feb 2021 16:52:35 +0100
Subject: [PATCH] Add assertion check for potential UAF trigger

It actually currently triggers if an INCLUDE directive is given at EOF,
but #742 will fix that.
---
 src/asm/fstack.c |  3 ++-
 src/asm/lexer.c  | 21 ++++++++++++++++-----
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/asm/fstack.c b/src/asm/fstack.c
index bfd10171..a938f6bb 100644
--- a/src/asm/fstack.c
+++ b/src/asm/fstack.c
@@ -24,7 +24,7 @@
 #define MAXINCPATHS 128
 
 #ifdef LEXER_DEBUG
-  #define dbgPrint(...) fprintf(stderr, "[lexer] " __VA_ARGS__)
+  #define dbgPrint(...) fprintf(stderr, "[fstack] " __VA_ARGS__)
 #else
   #define dbgPrint(...)
 #endif
@@ -270,6 +270,7 @@ bool yywrap(void)
 	struct Context *context = contextStack;
 
 	contextStack = contextStack->parent;
+	assert(contextDepth != 0); // This is never supposed to underflow
 	contextDepth--;
 
 	lexer_DeleteState(context->lexerState);
diff --git a/src/asm/lexer.c b/src/asm/lexer.c
index 9ff949e5..138471ce 100644
--- a/src/asm/lexer.c
+++ b/src/asm/lexer.c
@@ -479,6 +479,19 @@ void lexer_RestartRept(uint32_t lineNo)
 
 void lexer_DeleteState(struct LexerState *state)
 {
+	// A big chunk of the lexer state soundness is the file stack ("fstack").
+	// Each context in the fstack has its own *unique* lexer state; thus, we always guarantee
+	// that lexer states lifetimes are always properly managed, since they're handled solely
+	// by the fstack... with *one* exception.
+	// Assume a context is pushed on top of the fstack, and the corresponding lexer state gets
+	// scheduled at EOF; `lexerStateAtEOL` thus becomes a (weak) ref to that lexer state...
+	// It has been possible, due to a bug, that the corresponding fstack context gets popped
+	// before EOL, deleting the associated state... but it would still be switched to at EOL.
+	// This assertion checks that this doesn't happen again.
+	// It could be argued that deleting a state that's scheduled for EOF could simply clear
+	// `lexerStateEOL`, but there's currently no situation in which this should happen.
+	assert(state != lexerStateEOL);
+
 	if (!state->isMmapped)
 		close(state->fd);
 	else if (state->isFile && !state->isReferenced)
@@ -2329,13 +2342,13 @@ restart:
 		[LEXER_RAW]          = yylex_RAW,
 		[LEXER_SKIP_TO_ELIF] = yylex_SKIP_TO_ELIF,
 		[LEXER_SKIP_TO_ENDC] = yylex_SKIP_TO_ENDC,
-		[LEXER_SKIP_TO_ENDR] = yylex_SKIP_TO_ENDR
+		[LEXER_SKIP_TO_ENDR] = yylex_SKIP_TO_ENDR,
 	};
 	int token = lexerModeFuncs[lexerState->mode]();
 
 	if (token == T_EOF) {
 		/* Try to switch to new buffer; if it succeeds, scan again */
-		dbgPrint("Reached EOF!\n");
+		dbgPrint("Reached EOB!\n");
 		/* Captures end at their buffer's boundary no matter what */
 		if (!lexerState->capturing) {
 			if (!yywrap())
@@ -2345,9 +2358,7 @@ restart:
 		}
 	}
 
-	lexerState->atLineStart = false;
-	if (token == T_NEWLINE)
-		lexerState->atLineStart = true;
+	lexerState->atLineStart = token == T_NEWLINE;
 
 	return token;
 }