Merge pull request #39 from levitte/bugfix_gen_prime

Bugfixes to fp_prime_random_ex [reboot]

Author: sjaeckel

src/headers/tfm.h

@@ -318,8 +318,8 @@
 
    typedef ulong64            fp_digit;
 #define SIZEOF_FP_DIGIT 8
+#define DIGIT_SHIFT     6
    typedef unsigned long      fp_word __attribute__ ((mode(TI)));
-
 #else
 
    /* this is to make porting into LibTomCrypt easier :-) */
@@ -335,6 +335,7 @@
 
    typedef unsigned int       fp_digit;
 #define SIZEOF_FP_DIGIT 4
+#define DIGIT_SHIFT     5
    typedef ulong64            fp_word;
 #endif /* FP_64BIT */
 

src/numtheory/fp_prime_random_ex.c

@@ -2,11 +2,18 @@
 /* SPDX-License-Identifier: Unlicense */
 #include <tfm_private.h>
 
+#define fp_on_bitnum(a, bitnum) \
+    a->dp[(bitnum) >> DIGIT_SHIFT] |= (fp_digit)1 << ((bitnum) & (DIGIT_BIT-1))
+
+#define fp_off_bitnum(a, bitnum) \
+    a->dp[(bitnum) >> DIGIT_SHIFT] &= ~((fp_digit)1 << ((bitnum) & (DIGIT_BIT-1)))
+
 /* This is possibly the mother of all prime generation functions, muahahahahaha! */
 int fp_prime_random_ex(fp_int *a, int t, int size, int flags, tfm_prime_callback cb, void *dat)
 {
-   unsigned char *tmp, maskAND, maskOR_msb, maskOR_lsb;
-   int res, err, bsize, maskOR_msb_offset;
+   fp_digit maskAND_msb, maskOR_lsb;
+   int res, bsize, dsize;
+   unsigned char buf[FP_SIZE * sizeof(fp_digit)];
 
    /* sanity check the input */
    if (size <= 1 || cb == NULL || t <= 0 || t > FP_PRIME_SIZE) {
@@ -18,26 +25,13 @@ int fp_prime_random_ex(fp_int *a, int t, int size, int flags, tfm_prime_callback
       flags |= TFM_PRIME_BBS;
    }
 
-   /* calc the byte size */
-   bsize = (size>>3)+(size&7?1:0);
+   /* calc the size in fp_digit */
+   dsize = (size + DIGIT_BIT - 1) >> DIGIT_SHIFT;
+   /* calc the size in bytes */
+   bsize = (size + 7) >> 3;
 
-   /* we need a buffer of bsize bytes */
-   tmp = malloc(bsize);
-   if (tmp == NULL) {
-      return FP_MEM;
-   }
-
-   /* calc the maskAND value for the MSbyte*/
-   maskAND = 0xFF >> ((8 - (size & 7)) & 7);
-
-   /* calc the maskOR_msb */
-   maskOR_msb        = 0;
-   maskOR_msb_offset = (size - 2) >> 3;
-   if (flags & TFM_PRIME_2MSB_ON) {
-      maskOR_msb     |= 1 << ((size - 2) & 7);
-   } else if (flags & TFM_PRIME_2MSB_OFF) {
-      maskAND        &= ~(1 << ((size - 2) & 7));
-   }
+   /* calc the maskAND value for the MSbyte */
+   maskAND_msb = FP_MASK >> ((DIGIT_BIT - size) & (DIGIT_BIT-1));
 
    /* get the maskOR_lsb */
    maskOR_lsb         = 1;
@@ -47,21 +41,30 @@ int fp_prime_random_ex(fp_int *a, int t, int size, int flags, tfm_prime_callback
 
    do {
       /* read the bytes */
-      if (cb(tmp, bsize, dat) != bsize) {
-         err = FP_VAL;
-         goto error;
+      if (cb(buf, bsize, dat) != bsize) {
+         return FP_VAL;
       }
+      fp_read_unsigned_bin(a, buf, bsize);
+
+      /* make sure the MSbyte has the required number of bits */
+      a->dp[dsize-1]    &= maskAND_msb;
 
-      /* work over the MSbyte */
-      tmp[0]    &= maskAND;
-      tmp[0]    |= 1 << ((size - 1) & 7);
+      /* Force a->used as well, it could be smaller if the highest bits were
+         generated as 0 by the callback. */
+      a->used           = dsize;
 
-      /* mix in the maskORs */
-      tmp[maskOR_msb_offset]   |= maskOR_msb;
-      tmp[bsize-1]             |= maskOR_lsb;
+      /* modify the LSbyte as requested */
+      a->dp[0]          |= maskOR_lsb;
 
-      /* read it in */
-      fp_read_unsigned_bin(a, tmp, bsize);
+      /* turn on the MSbit to force the requested magnitude */
+      fp_on_bitnum(a, size-1);
+
+      /* modify the 2nd MSBit */
+      if (flags & TFM_PRIME_2MSB_ON) {
+          fp_on_bitnum(a, size-2);
+      } else if (flags & TFM_PRIME_2MSB_OFF) {
+          fp_off_bitnum(a, size-2);
+      }
 
       /* is it prime? */
       res = fp_isprime_ex(a, t);
@@ -83,8 +86,5 @@ int fp_prime_random_ex(fp_int *a, int t, int size, int flags, tfm_prime_callback
       fp_add_d(a, 1, a);
    }
 
-   err = FP_OKAY;
-error:
-   free(tmp);
-   return err;
+   return FP_OKAY;
 }