/* Copyright 2005-2006, Voltage Security, all rights reserved.
 */
#include "vibecrypto.h"
#include "environment.h"
#include "base.h"
#include "libctx.h"
#include "algobj.h"
#include "cipher.h"
#include "p1pad.h"
#include "errorctx.h"

/* Generate mask bytes based on the seed and XOR the dataToMask.
 */
static int VOLT_CALLING_CONV VoltMgf1 VOLT_PROTO_LIST ((
   VoltLibCtx *libCtx,
   VtAlgorithmObject digester,
   unsigned int digestLen,
   unsigned char *seed,
   unsigned int seedLen,
   unsigned char *dataToMask,
   unsigned int dataToMaskLen
));

int VtPaddingPkcs1OAEP (
   VtAlgorithmObject object,
   VtPaddingInfo *info,
   unsigned int flag
   )
{
  int status;
  unsigned int digestLen;
  VoltAlgorithmObject *obj = (VoltAlgorithmObject *)object;
  VoltLibCtx *libCtx = (VoltLibCtx *)(obj->voltObject.libraryCtx);
  VoltCipherClassCtx *cipherCtx;
  VoltPaddingInfo *padInfo = (VoltPaddingInfo *)info;
  VtOAEPInfo *oaepInfo;
  VoltP1OaepCtx *ctx;
  unsigned char *buffer = (unsigned char *)0;
  VtAlgorithmObject digester = (VtAlgorithmObject)0;
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  do
  {
    /* Check the flag, it should be VOLT_PADDING_SET_TYPE_FLAG.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_TYPE;
    if (flag != VOLT_PADDING_SET_TYPE_FLAG)
      break;

    /* The associated info should be a pointer to a VtOAEPInfo struct.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_ASSOCIATED_INFO;
    if (padInfo->info == (Pointer)0)
      break;

    oaepInfo = (VtOAEPInfo *)(padInfo->info);

    VOLT_SET_FNCT_LINE (fnctLine)
    if (oaepInfo->digestImpl == (VtAlgorithmImpl *)0)
      break;

    /* Check the class of the object. It should be
     * VOLT_CLASS_ASYMMETRIC_CIPHER.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_SET;
    if (obj->algClass != VOLT_CLASS_ASYMMETRIC_CIPHER)
      break;

    /* We have an asymmetric cipher object, which means we have a
     * CipherCtx.
     */
    cipherCtx = (VoltCipherClassCtx *)(obj->classCtx);

    /* This only works with RSA.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    if (cipherCtx->setState != VOLT_CIPHER_SET_STATE_RSA)
      break;

    /* Build the digest object we will use to do the OAEP work. Do this
     * first so we can digest the label. We want to get the digestLen
     * before we allocate memory for the ctx.
     */
    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtCreateAlgorithmObject (
      (VtLibCtx)libCtx, oaepInfo->digestImpl, oaepInfo->digestInfo,
      &digester);
    if (status != 0)
      break;

    /* Make sure the object created is indeed a digest.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_ASSOCIATED_INFO;
    if (((VoltAlgorithmObject *)digester)->algClass != VOLT_CLASS_DIGEST)
      break;

    /* Digest the label. Actually, just get the output length first.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtDigestInit (digester);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtDigestFinal (
      digester, (unsigned char *)0, 0, (unsigned char *)0, 0, &digestLen);
    if (status == 0)
      status = VT_ERROR_GENERAL;
    if (status != VT_ERROR_BUFFER_TOO_SMALL)
      break;

    /* Build the OAEP ctx.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_MEMORY;
    buffer = (unsigned char *)Z2Malloc (
      sizeof (VoltP1OaepCtx) + digestLen, VOLT_MEMORY_SENSITIVE);
    if (buffer == (unsigned char *)0)
      break;
    Z2Memset (buffer, 0, sizeof (VoltP1OaepCtx) + digestLen);

    ctx = (VoltP1OaepCtx *)buffer;
    ctx->digestOfLabel = buffer + sizeof (VoltP1OaepCtx);
    ctx->digestLen = digestLen;

    /* Digest the label now.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtDigestFinal (
      digester, oaepInfo->label.data, oaepInfo->label.len,
      ctx->digestOfLabel, digestLen, &digestLen);
    if (status != 0)
      break;

    ctx->digester = digester;

    cipherCtx->Pad = P1OAEPad;
    cipherCtx->Unpad = P1OAEPUnpad;
    cipherCtx->padCtx = (Pointer)ctx;
    cipherCtx->PadCtxDestroy = P1OAEPCtxDestroy;

    obj->subAlg2 |= VOLT_SUB_ALG_P1_OAEP_PAD;

    status = 0;

  } while (0);

  if (status == 0)
    return (0);

  /* If there was an error, destrow anything we created and would have
   * returned, but didn't.
   */
  VtDestroyAlgorithmObject (&digester);

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

  VOLT_LOG_ERROR_INFO_COMPARE (
    status, 0, obj, status, 0, errorType,
    (char *)0, "VtPaddingPkcs1OAEP", fnctLine, (char *)0)

  return (status);
}

int P1OAEPad (
   VoltAlgorithmObject *obj,
   VtRandomObject random,
   Pointer context,
   unsigned char *block,
   unsigned int inputLen,
   unsigned int blockSize
   )
{
  int status;
  unsigned int zeroLen;
  VoltLibCtx *libCtx = (VoltLibCtx *)(obj->voltObject.libraryCtx);
  VoltP1OaepCtx *ctx = (VoltP1OaepCtx *)context;
  VtRandomObject randomToUse;
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  do
  {
    /* If there's no random object, get one from the libCtx.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_NO_RANDOM_OBJECT;
    randomToUse = random;
    if (random == (VtRandomObject)0)
    {
      randomToUse = (VtRandomObject)VoltGetLibCtxInfo (
        (VtLibCtx)libCtx, VOLT_LIB_CTX_INFO_TYPE_RANDOM);

      if (randomToUse == (VtRandomObject)0)
        break;
    }

    /* Make sure the random object is valid.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_RANDOM_OBJ;
    if (VOLT_OBJECT_TYPE_NOT_EQUAL (randomToUse, VOLT_OBJECT_TYPE_RANDOM))
      break;

    /* Start by creating a block that looks like this.
     *    00
     *    <digestLen random bytes: known as seed>
     *    <digest of label>
     *    00 ... 00
     *    01
     *    <data to encrypt>
     * It is possible to have no 00 bytes (after digest of label and
     * before the 01).
     * That means the length of the data to encrypt must be less than
     *    blockSize - 2*digestLen - 2
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_INPUT_LENGTH;
    zeroLen = inputLen + (2 * ctx->digestLen) + 2;
    if (zeroLen > blockSize)
      break;

    zeroLen = blockSize - zeroLen;

    /* Move the data to the end.
     */
    Z2Memmove (block + (blockSize - inputLen), block, inputLen);

    /* The first byte is 00.
     */
    block[0] = 0;

    /* The next digestLen bytes are random. This is the seed.
     */
    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtGenerateRandomBytes (randomToUse, block + 1, ctx->digestLen);
    if (status != 0)
      break;

    /* The next digestLen bytes are the digest of the label.
     */
    Z2Memcpy (block + ctx->digestLen + 1, ctx->digestOfLabel, ctx->digestLen);

    /* The next zeroLen bytes are 00.
     */
    Z2Memset (block + (2 * ctx->digestLen) + 1, 0, zeroLen);

    /* Finally, there's a 01 byte. Use zeroLen as a temp variable.
     */
    zeroLen += (2 * ctx->digestLen) + 1;
    block[zeroLen] = 1;

    /* Call on the mask generating function to generate a mask for
     * everything after the seed, using the seed as the MGF seed.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VoltMgf1 (
      libCtx, ctx->digester, ctx->digestLen, block + 1, ctx->digestLen,
      block + ctx->digestLen + 1, blockSize - (ctx->digestLen + 1));
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = VoltMgf1 (
      libCtx, ctx->digester, ctx->digestLen,
      block + ctx->digestLen + 1, blockSize - (ctx->digestLen + 1),
      block + 1, ctx->digestLen);

  } while (0);

  VOLT_LOG_ERROR_INFO_COMPARE (
    status, 0, obj, status, 0, errorType,
    (char *)0, "P1OAEPad", fnctLine, (char *)0)

  return (status);
}

int P1OAEPUnpad (
   VoltAlgorithmObject *obj,
   Pointer context,
   unsigned char *block,
   unsigned int blockSize,
   unsigned int *outputLen
   )
{
  int status;
  unsigned int index;
  VoltLibCtx *libCtx = (VoltLibCtx *)(obj->voltObject.libraryCtx);
  VoltP1OaepCtx *ctx = (VoltP1OaepCtx *)context;
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  do
  {
    /* Use the data after the seed as the seed for MGF1 to unmask the
     * seed (whoever came up with the terminology for OAEP and MGF,
     * didn't do a very good job).
     */
    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VoltMgf1 (
      libCtx, ctx->digester, ctx->digestLen,
      block + ctx->digestLen + 1, blockSize - (ctx->digestLen + 1),
      block + 1, ctx->digestLen);
    if (status != 0)
      break;

    /* Use the seed as the seed for MGF1 to unmask the data after the
     * seed.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VoltMgf1 (
      libCtx, ctx->digester, ctx->digestLen, block + 1, ctx->digestLen,
      block + ctx->digestLen + 1, blockSize - (ctx->digestLen + 1));
    if (status != 0)
      break;

    /* To help avoid timing attacks, perform all checks.
     */
    VOLT_SET_FNCT_LINE (fnctLine)

    /* Make sure the first byte is 00.
     */
    if (block[0] != 0)
      status = VT_ERROR_INVALID_PAD;

    /* Skip the next digestLen bytes. Then the next digestLen bytes
     * should be the digest of the label.
     */
    if (Z2Memcmp (
      block + ctx->digestLen + 1, ctx->digestOfLabel, ctx->digestLen) != 0)
      status = VT_ERROR_INVALID_PAD;

    /* The next bytes should be 00, until we hit 01.
     */
    index = (2 * ctx->digestLen) + 1;
    for (; index < blockSize; ++index)
      if (block[index] != 0)
        break;

    /* If the 00s never ended, error. Otherwise, continue checking.
     */
    if (index >= blockSize)
    {
      status = VT_ERROR_INVALID_PAD;
    }
    else
    {
      if (block[index] != 1)
        status = VT_ERROR_INVALID_PAD;

      /* Move the message block to the front.
       */
      index++;
      Z2Memmove (block, block + index, blockSize - index);
      *outputLen = blockSize - index;
    }

  } while (0);

  VOLT_LOG_ERROR_INFO_COMPARE (
    status, 0, obj, status, 0, errorType,
    (char *)0, "P1OAEPUnpad", fnctLine, (char *)0)

  return (status);
}

void P1OAEPCtxDestroy (
   Pointer object,
   Pointer context
   )
{
  VoltAlgorithmObject *obj;
  VoltLibCtx *libCtx;
  VoltCipherClassCtx *cipherCtx;
  VoltP1OaepCtx *ctx;

  /* Anything to destroy?
   */
  if ( (object == (Pointer)0) || (context == (Pointer)0) )
    return;

  obj = (VoltAlgorithmObject *)object;
  libCtx = (VoltLibCtx *)(obj->voltObject.libraryCtx);
  cipherCtx = (VoltCipherClassCtx *)(obj->classCtx);
  ctx = (VoltP1OaepCtx *)(cipherCtx->padCtx);

  VtDestroyAlgorithmObject (&(ctx->digester));

  Z2Free (ctx);
}

static int VoltMgf1 (
   VoltLibCtx *libCtx,
   VtAlgorithmObject digester,
   unsigned int digestLen,
   unsigned char *seed,
   unsigned int seedLen,
   unsigned char *dataToMask,
   unsigned int dataToMaskLen
   )
{
  int status;
  unsigned int index, currentLen, outputLen;
  unsigned char *buffer = (unsigned char *)0;
  unsigned char counter[4] = { 0, 0, 0, 0 };
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  /* Allocate a buffer to hold the results of digesting. If the Malloc
   * works, go ahead and do the mask operation.
   */
  VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
  VOLT_SET_FNCT_LINE (fnctLine)
  status = VT_ERROR_MEMORY;
  buffer = (unsigned char *)Z2Malloc (digestLen, VOLT_MEMORY_SENSITIVE);
  if (buffer != (unsigned char *)0)
  {
    VOLT_SET_ERROR_TYPE (errorType, 0)

    do
    {
      /* How many bytes do we still need? If we can't get them all, get
       * digestLen more mask bytes.
       */
      currentLen = dataToMaskLen;
      if (dataToMaskLen >= digestLen)
        currentLen = digestLen;

      dataToMaskLen -= currentLen;

      /* Find digest (seed || counter).
       */
      VOLT_SET_FNCT_LINE (fnctLine)
      status = VtDigestInit (digester);
      if (status != 0)
        break;

      VOLT_SET_FNCT_LINE (fnctLine)
      status = VtDigestUpdate (digester, seed, seedLen);
      if (status != 0)
        break;

      status = VtDigestFinal (
        digester, counter, 4, buffer, digestLen, &outputLen);
      if (status != 0)
        break;

      /* Use this result to mask the next currentLen bytes of the
       * dataToMask.
       */
      for (index = 0; index < currentLen; ++index)
      {
        *dataToMask ^= buffer[index];
        dataToMask++;
      }

      if (dataToMaskLen == 0)
        break;

      /* Increment the counter. A true implementation would worry about
       * carry from one byte to another (add 1 to 0xff and get 0, add
       * the carry bit to the next byte). However, that will only
       * happen if a block to be masked is > 5100 bytes (assuming a
       * minimum 20-byte digest). That would mean we were dealing with
       * a 40,000 bit key (instead of the 1024, 2048 and possibly 4096
       * bit keys in use today).
       * So this implementation does not look at the carry. If we ever
       * have to mask more than 5100 bytes with a 20-byte digest (or an
       * 8160 bytes with a 32-byte digest, etc.), we'll have to change
       * this code.
       */
      counter[3]++;

    } while (1);
  }

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

  VOLT_LOG_ERROR_INFO_COMPARE (
    status, libCtx, 0, status, 0, errorType,
    (char *)0, "VoltMgf1", fnctLine, (char *)0)

  return (status);
}