Make gendata and co safer by checking buffer size

5 files changed, 84 insertions, 69 deletions

diff --git a/src/solvers/h48/gendata_cocsep.h b/src/solvers/h48/gendata_cocsep.h
index 18bfc38..b46bcf3 100644
--- a/src/solvers/h48/gendata_cocsep.h
+++ b/src/solvers/h48/gendata_cocsep.h
@@ -1,7 +1,8 @@
 STATIC_INLINE bool get_visited(const uint8_t *, int64_t);
 STATIC_INLINE void set_visited(uint8_t *, int64_t);
 
-STATIC size_t gendata_cocsep(void *, uint64_t *, cube_t *);
+STATIC size_t gendata_cocsep(
+    char [static COCSEP_FULLSIZE+INFOSIZE], uint64_t *, cube_t *);
 STATIC uint32_t gendata_cocsep_dfs(cocsep_dfs_arg_t *);
 STATIC void getdistribution_cocsep(const uint32_t *, uint64_t [static 21]);
 
@@ -19,7 +20,11 @@ After the data as described above, more auxiliary information is appended:
     of positions having that pruning value.
 */
 STATIC size_t
-gendata_cocsep(void *buf, uint64_t *selfsim, cube_t *rep)
+gendata_cocsep(
+        char buf[static COCSEP_FULLSIZE+INFOSIZE],
+        uint64_t *selfsim,
+        cube_t *rep
+)
 {
         uint32_t *buf32, cc;
         uint16_t n;
@@ -64,7 +69,7 @@ gendata_cocsep(void *buf, uint64_t *selfsim, cube_t *rep)
                 info.distribution[i] = cc;
         }
 
-        writetableinfo(&info, buf);
+        writetableinfo(&info, COCSEP_FULLSIZE+INFOSIZE, buf);
 
         DBG_ASSERT(n == COCSEP_CLASSES, 0,
             "cocsep: computed %" PRIu16 " symmetry classes, "
diff --git a/src/solvers/h48/gendata_h48.h b/src/solvers/h48/gendata_h48.h
index 50ab01c..2a4469f 100644
--- a/src/solvers/h48/gendata_h48.h
+++ b/src/solvers/h48/gendata_h48.h
@@ -1,14 +1,16 @@
 STATIC uint64_t gendata_h48short(gendata_h48short_arg_t *);
-STATIC size_t gendata_h48(gendata_h48_arg_t *);
-STATIC size_t gendata_h48h0k4(gendata_h48_arg_t *);
-STATIC size_t gendata_h48k2(gendata_h48_arg_t *);
+STATIC int64_t gendata_h48(gendata_h48_arg_t *);
+STATIC void gendata_h48h0k4(gendata_h48_arg_t *);
+STATIC void gendata_h48k2(gendata_h48_arg_t *);
+STATIC void gendata_h48k2_realcoord(gendata_h48_arg_t *);
+
 STATIC void * gendata_h48h0k4_runthread(void *);
+STATIC void * gendata_h48k2_runthread(void *);
+
 STATIC_INLINE void gendata_h48_mark_atomic(gendata_h48_mark_t *);
 STATIC_INLINE void gendata_h48_mark(gendata_h48_mark_t *);
 STATIC_INLINE bool gendata_h48k2_dfs_stop(cube_t, int8_t, h48k2_dfs_arg_t *);
-STATIC size_t gendata_h48k2_realcoord(gendata_h48_arg_t *);
 STATIC void gendata_h48k2_dfs(h48k2_dfs_arg_t *arg);
-STATIC void * gendata_h48k2_runthread(void *);
 STATIC tableinfo_t makeinfo_h48k2(gendata_h48_arg_t *);
 STATIC void getdistribution_h48(const uint8_t *,
     uint64_t [static INFO_DISTRIBUTION_LEN], uint8_t, uint8_t);
@@ -62,75 +64,77 @@ gendata_h48short(gendata_h48short_arg_t *arg)
 }
 
 /* Generic function that dispatches to the data generators */
-STATIC size_t
+STATIC int64_t
 gendata_h48(gendata_h48_arg_t *arg)
 {
+        uint64_t size;
         void *cocsepdata_offset;
-        size_t cocsepsize, h48size;
+        size_t cocsepsize;
         tableinfo_t cocsepinfo;
 
-        cocsepsize = gendata_cocsep(arg->buf, arg->selfsim, arg->crep);
+        if (arg == NULL) {
+                LOG("Error computing H48 data: arg is NULL.\n");
+                return NISSY_ERROR_UNKNOWN;
+        }
 
-        if (arg->buf == NULL) {
-                cocsepdata_offset = NULL;
-                arg->cocsepdata = NULL;
-                arg->h48buf = NULL;
-        } else {
-                cocsepdata_offset = (char *)arg->buf + INFOSIZE;
-                arg->cocsepdata = (uint32_t *)cocsepdata_offset;
-                arg->h48buf = (char *)arg->buf + cocsepsize;
+        size = 2*INFOSIZE + COCSEP_FULLSIZE + H48_TABLESIZE(arg->h, arg->k);
+
+        if (arg->buf == NULL)
+                return size; /* Dry-run */
+
+        if (arg->buf_size < size) {
+                LOG("Error computing H48 data: buffer is too small "
+                    "(needed %" PRId64 " bytes but received %" PRId64 ")\n",
+                    size, arg->buf_size);
+                return NISSY_ERROR_BUFFER_SIZE;
         }
 
-        arg->base = 99; // TODO: set this somewhere else
+        cocsepsize = gendata_cocsep(arg->buf, arg->selfsim, arg->crep);
+
+        cocsepdata_offset = (char *)arg->buf + INFOSIZE;
+        arg->cocsepdata = (uint32_t *)cocsepdata_offset;
+        arg->h48buf = (char *)arg->buf + cocsepsize;
+
+        arg->base = 99; /* TODO: set this somewhere else */
 
         if (arg->h == 0 && arg->k == 4) {
-                h48size = gendata_h48h0k4(arg);
+                gendata_h48h0k4(arg);
         } else if ((arg->h == 0 || arg->h == 11) && arg->k == 2) {
-                h48size = gendata_h48k2_realcoord(arg);
+                gendata_h48k2_realcoord(arg);
         } else if (arg->k == 2) {
-                h48size = gendata_h48k2(arg);
+                gendata_h48k2(arg);
         } else {
                 LOG("Cannot generate data for h = %" PRIu8 " and k = %" PRIu8
                     " (not implemented yet)\n", arg->h, arg->k);
-                goto gendata_h48_error;
+                return NISSY_ERROR_INVALID_SOLVER;
         }
 
-        if (arg->buf == NULL)
-                goto gendata_h48_return_size;
-
-        if (!readtableinfo(arg->buf, &cocsepinfo)) {
+        if (readtableinfo(arg->buf_size, arg->buf, &cocsepinfo) != NISSY_OK) {
                 LOG("gendata_h48: could not read info for cocsep table\n");
-                goto gendata_h48_error;
+                return NISSY_ERROR_UNKNOWN;
         }
 
         cocsepinfo.next = cocsepsize;
-        if (!writetableinfo(&cocsepinfo, arg->buf)) {
+        if (writetableinfo(&cocsepinfo, arg->buf_size, arg->buf) != NISSY_OK) {
                 LOG("gendata_h48: could not write info for cocsep table"
                     " with updated 'next' value\n");
-                goto gendata_h48_error;
+                return NISSY_ERROR_UNKNOWN;
         }
 
-gendata_h48_return_size:
-        return cocsepsize + h48size;
-
-gendata_h48_error:
-        return 0;
+        return size;
 }
 
-STATIC size_t
+STATIC void
 gendata_h48h0k4(gendata_h48_arg_t *arg)
 {
         _Atomic uint8_t *table;
         uint8_t val;
         int64_t i, sc, done, d, h48max;
-        uint64_t t, tt, isize, cc;
+        uint64_t t, tt, isize, cc, bufsize;
         h48h0k4_bfs_arg_t bfsarg[THREADS];
         pthread_t thread[THREADS];
         pthread_mutex_t table_mutex[CHUNKS];
 
-        if (arg->buf == NULL)
-                goto gendata_h48h0k4_return_size;
-
         arg->info = (tableinfo_t) {
                 .solver = "h48 solver h = 0, k = 4",
                 .type = TABLETYPE_PRUNING,
@@ -194,10 +198,8 @@ gendata_h48h0k4(gendata_h48_arg_t *arg)
         }
 
         arg->info.maxvalue = d - 1;
-        writetableinfo(&arg->info, arg->h48buf);
-
-gendata_h48h0k4_return_size:
-        return H48_TABLESIZE(0, 4) + INFOSIZE;
+        bufsize = arg->buf_size - COCSEP_FULLSIZE - INFOSIZE;
+        writetableinfo(&arg->info, bufsize, arg->h48buf);
 }
 
 STATIC void *
@@ -258,7 +260,7 @@ gendata_h48h0k4_runthread(void *arg)
         return NULL;
 }
 
-STATIC size_t
+STATIC void
 gendata_h48k2(gendata_h48_arg_t *arg)
 {
         static const uint8_t shortdepth = 8;
@@ -313,19 +315,15 @@ gendata_h48k2(gendata_h48_arg_t *arg)
         uint8_t t;
         uint8_t *table;
         int64_t j;
-        uint64_t i, ii, inext, count;
+        uint64_t i, ii, inext, count, bufsize;
         h48map_t shortcubes;
         gendata_h48short_arg_t shortarg;
         h48k2_dfs_arg_t dfsarg[THREADS];
         pthread_t thread[THREADS];
         pthread_mutex_t shortcubes_mutex, table_mutex[CHUNKS];
 
-        if (arg->buf == NULL)
-                goto gendata_h48k2_return_size;
-
         table = (uint8_t *)arg->h48buf + INFOSIZE;
-        if (arg->buf != NULL)
-                memset(table, 0xFF, H48_TABLESIZE(arg->h, arg->k));
+        memset(table, 0xFF, H48_TABLESIZE(arg->h, arg->k));
 
         LOG("Computing depth <=%" PRIu8 "\n", shortdepth)
         h48map_create(&shortcubes, capacity, randomizer);
@@ -379,10 +377,8 @@ gendata_h48k2(gendata_h48_arg_t *arg)
                 arg->info.distribution[t]++;
         }
 
-        writetableinfo(&arg->info, arg->h48buf);
-
-gendata_h48k2_return_size:
-        return H48_TABLESIZE(arg->h, 2) + INFOSIZE;
+        bufsize = arg->buf_size - COCSEP_FULLSIZE - INFOSIZE;
+        writetableinfo(&arg->info, bufsize, arg->h48buf);
 }
 
 STATIC void *
@@ -565,11 +561,11 @@ gendata_h48k2_dfs_stop(cube_t cube, int8_t depth, h48k2_dfs_arg_t *arg)
         }
 }
 
-STATIC size_t
+STATIC void
 gendata_h48k2_realcoord(gendata_h48_arg_t *arg)
 {
         /* TODO */
-        return gendata_h48k2(arg);
+        gendata_h48k2(arg);
 }
 
 STATIC void *
@@ -684,6 +680,7 @@ gendata_h48_derive(uint8_t h, const void *fulltable, void *buf)
         const uint8_t *h48full;
         uint8_t *h48derive;
         int64_t i, j, h48max;
+        uint64_t bufsize;
         gendata_h48_arg_t arg;
         tableinfo_t cocsepinfo, fulltableinfo;
 
@@ -693,7 +690,9 @@ gendata_h48_derive(uint8_t h, const void *fulltable, void *buf)
         fulltableinfo.bits = 2;
         fulltableinfo.base = 8;
 
-        readtableinfo_n(fulltable, 2, &fulltableinfo);
+        int64_t TODOlarge = 999999999999; /* TODO: cleanup here */
+
+        readtableinfo_n(TODOlarge, fulltable, 2, &fulltableinfo);
         arg.h = h;
         arg.k = fulltableinfo.bits;
         arg.maxdepth = 20;
@@ -711,13 +710,15 @@ gendata_h48_derive(uint8_t h, const void *fulltable, void *buf)
         if (buf == NULL)
                 goto gendata_h48_derive_return_size;
 
-        if (!readtableinfo(buf, &cocsepinfo)) {
+        bufsize = COCSEP_FULLSIZE + INFOSIZE;
+        if (readtableinfo(bufsize, buf, &cocsepinfo) != NISSY_OK) {
                 LOG("gendata_h48: could not read info for cocsep table\n");
                 goto gendata_h48_derive_error;
         }
 
         cocsepinfo.next = cocsepsize;
-        if (!writetableinfo(&cocsepinfo, buf)) {
+        bufsize = COCSEP_FULLSIZE + INFOSIZE;
+        if (writetableinfo(&cocsepinfo, bufsize, buf) != NISSY_OK) {
                 LOG("gendata_h48_derive: could not write info for cocsep table"
                     " with updated 'next' value\n");
                 goto gendata_h48_derive_error;
@@ -742,7 +743,8 @@ gendata_h48_derive(uint8_t h, const void *fulltable, void *buf)
 
         getdistribution_h48(h48derive, arg.info.distribution, h, arg.k);
 
-        if (!writetableinfo(&arg.info, arg.h48buf)) {
+        bufsize = arg.buf_size - COCSEP_FULLSIZE - INFOSIZE;
+        if (writetableinfo(&arg.info, bufsize, arg.h48buf) != NISSY_OK) {
                 LOG("gendata_h48_derive: could not write info for table\n");
                 goto gendata_h48_derive_error;
         }
diff --git a/src/solvers/h48/gendata_types_macros.h b/src/solvers/h48/gendata_types_macros.h
index e0e4df7..73a46ff 100644
--- a/src/solvers/h48/gendata_types_macros.h
+++ b/src/solvers/h48/gendata_types_macros.h
@@ -58,6 +58,7 @@ typedef struct {
         uint8_t base;
         uint8_t maxdepth;
         tableinfo_t info;
+        uint64_t buf_size;
         void *buf;
         void *h48buf;
         uint32_t *cocsepdata;
diff --git a/src/solvers/h48/solve.h b/src/solvers/h48/solve.h
index 3948b15..f2babdf 100644
--- a/src/solvers/h48/solve.h
+++ b/src/solvers/h48/solve.h
@@ -21,7 +21,8 @@ STATIC uint32_t allowednextmove_h48(uint8_t *, uint8_t, uint32_t);
 STATIC void solve_h48_appendsolution(dfsarg_solveh48_t *);
 STATIC_INLINE bool solve_h48_stop(dfsarg_solveh48_t *);
 STATIC int64_t solve_h48_dfs(dfsarg_solveh48_t *);
-STATIC int64_t solve_h48(cube_t, int8_t, int8_t, int8_t, const void *, char *);
+STATIC int64_t solve_h48(cube_t, int8_t, int8_t,
+    int8_t, uint64_t, const void *, uint64_t, char *);
 
 STATIC uint32_t
 allowednextmove_h48(uint8_t *moves, uint8_t n, uint32_t h48branch)
@@ -163,7 +164,9 @@ solve_h48(
         int8_t minmoves,
         int8_t maxmoves,
         int8_t maxsolutions,
+        uint64_t data_size,
         const void *data,
+        uint64_t solutions_size,
         char *solutions
 )
 {
@@ -171,9 +174,9 @@ solve_h48(
         dfsarg_solveh48_t arg;
         tableinfo_t info;
 
-        if(!readtableinfo_n(data, 2, &info)) {
+        if(readtableinfo_n(data_size, data, 2, &info) != NISSY_OK) {
                 LOG("solve_h48: error reading table\n");
-                return 0;
+                return NISSY_ERROR_DATA;
         }
 
         arg = (dfsarg_solveh48_t) {
diff --git a/src/solvers/h48/solve_multithread.h b/src/solvers/h48/solve_multithread.h
index 766d84a..3ba914a 100644
--- a/src/solvers/h48/solve_multithread.h
+++ b/src/solvers/h48/solve_multithread.h
@@ -20,7 +20,8 @@ STATIC void copy_queue(task_queue_t *, task_queue_t *, int, _Atomic int64_t *);
 STATIC void *start_thread(void *);
 STATIC int64_t solve_h48_bfs(dfsarg_solveh48_t *, task_queue_t *, int8_t);
 STATIC int64_t solve_h48_single(dfsarg_solveh48_t *, task_queue_t *);
-STATIC int64_t solve_h48_multithread(cube_t, int8_t, int8_t, int8_t, const void *, char *);
+STATIC int64_t solve_h48_multithread(cube_t, int8_t, int8_t, int8_t,
+    uint64_t, const void *, uint64_t, char *);
 
 STATIC void
 solve_h48_appendsolution_thread(dfsarg_solveh48_t *arg, task_queue_t *tq)
@@ -236,8 +237,11 @@ solve_h48_multithread(
         int8_t minmoves,
         int8_t maxmoves,
         int8_t maxsolutions,
+        uint64_t data_size,
         const void *data,
-        char *solutions)
+        uint64_t solutions_size,
+        char *solutions
+)
 {
         _Atomic int64_t nsols = 0;
         int p_depth = 0;
@@ -245,9 +249,9 @@ solve_h48_multithread(
         tableinfo_t info;
         pthread_t threads[THREADS];
 
-        if (!readtableinfo_n(data, 2, &info)){
+        if (readtableinfo_n(data_size, data, 2, &info) != NISSY_OK) {
                 LOG("solve_h48: error reading table\n");
-                return 0;
+                return NISSY_ERROR_DATA;
         }
 
         arg = (dfsarg_solveh48_t){