Fix overflow issue + odd behaviour

It is best to check for an overflow risk and lower the step
BEFORE we actually perform the inversion

Signed-off-by: sigsegv7 <ian@vegaa.systems>

1 files changed, 5 insertions, 5 deletions

diff --git a/src/main.c b/src/main.c
index 259b294..a42d05b 100644
--- a/src/main.c
+++ b/src/main.c
@@ -152,6 +152,11 @@ encrypt(const struct cpu_info *info, char *buf, size_t buf_size)
             assert((step & 1) == 0 && step <= 16);
         }
 
+        /* Ensure we don't cause any overflows */
+        while (((current_pos + step) >= buf_size) && step > 1)
+            /* Essentially divide the step by 2, just faster */
+            step >>= 1;
+
         switch (step) {
         case 16:
             accel_invert128((uintptr_t)buf + current_pos);
@@ -171,11 +176,6 @@ encrypt(const struct cpu_info *info, char *buf, size_t buf_size)
         }
 
         current_pos += step;
-
-        /* Ensure we don't cause any overflows */
-        while (((current_pos + step) >= buf_size) && step > 1)
-            /* Essentially divide the step by 2, just faster */
-            step >>= 1;
     }
 }