/* Copyright 2003-2006, Voltage Security, all rights reserved.
 */

#include "vibecrypto.h"
#include "environment.h"
#include "base.h"
#include "libctx.h"
#include "algobj.h"
#include "mpint.h"
#include "prime.h"
#include "errorctx.h"
#include "surrender.h"

/* Get a random value to use in the prime test. This function will
 * digest the primeCandidate along with the current contents of the
 * buffer to get some random bytes. This is used so that we don't need
 * to obtain random bytes from the PRNG. THis is necessary so we can
 * pass FIPS tests.
 * <p>The random values generated using this technique are to be used
 * only for getting random values used in testing primes. It must
 * generate new values each time.
 */
static int VOLT_CALLING_CONV GetRandomValue VOLT_PROTO_LIST ((
   VoltLibCtx *libCtx,
   VoltMpIntCtx *mpCtx,
   VoltMpInt *primeCandidate,
   unsigned char *buffer,
   unsigned int bufferLen
));

int VoltGeneratePrimeRabinMiller (
   unsigned int primeSizeBits,
   unsigned int leadingBits,
   unsigned int iterationCount,
   unsigned int oneRandom,
   VtRandomObject random,
   VoltSurrenderCtx *surrCtx,
   unsigned int surrFlag,
   unsigned int *callNumber,
   VoltMpInt *relativePrime,
   VoltMpInt *prime
   )
{
  int status, count;
  unsigned int bufferSize, index, sieveSize, isPrime, relPrime;
  unsigned int bitMask, bitSet0, bitSet1, increment, callNum;
  unsigned char *buffer = (unsigned char *)0;
  unsigned char *sieve = (unsigned char *)0;
  VoltMpIntCtx *mpCtx = (VoltMpIntCtx *)(prime->mpCtx);
  VoltLibCtx *libCtx = (VoltLibCtx *)(mpCtx->voltObject.libraryCtx);
  VoltMpInt *mpIncrement = (VoltMpInt *)0;
  VoltMpInt *temp1 = (VoltMpInt *)0;
  VoltMpInt *temp2 = (VoltMpInt *)0;
  unsigned int firstPrimes[100] = FIRST_100_ODD_PRIMES;
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  /* This function accepts only 1 or 2 as valid leadingBits values.
   */
  if (leadingBits != 1)
    leadingBits = 2;

  callNum = 1;
  if (callNumber != (unsigned int *)0)
    callNum = *callNumber;

  do
  {
    /* Semi arbitrary limit on the size of the prime.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_INPUT_LENGTH;
    if (primeSizeBits <= 32)
      break;

    /* Create a buffer of the appropriate size. We'll fill it with
     * random bytes and use that as a starting point.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_MEMORY;
    bufferSize = (primeSizeBits + 7) / 8;
    buffer = (unsigned char *)Z2Malloc (bufferSize, VOLT_MEMORY_SENSITIVE);
    if (buffer == (unsigned char *)0)
      break;

    /* Create a mask for the first byte. Depending on the bit length,
     * we may want to trim some bits off the top.
     * Then make sure the msbit is set. Create a value to OR with the
     * first byte that will set the msbit.
     * Finally, create a value to OR with the second byte. If the
     * caller wants the two most significant bits set, and the bit
     * length is such that the second most significant bit appears in
     * the second byte, set that bit.
     */
    bitSet0 = primeSizeBits % 8;
    bitSet1 = 0;
    if (bitSet0 == 0)
    {
      bitSet0 = 0x80;
      if (leadingBits == 2)
        bitSet0 = 0xc0;
      bitMask = 0xff;
    }
    else if (bitSet0 == 1)
    {
      bitSet0 = 1;
      if (leadingBits == 2)
        bitSet1 = 0x80;
      bitMask = 1;
    }
    else
    {
      bitMask = 0xff >> (8 - bitSet0);
      increment = bitSet0 - leadingBits;
      bitSet0 = 1;
      if (leadingBits == 2)
        bitSet0 = 3;
      bitSet0 = (bitSet0 << increment);
    }

    /* Create a sieve buffer.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    sieveSize = VOLT_SIEVE_SIZE;
    sieve = (unsigned char *)Z2Malloc (sieveSize, VOLT_MEMORY_SENSITIVE);
    if (sieve == (unsigned char *)0)
      break;

    /* We'll need the increment as an MpInt later on.
     */
    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->CreateMpInt ((Pointer)mpCtx, &mpIncrement);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->CreateMpInt ((Pointer)mpCtx, &temp1);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->CreateMpInt ((Pointer)mpCtx, &temp2);
    if (status != 0)
      break;

    count = 0;
    isPrime = 1;
    relPrime = 1;
    do
    {
      callNum++;
      VOLT_CALL_SURRENDER (surrCtx, surrFlag, 0, callNum)

      if ( (count == 0) || (oneRandom == 0) )
      {
        /* Generate random bytes of the appropriate bit length.
         */
        VOLT_SET_ERROR_TYPE (errorType, 0)
        VOLT_SET_FNCT_LINE (fnctLine)
        status = VtGenerateRandomBytes (random, buffer, bufferSize);
        if (status != 0)
          break;

        /* Make sure the msbytes are set properly.
         */
        buffer[0] &= (unsigned char)bitMask;
        buffer[0] |= (unsigned char)bitSet0;
        buffer[1] |= (unsigned char)bitSet1;

        /* Set the lsbit to guarantee it is odd.
         */
        buffer[bufferSize - 1] |= 1;

        /* Set the MpInt to this value.
         */
        VOLT_SET_FNCT_LINE (fnctLine)
          status = mpCtx->OctetStringToMpInt (0, buffer, bufferSize, prime);
        if (status != 0)
          break;
      }
      else
      {
        /* If this is not the first loop, and if the caller wants one
         * random, just add increment to get the next number after the
         * previously tested value.
         */
        VOLT_SET_FNCT_LINE (fnctLine)
        status = mpCtx->IntToMpInt (0, increment, mpIncrement);
        if (status != 0)
          break;

        VOLT_SET_FNCT_LINE (fnctLine)
        status = mpCtx->Add (prime, mpIncrement, prime);
        if (status != 0)
          break;
      }

      /* Create the sieve for this candidate.
       */
      VOLT_SET_FNCT_LINE (fnctLine)
      status = VoltSieveCandidate (
        prime, firstPrimes, 100, sieve, sieveSize);
      if (status != 0)
        break;

      /* Now run Rabin-Miller on numbers that fell through the sieve.
       */
      increment = 0;
      for (index = 0; index < sieveSize; ++index)
      {
        if (sieve[index] != 0)
        {
          increment += 2;
          continue;
        }

        /* Add increment to prime, this is the number we want to check.
         */
        VOLT_SET_FNCT_LINE (fnctLine)
        status = mpCtx->IntToMpInt (0, increment, mpIncrement);
        if (status != 0)
          break;

        VOLT_SET_FNCT_LINE (fnctLine)
        status = mpCtx->Add (prime, mpIncrement, prime);
        if (status != 0)
          break;

        if (relativePrime != (VoltMpInt *)0)
        {
          VOLT_SET_FNCT_LINE (fnctLine)
          status = VoltCheckRelativePrime (
            mpCtx, relativePrime, prime, temp1, temp2, &relPrime);
          if (status != 0)
            break;
        }

        /* If TRUE (non zero), the two numbers pass the relative
         * prime test, continue. If FALSE (zero) the two numbers fail
         * the test, move on.
         */
        if (relPrime != 0)
        {
          callNum++;
          VOLT_CALL_SURRENDER (surrCtx, surrFlag, 0, callNum)

          VOLT_SET_FNCT_LINE (fnctLine)
          status = VoltRabinMillerTest (
            prime, primeSizeBits, iterationCount, random, &isPrime);
          if (status != 0)
            break;

          /* If this came back prime, we're done.
           */
          if (isPrime != 0)
            break;
        }

        /* Reset increment.
         */
        increment = 2;
      }
      if (status != 0)
        break;

      if (isPrime != 0)
        break;

      /* The Rabin-Miller test indicates that none of the numbers that
       * fell through the sieve were prime, get a new random starting
       * point. But first, don't run this test forever.  Try up to 100
       * times. Actually, it will almost certainly pass eventually,
       * but put the count in there just to be safe.
       */
      VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
      VOLT_SET_FNCT_LINE (fnctLine)
      status = VT_ERROR_NO_PRIME_FOUND;
      count++;
      if (count > 100)
        break;

    } while (1);

  } while (0);

  if (callNumber != (unsigned int *)0)
    *callNumber = callNum;

  if (buffer != (unsigned char *)0)
    Z2Free (buffer);

  if (sieve != (unsigned char *)0)
    Z2Free (sieve);

  mpCtx->DestroyMpInt (&mpIncrement);
  mpCtx->DestroyMpInt (&temp1);
  mpCtx->DestroyMpInt (&temp2);

  VOLT_LOG_ERROR_COMPARE (
    status, (VtLibCtx)libCtx, status, errorType, fnctLine,
    "VoltGeneratePrimeRabinMiller", (char *)0)

  return (status);
}

/* The Rabin-Miller algorithm as described in FIPS 186-2.
 * Step 1. Set i = 1 and n >= 50.
 * Step 2. Set w = the integer to be tested,
 *             w = 1 + (2^a)m, where m is odd and 2^a is the largest
 *             power of 2 dividing w - 1.
 * Step 3. Generate a random integer b in the range 1 < b < w.
 * Step 4. Set j = 0 and z = b^m mod w.
 * Step 5. If j = 0 and z = 1, or if z = w - 1, go to step 9.
 * Step 6. If j > 0 and z = 1, go to step 8.
 * Step 7. j = j + 1. If j < a, set z = z^2 mod w and go to step 5.
 * Step 8. w is not prime. Stop.
 * Step 9. If i < n, set i = i + 1 and go to step 3. Otherwise,
 *         w is probably prime.
 */
int VoltRabinMillerTest (
   VoltMpInt *primeCandidate,
   unsigned int primeSizeBits,
   unsigned int iterationCount,
   VtRandomObject random,
   unsigned int *isPrime
   )
{
  int status, compareResult;
  unsigned int index, count, expo, bufferSize, value;
  unsigned char *buffer = (unsigned char *)0;
  VoltMpIntCtx *mpCtx = (VoltMpIntCtx *)(primeCandidate->mpCtx);
  VoltLibCtx *libCtx = (VoltLibCtx *)(mpCtx->voltObject.libraryCtx);
  VoltMpInt *candidateMinusOne = (VoltMpInt *)0;
  VoltMpInt *mVal = (VoltMpInt *)0;
  VoltMpInt *zVal = (VoltMpInt *)0;
  VoltMpInt *randVal = (VoltMpInt *)0;
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  /* Init the isPrime return to not prime, change it if it passes the
   * tests.
   */
  *isPrime = 0;

  do
  {
    /* Create a buffer big enough to hold the random values we'll
     * generate as part of the test.
     * The random values must be < the prime candidate. So set the
     * length in bytes to be primeSizeBits / 8. If the bit size is a
     * multiple of 8, the buffer size will be the same byte size as the
     * prime. If the bit size is not a multiple of 8, then the buffer
     * size will be one byte shorter than the byte size of the prime.
     * If the byte size is smaller, then the random value we generate
     * will be smaller. If the byte size is the same, then the random
     * value we generate might be bigger than the prime. If the byte
     * size is the same, that means the bit length is a multiple of 8
     * which means the most significant bit of the most significant
     * byte of the prime is set. So if we clear the most significant
     * bit of the most significant byte of the random value, it will be
     * < the prime.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_MEMORY;
    bufferSize = primeSizeBits / 8;
    buffer = (unsigned char *)Z2Malloc (bufferSize, VOLT_MEMORY_SENSITIVE);
    if (buffer == (unsigned char *)0)
      break;
    Z2Memset (buffer, 0, bufferSize);