From 05b5f187230d4d1c9853269ba3d9c7d8a72d6086 Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Wed, 1 Sep 2021 07:00:50 +0100
Subject: [PATCH] Check for size_max on 32 bit

---
 c/tests/test_core.c | 16 +++++++++++++++-
 c/tskit/core.c      | 18 +++++++++++++++---
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/c/tests/test_core.c b/c/tests/test_core.c
index 6376599068..8004be3122 100644
--- a/c/tests/test_core.c
+++ b/c/tests/test_core.c
@@ -1,7 +1,7 @@
 /*
  * MIT License
  *
- * Copyright (c) 2019 Tskit Developers
+ * Copyright (c) 2019-2021 Tskit Developers
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -358,6 +358,19 @@ test_malloc_zero(void)
     free(p);
 }
 
+static void
+test_malloc_overflow(void)
+{
+#if TSK_MAX_SIZE > SIZE_MAX
+    tsk_size_t size_max = SIZE_MAX;
+    void *p = tsk_malloc(size_max + 1);
+    CU_ASSERT_FATAL(p == NULL);
+
+    p = tsk_calloc(size_max + 1, 1);
+    CU_ASSERT_FATAL(p == NULL);
+#endif
+}
+
 int
 main(int argc, char **argv)
 {
@@ -369,6 +382,7 @@ main(int argc, char **argv)
         { "test_blkalloc", test_blkalloc },
         { "test_unknown_time", test_unknown_time },
         { "test_malloc_zero", test_malloc_zero },
+        { "test_malloc_overflow", test_malloc_overflow },
         { NULL, NULL },
     };
 
diff --git a/c/tskit/core.c b/c/tskit/core.c
index fa2daff909..c71952cc29 100644
--- a/c/tskit/core.c
+++ b/c/tskit/core.c
@@ -717,9 +717,11 @@ tsk_malloc(tsk_size_t size)
     if (size == 0) {
         size = 1;
     }
-    /* TODO
-     * 1. check if size > SIZE_MAX on 32 bit
-     */
+#if TSK_MAX_SIZE > SIZE_MAX
+    if (size > SIZE_MAX) {
+        return NULL;
+    }
+#endif
     return malloc((size_t) size);
 }
 
@@ -731,6 +733,11 @@ tsk_realloc(void *ptr, tsk_size_t size)
     return realloc(ptr, (size_t) size);
 }
 
+/* We keep the size argument here as a size_t because we'd have to
+ * cast the outputs of sizeof() otherwise, which would lead to
+ * less readable code. We need to be careful to use calloc within
+ * the library accordingly, so that size can't overflow on 32 bit.
+ */
 void *
 tsk_calloc(tsk_size_t n, size_t size)
 {
@@ -738,6 +745,11 @@ tsk_calloc(tsk_size_t n, size_t size)
     if (n == 0) {
         n = 1;
     }
+#if TSK_MAX_SIZE > SIZE_MAX
+    if (n > SIZE_MAX) {
+        return NULL;
+    }
+#endif
     return calloc((size_t) n, size);
 }