From fa18a75dc55c58505b7c643560bc7d8f1198995b Mon Sep 17 00:00:00 2001
From: Rangi <remy.oukaour+rangi42@gmail.com>
Date: Tue, 15 Aug 2023 19:26:41 -0400
Subject: [PATCH] Fix segfaults at EOF, add token/newline checks, and support
 '\r' in scan_includes.c

---
 tools/scan_includes.c | 58 ++++++++++++++++++++++++++++++-------------
 1 file changed, 41 insertions(+), 17 deletions(-)

diff --git a/tools/scan_includes.c b/tools/scan_includes.c
index e57ddc35..9ec5f3c5 100644
--- a/tools/scan_includes.c
+++ b/tools/scan_includes.c
@@ -3,6 +3,8 @@
 
 #include "common.h"
 
+#include <ctype.h>
+
 void parse_args(int argc, char *argv[], bool *strict) {
 	struct option long_options[] = {
 		{"strict", no_argument, 0, 's'},
@@ -40,31 +42,47 @@ void scan_file(const char *filename, bool strict) {
 	fclose(f);
 	contents[size] = '\0';
 
-	for (char *ptr = contents; ptr && ptr - contents < size; ptr++) {
-		bool is_incbin = false, is_include = false;
+	for (char *ptr = contents; ptr && ptr < contents + size; ptr++) {
+		ptr = strpbrk(ptr, ";\"Ii");
+		if (!ptr) {
+			break;
+		}
 		switch (*ptr) {
 		case ';':
-			ptr = strchr(ptr, '\n');
-			if (!ptr) {
-				fprintf(stderr, "%s: no newline at end of file\n", filename);
-			}
-			break;
-		case '"':
-			ptr++;
-			ptr = strchr(ptr, '"');
-			if (ptr) {
+			// Skip comments until the end of the line
+			ptr += strcspn(ptr + 1, "\r\n");
+			if (*ptr) {
 				ptr++;
-			} else {
-				fprintf(stderr, "%s: unterminated string\n", filename);
 			}
 			break;
+
+		case '"':
+			// Skip string literal until the closing quote
+			ptr += strcspn(ptr + 1, "\"");
+			if (*ptr) {
+				ptr++;
+			}
+			break;
+
 		case 'I':
 		case 'i':
-			is_incbin = !strncmp(ptr, "INCBIN", 6) || !strncmp(ptr, "incbin", 6);
-			is_include = !strncmp(ptr, "INCLUDE", 7) || !strncmp(ptr, "include", 7);
+			/* empty statement between the label and the variable declaration */;
+			// Check that an INCLUDE/INCBIN starts as its own token
+			char before = ptr > contents ? *(ptr - 1) : '\n';
+			if (!isspace((unsigned)before) && before != ':') {
+				break;
+			}
+			bool is_incbin = !strncmp(ptr, "INCBIN", 6) || !strncmp(ptr, "incbin", 6);
+			bool is_include = !strncmp(ptr, "INCLUDE", 7) || !strncmp(ptr, "include", 7);
 			if (is_incbin || is_include) {
-				ptr = strchr(ptr, '"');
-				if (ptr) {
+				// Check that an INCLUDE/INCBIN ends as its own token
+				ptr += is_include ? 7 : 6;
+				if (!isspace((unsigned)*ptr) && *ptr != '"') {
+					break;
+				}
+				ptr += strspn(ptr, " \t");
+				if (*ptr == '"') {
+					// Print the file path and recursively scan INCLUDEs
 					ptr++;
 					char *include_path = ptr;
 					size_t length = strcspn(ptr, "\"");
@@ -74,6 +92,12 @@ void scan_file(const char *filename, bool strict) {
 					if (is_include) {
 						scan_file(include_path, strict);
 					}
+				} else {
+					fprintf(stderr, "%s: no file path after INC%s\n", filename, is_include ? "LUDE" : "BIN");
+					// Continue to process a comment
+					if (*ptr == ';') {
+						ptr--;
+					}
 				}
 			}
 			break;