Use the same memory-failure mechanism for mallocs

FossilOrigin-Name: fc4f944f1ef599321fac8aa6af72bbe55fb438b4f43800ba37dcfdfa2e6b4c4e

1 files changed, 21 insertions, 13 deletions

diff --git a/src/bin/main.c b/src/bin/main.c
index 481874c..a7acc77 100644
--- a/src/bin/main.c
+++ b/src/bin/main.c
@@ -26,13 +26,31 @@
 static uint64_t alloc_count;
 static long memdbg;
 
+static int
+hook_alloc(void)
+{
+	alloc_count++;
+	if((memdbg >= 0) && (alloc_count > ((uint64_t)memdbg)))
+		return 0;
+	else
+		return 1;
+}
+
+static void *
+my_malloc(size_t size)
+{
+	if(!hook_alloc())
+		return NULL;
+	return malloc(size);
+}
+
 static void *
 my_alloc(size_t min, size_t max, size_t *actual)
 {
 	size_t target = min;
 	void *p;
 
-	if(target == 0)
+	if(NEVER(target == 0))
 		target = 1;
 	else{
 		target--;
@@ -48,10 +66,10 @@ my_alloc(size_t min, size_t max, size_t *actual)
 	if(max && (target > max))
 		target = max;
 
-	p = malloc(target);
+	p = my_malloc(target);
 	if((p == NULL) && (min < target)){
 		target = min;
-		p = malloc(target);
+		p = my_malloc(target);
 	}
 
 	*actual = target;
@@ -64,16 +82,6 @@ my_dealloc(void *p)
 	free(p);
 }
 
-static int
-hook_alloc(void)
-{
-	alloc_count++;
-	if((memdbg >= 0) && (alloc_count > ((uint64_t)memdbg)))
-		return 0;
-	else
-		return 1;
-}
-
 int
 main(int argc, char *argv[])
 {