Cleanup and fixes

4 files changed, 140 insertions, 96 deletions

diff --git a/README.md b/README.md
index d56be50..0fd1aa4 100644
--- a/README.md
+++ b/README.md
@@ -158,10 +158,11 @@ if (content) {
     FREE(content); // Standard free (unless using arena)
 }
 
-// Or read directly into a string view or builder
-stringv file_sv = read_entire_file_sv("config.ini");
-// ... use file_sv ...
-FREE((void*)file_sv.data); // Note: read_entire_file_sv allocates the data
+// Or read directly into a string builder
+stringb file_sb = read_entire_file_sb("config.ini");
+// ... use file_sb ...
+sb_free(&file_sb);
+
 ```
 
 ## Testing
diff --git a/examples/files.c b/examples/files.c
index 663bbf9..2e637de 100644
--- a/examples/files.c
+++ b/examples/files.c
@@ -23,12 +23,6 @@ int main(void) {
 		printf("Written test_sv.txt\n");
 	}
 
-	stringv sv_read = read_entire_file_sv("test_sv.txt");
-	if (sv_read.data) {
-		printf("Read sv: %.*s", (int)sv_read.length, sv_read.data);
-		free((char *)sv_read.data);
-	}
-
 	// 3. usage with string builder
 	stringb sb;
 	sb_init(&sb, 0);
diff --git a/nonstd.h b/nonstd.h
index 631f7f1..57e8da4 100644
--- a/nonstd.h
+++ b/nonstd.h
@@ -35,14 +35,17 @@ typedef char c8;
 	for (size_t _i_##var = 0, _n_##var = countof(array); \
 		 _i_##var < _n_##var && ((var) = (array)[_i_##var], 1); ++_i_##var)
 
-#define ALLOC(type, n) ((type *)malloc((n) * sizeof(type)))
-#define REALLOC(ptr, type, n) ((type *)realloc((ptr), (n) * sizeof(type)))
+#define ALLOC(type, n) ((type *)safe_malloc(sizeof(type), (n)))
+#define REALLOC(ptr, type, n) ((type *)safe_realloc((ptr), sizeof(type), (n)))
 #define FREE(ptr)   \
 	do {            \
 		free(ptr);  \
 		ptr = NULL; \
 	} while (0)
 
+NONSTD_DEF void *safe_malloc(size_t item_size, size_t count);
+NONSTD_DEF void *safe_realloc(void *ptr, size_t item_size, size_t count);
+
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 #define CLAMP(x, lo, hi) (MIN((hi), MAX((lo), (x))))
@@ -131,24 +134,36 @@ NONSTD_DEF stringv sb_as_sv(const stringb *sb);
 		(arr).capacity = 0; \
 	} while (0)
 
-#define array_ensure(arr, additional)                                   \
-	do {                                                                \
-		size_t _needed = (arr).length + (additional);                   \
-		if (_needed > (arr).capacity) {                                 \
-			size_t _new_cap = (arr).capacity ? (arr).capacity : 16;     \
-			while (_new_cap < _needed) {                                \
-				_new_cap *= 2;                                          \
-			}                                                           \
-			(arr).data =                                                \
-				REALLOC((arr).data, __typeof__(*(arr).data), _new_cap); \
-			(arr).capacity = _new_cap;                                  \
-		}                                                               \
+#define array_ensure(arr, additional)                                    \
+	do {                                                                 \
+		size_t _needed = (arr).length + (additional);                    \
+		if (_needed > (arr).capacity) {                                  \
+			size_t _new_cap = (arr).capacity ? (arr).capacity : 16;      \
+			while (_new_cap < _needed) {                                 \
+				if (_new_cap > SIZE_MAX / 2) {                           \
+					_new_cap = SIZE_MAX;                                 \
+					break;                                               \
+				}                                                        \
+				_new_cap *= 2;                                           \
+			}                                                            \
+			if (_new_cap < _needed) { /* Overflow or OOM */              \
+				break;                                                   \
+			}                                                            \
+			void *_new_data =                                            \
+				safe_realloc((arr).data, sizeof(*(arr).data), _new_cap); \
+			if (_new_data) {                                             \
+				(arr).data = _new_data;                                  \
+				(arr).capacity = _new_cap;                               \
+			}                                                            \
+		}                                                                \
 	} while (0)
 
-#define array_push(arr, value)                \
-	do {                                      \
-		array_ensure((arr), 1);               \
-		(arr).data[(arr).length++] = (value); \
+#define array_push(arr, value)                    \
+	do {                                          \
+		array_ensure((arr), 1);                   \
+		if ((arr).length < (arr).capacity) {      \
+			(arr).data[(arr).length++] = (value); \
+		}                                         \
 	} while (0)
 
 #define array_pop(arr) ((arr).length > 0 ? (arr).data[--(arr).length] : 0)
@@ -162,16 +177,18 @@ NONSTD_DEF stringv sb_as_sv(const stringb *sb);
 		}                                \
 	} while (0)
 
-#define array_insert(arr, index, value)                          \
-	do {                                                         \
-		if ((index) <= (arr).length) {                           \
-			array_ensure((arr), 1);                              \
-			for (size_t _i = (arr).length; _i > (index); --_i) { \
-				(arr).data[_i] = (arr).data[_i - 1];             \
-			}                                                    \
-			(arr).data[index] = (value);                         \
-			(arr).length++;                                      \
-		}                                                        \
+#define array_insert(arr, index, value)                              \
+	do {                                                             \
+		if ((index) <= (arr).length) {                               \
+			array_ensure((arr), 1);                                  \
+			if ((arr).length < (arr).capacity) {                     \
+				for (size_t _i = (arr).length; _i > (index); --_i) { \
+					(arr).data[_i] = (arr).data[_i - 1];             \
+				}                                                    \
+				(arr).data[index] = (value);                         \
+				(arr).length++;                                      \
+			}                                                        \
+		}                                                            \
 	} while (0)
 
 #define array_remove(arr, index)                                     \
@@ -189,13 +206,16 @@ NONSTD_DEF stringv sb_as_sv(const stringb *sb);
 		(arr).length = 0; \
 	} while (0)
 
-#define array_reserve(arr, new_capacity)                                      \
-	do {                                                                      \
-		if ((new_capacity) > (arr).capacity) {                                \
-			(arr).data =                                                      \
-				REALLOC((arr).data, __typeof__(*(arr).data), (new_capacity)); \
-			(arr).capacity = (new_capacity);                                  \
-		}                                                                     \
+#define array_reserve(arr, new_capacity)                                       \
+	do {                                                                       \
+		if ((new_capacity) > (arr).capacity) {                                 \
+			void *_new_data =                                                  \
+				safe_realloc((arr).data, sizeof(*(arr).data), (new_capacity)); \
+			if (_new_data) {                                                   \
+				(arr).data = _new_data;                                        \
+				(arr).capacity = (new_capacity);                               \
+			}                                                                  \
+		}                                                                      \
 	} while (0)
 
 #define array_foreach(arr, var)                                        \
@@ -204,7 +224,7 @@ NONSTD_DEF stringv sb_as_sv(const stringb *sb);
 		 ++_i_##var)
 
 #define array_foreach_idx(arr, var, index) \
-	for (size_t index = 0;               \
+	for (size_t index = 0;                 \
 		 index < (arr).length && ((var) = (arr).data[index], 1); ++index)
 
 // Arena - block-based memory allocator
@@ -226,7 +246,7 @@ NONSTD_DEF void arena_free(Arena *a);
 // File I/O helpers
 NONSTD_DEF char *read_entire_file(const char *filepath, size_t *out_size);
 NONSTD_DEF int write_entire_file(const char *filepath, const void *data, size_t size);
-NONSTD_DEF stringv read_entire_file_sv(const char *filepath);
+// read_entire_file_sv removed for security (ownership confusion)
 NONSTD_DEF stringb read_entire_file_sb(const char *filepath);
 NONSTD_DEF int write_file_sv(const char *filepath, stringv sv);
 NONSTD_DEF int write_file_sb(const char *filepath, const stringb *sb);
@@ -235,6 +255,20 @@ NONSTD_DEF int write_file_sb(const char *filepath, const stringb *sb);
 
 #ifdef NONSTD_IMPLEMENTATION
 
+NONSTD_DEF void *safe_malloc(size_t item_size, size_t count) {
+	if (count != 0 && item_size > SIZE_MAX / count) {
+		return NULL;
+	}
+	return malloc(item_size * count);
+}
+
+NONSTD_DEF void *safe_realloc(void *ptr, size_t item_size, size_t count) {
+	if (count != 0 && item_size > SIZE_MAX / count) {
+		return NULL;
+	}
+	return realloc(ptr, item_size * count);
+}
+
 NONSTD_DEF stringv sv_from_cstr(const char *s) {
 	return (stringv){.data = s, .length = s ? strlen(s) : 0};
 }
@@ -285,11 +319,24 @@ NONSTD_DEF void sb_free(stringb *sb) {
 
 NONSTD_DEF void sb_ensure(stringb *sb, size_t additional) {
 	size_t needed = sb->length + additional + 1;
-	if (needed > sb->capacity) {
-		while (sb->capacity < needed) {
-			sb->capacity *= 2;
+	size_t new_cap = sb->capacity;
+
+	if (needed > new_cap) {
+		while (new_cap < needed) {
+			if (new_cap > SIZE_MAX / 2) {
+				new_cap = SIZE_MAX;
+				break;
+			}
+			new_cap *= 2;
+		}
+		if (new_cap < needed)
+			return; // Overflow
+
+		char *new_data = safe_realloc(sb->data, sizeof(char), new_cap);
+		if (new_data) {
+			sb->data = new_data;
+			sb->capacity = new_cap;
 		}
-		sb->data = REALLOC(sb->data, char, sb->capacity);
 	}
 }
 
@@ -299,9 +346,11 @@ NONSTD_DEF void sb_append_cstr(stringb *sb, const char *s) {
 	}
 	size_t slength = strlen(s);
 	sb_ensure(sb, slength);
-	memcpy(sb->data + sb->length, s, slength);
-	sb->length += slength;
-	sb->data[sb->length] = '\0';
+	if (sb->length + slength + 1 <= sb->capacity) {
+		memcpy(sb->data + sb->length, s, slength);
+		sb->length += slength;
+		sb->data[sb->length] = '\0';
+	}
 }
 
 NONSTD_DEF void sb_append_sv(stringb *sb, stringv sv) {
@@ -309,15 +358,19 @@ NONSTD_DEF void sb_append_sv(stringb *sb, stringv sv) {
 		return;
 	}
 	sb_ensure(sb, sv.length);
-	memcpy(sb->data + sb->length, sv.data, sv.length);
-	sb->length += sv.length;
-	sb->data[sb->length] = '\0';
+	if (sb->length + sv.length + 1 <= sb->capacity) {
+		memcpy(sb->data + sb->length, sv.data, sv.length);
+		sb->length += sv.length;
+		sb->data[sb->length] = '\0';
+	}
 }
 
 NONSTD_DEF void sb_append_char(stringb *sb, char c) {
 	sb_ensure(sb, 1);
-	sb->data[sb->length++] = c;
-	sb->data[sb->length] = '\0';
+	if (sb->length + 2 <= sb->capacity) {
+		sb->data[sb->length++] = c;
+		sb->data[sb->length] = '\0';
+	}
 }
 
 NONSTD_DEF stringv sb_as_sv(const stringb *sb) {
@@ -343,12 +396,20 @@ NONSTD_DEF void *arena_alloc(Arena *a, size_t size) {
 	size_t align = sizeof(void *);
 	uintptr_t current = (uintptr_t)a->ptr;
 	uintptr_t aligned = (current + align - 1) & ~(align - 1);
-	uintptr_t available = (uintptr_t)a->end - aligned;
+	uintptr_t end = (uintptr_t)a->end;
 
-	if (available < size) {
+	// Check for overflow (aligned wrapped around) or out of bounds (aligned >= end)
+	// or not enough space ((end - aligned) < size)
+	if (aligned < current || aligned >= end || (end - aligned) < size) {
 		arena_grow(a, size);
 		current = (uintptr_t)a->ptr;
 		aligned = (current + align - 1) & ~(align - 1);
+		end = (uintptr_t)a->end;
+	}
+
+	// Double check after grow (in case grow failed or size is just too huge)
+	if (aligned < current || aligned >= end || (end - aligned) < size) {
+		return NULL;
 	}
 
 	a->ptr = (char *)(aligned + size);
@@ -413,14 +474,7 @@ NONSTD_DEF int write_entire_file(const char *filepath, const void *data, size_t
 	return written == size;
 }
 
-NONSTD_DEF stringv read_entire_file_sv(const char *filepath) {
-	size_t size = 0;
-	char *data = read_entire_file(filepath, &size);
-	if (!data) {
-		return (stringv){0};
-	}
-	return (stringv){.data = data, .length = size};
-}
+// read_entire_file_sv removed
 
 NONSTD_DEF stringb read_entire_file_sb(const char *filepath) {
 	size_t size = 0;
diff --git a/tests.c b/tests.c
index e865e18..32da92f 100644
--- a/tests.c
+++ b/tests.c
@@ -784,6 +784,27 @@ MU_TEST(test_arena_alignment) {
 	arena_free(&a);
 }
 
+MU_TEST(test_arena_safety) {
+	Arena a = arena_make();
+
+	// Test 1: Integer overflow in allocation should be caught by safe_alloc
+	// Requesting 2 items of size (SIZE_MAX/2 + 1) should fail
+	// This tests ns_safe_malloc directly or via macro if we exposed it,
+	// but arena_alloc calls ensure/grow which calls ALLOC(char, size).
+	// To test ALLOC overflow: ALLOC(char, SIZE_MAX) is just 1*SIZE_MAX.
+	// We need a way to trigger safe_alloc failure.
+
+	// Let's test the underlying safe allocator directly since we exposed it in nonstd.h
+	void *p = safe_malloc(SIZE_MAX / 2 + 100, 2);
+	mu_check(p == NULL);
+
+	// Test 2: Ensure normal allocations still work
+	void *p2 = arena_alloc(&a, 100);
+	mu_check(p2 != NULL);
+
+	arena_free(&a);
+}
+
 // File I/O tests
 MU_TEST(test_file_io_basic) {
 	const char *filename = "test_io_basic.txt";
@@ -804,25 +825,6 @@ MU_TEST(test_file_io_basic) {
 	remove(filename);
 }
 
-MU_TEST(test_file_io_sv) {
-	const char *filename = "test_io_sv.txt";
-	stringv sv = sv_from_cstr("Hello from sv!");
-
-	// Write
-	mu_check(write_file_sv(filename, sv));
-
-	// Read
-	stringv read_sv = read_entire_file_sv(filename);
-	mu_check(read_sv.data != NULL);
-	mu_assert_int_eq(sv.length, read_sv.length);
-	mu_check(sv_equals(sv, read_sv));
-
-	// The data returned by read_entire_file_sv is malloc'd and needs freeing
-	// We cast away const to free it since we know it was allocated
-	free((char *)read_sv.data);
-	remove(filename);
-}
-
 MU_TEST(test_file_io_sb) {
 	const char *filename = "test_io_sb.txt";
 	stringb sb;
@@ -849,12 +851,6 @@ MU_TEST(test_file_io_read_missing) {
 	mu_check(content == NULL);
 }
 
-MU_TEST(test_file_io_read_missing_sv) {
-	stringv sv = read_entire_file_sv("non_existent_file.txt");
-	mu_check(sv.data == NULL);
-	mu_assert_int_eq(0, sv.length);
-}
-
 MU_TEST(test_file_io_read_missing_sb) {
 	stringb sb = read_entire_file_sb("non_existent_file.txt");
 	mu_check(sb.data == NULL);
@@ -953,15 +949,14 @@ MU_TEST_SUITE(test_suite_arena) {
 	RUN_TEST_WITH_NAME(test_arena_basic);
 	RUN_TEST_WITH_NAME(test_arena_growth);
 	RUN_TEST_WITH_NAME(test_arena_alignment);
+	RUN_TEST_WITH_NAME(test_arena_safety);
 }
 
 MU_TEST_SUITE(test_suite_files) {
 	printf("\n[File I/O Tests]\n");
 	RUN_TEST_WITH_NAME(test_file_io_basic);
-	RUN_TEST_WITH_NAME(test_file_io_sv);
 	RUN_TEST_WITH_NAME(test_file_io_sb);
 	RUN_TEST_WITH_NAME(test_file_io_read_missing);
-	RUN_TEST_WITH_NAME(test_file_io_read_missing_sv);
 	RUN_TEST_WITH_NAME(test_file_io_read_missing_sb);
 }