/*
 * aeskeywrap.c - implement AES Key Wrap algorithm from RFC 3394
 *
 * The contents of this file are subject to the Mozilla Public
 * License Version 1.1 (the "License"); you may not use this file
 * except in compliance with the License. You may obtain a copy of
 * the License at http://www.mozilla.org/MPL/
 * 
 * Software distributed under the License is distributed on an "AS
 * IS" basis, WITHOUT WARRANTY OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * rights and limitations under the License.
 * 
 * The Original Code is the Netscape security libraries.
 * 
 * The Initial Developer of the Original Code is Netscape
 * Communications Corporation.  Portions created by Netscape are 
 * Copyright (C) 2002, 2003 Netscape Communications Corporation.  All
 * Rights Reserved.
 * 
 * Contributor(s):
 * 
 * Alternatively, the contents of this file may be used under the
 * terms of the GNU General Public License Version 2 or later (the
 * "GPL"), in which case the provisions of the GPL are applicable 
 * instead of those above.  If you wish to allow use of your 
 * version of this file only under the terms of the GPL and not to
 * allow others to use your version of this file under the MPL,
 * indicate your decision by deleting the provisions above and
 * replace them with the notice and other provisions required by
 * the GPL.  If you do not delete the provisions above, a recipient
 * may use your version of this file under either the MPL or the
 * GPL.
 *
 * $Id$
 */

#include "prcpucfg.h"
#if defined(IS_LITTLE_ENDIAN) || defined(SHA_NO_LONG_LONG)
#define BIG_ENDIAN_WITH_64_BIT_REGISTERS 0
#else
#define BIG_ENDIAN_WITH_64_BIT_REGISTERS 1
#endif
#include "prtypes.h"	/* for PRUintXX */
#include "secport.h"	/* for PORT_XXX */
#include "secerr.h"
#include "blapi.h"	/* for AES_ functions */


struct AESKeyWrapContextStr {
     AESContext *  aescx;
     unsigned char iv[AES_KEY_WRAP_IV_BYTES];
};

/******************************************/
/*
** AES key wrap algorithm, RFC 3394
*/

/*
** Create a new AES context suitable for AES encryption/decryption.
** 	"key" raw key data
** 	"keylen" the number of bytes of key data (16, 24, or 32)
*/
extern AESKeyWrapContext *
AESKeyWrap_CreateContext(const unsigned char *key, const unsigned char *iv, 
                         int encrypt, unsigned int keylen)
{
    AESKeyWrapContext * cx = PORT_ZNew(AESKeyWrapContext);
    if (!cx) 
    	return NULL;	/* error is already set */
    cx->aescx = AES_CreateContext(key, NULL, NSS_AES, encrypt, keylen, 
                                  AES_BLOCK_SIZE);
    if (!cx->aescx) {
        PORT_Free(cx);
	return NULL; 	/* error should already be set */
    }
    if (iv) {
    	memcpy(cx->iv, iv, AES_KEY_WRAP_IV_BYTES);
    } else {
	memset(cx->iv, 0xA6, AES_KEY_WRAP_IV_BYTES);
    }
    return cx;
}

/*
** Destroy a AES KeyWrap context.
**	"cx" the context
**	"freeit" if PR_TRUE then free the object as well as its sub-objects
*/
extern void 
AESKeyWrap_DestroyContext(AESKeyWrapContext *cx, PRBool freeit)
{
    if (cx) {
	if (cx->aescx)
	    AES_DestroyContext(cx->aescx, PR_TRUE);
	memset(cx, 0, sizeof *cx);
	if (freeit)
	    PORT_Free(cx);
    }
}

#if !BIG_ENDIAN_WITH_64_BIT_REGISTERS

/* The AES Key Wrap algorithm has 64-bit values that are ALWAYS big-endian
** (Most significant byte first) in memory.  The only ALU operations done
** on them are increment, decrement, and XOR.  So, on little-endian CPUs,
** and on CPUs that lack 64-bit registers, these big-endian 64-bit operations
** are simulated in the following code.  This is thought to be faster and
** simpler than trying to convert the data to little-endian and back.
*/

/* A and T point to two 64-bit values stored most signficant byte first
** (big endian).  This function increments the 64-bit value T, and then
** XORs it with A, changing A.
*/ 
static void
increment_and_xor(unsigned char *A, unsigned char *T)
{
    if (!++T[7])
        if (!++T[6])
	    if (!++T[5])
		if (!++T[4])
		    if (!++T[3])
			if (!++T[2])
			    if (!++T[1])
				 ++T[0];

    A[0] ^= T[0];
    A[1] ^= T[1];
    A[2] ^= T[2];
    A[3] ^= T[3];
    A[4] ^= T[4];
    A[5] ^= T[5];
    A[6] ^= T[6];
    A[7] ^= T[7];
}

/* A and T point to two 64-bit values stored most signficant byte first
** (big endian).  This function XORs T with A, giving a new A, then 
** decrements the 64-bit value T.
*/ 
static void
xor_and_decrement(unsigned char *A, unsigned char *T)
{
    A[0] ^= T[0];
    A[1] ^= T[1];
    A[2] ^= T[2];
    A[3] ^= T[3];
    A[4] ^= T[4];
    A[5] ^= T[5];
    A[6] ^= T[6];
    A[7] ^= T[7];

    if (!T[7]--)
        if (!T[6]--)
	    if (!T[5]--)
		if (!T[4]--)
		    if (!T[3]--)
			if (!T[2]--)
			    if (!T[1]--)
				 T[0]--;

}

/* Given an unsigned long t (in host byte order), store this value as a
** 64-bit big-endian value (MSB first) in *pt.
*/
static void
set_t(unsigned char *pt, unsigned long t)
{
    pt[7] = (unsigned char)t; t >>= 8;
    pt[6] = (unsigned char)t; t >>= 8;
    pt[5] = (unsigned char)t; t >>= 8;
    pt[4] = (unsigned char)t; t >>= 8;
    pt[3] = (unsigned char)t; t >>= 8;
    pt[2] = (unsigned char)t; t >>= 8;
    pt[1] = (unsigned char)t; t >>= 8;
    pt[0] = (unsigned char)t;
}

#endif

/*
** Perform AES key wrap.
**	"cx" the context
**	"output" the output buffer to store the encrypted data.
**	"outputLen" how much data is stored in "output". Set by the routine
**	   after some data is stored in output.
**	"maxOutputLen" the maximum amount of data that can ever be
**	   stored in "output"
**	"input" the input data
**	"inputLen" the amount of input data
*/
extern SECStatus 
AESKeyWrap_Encrypt(AESKeyWrapContext *cx, unsigned char *output,
            unsigned int *pOutputLen, unsigned int maxOutputLen,
            const unsigned char *input, unsigned int inputLen)
{
    PRUint64 *     R          = NULL;
    unsigned int   nBlocks;
    unsigned int   i, j;
    unsigned int   aesLen     = AES_BLOCK_SIZE;
    unsigned int   outLen     = inputLen + AES_KEY_WRAP_BLOCK_SIZE;
    SECStatus      s          = SECFailure;
    /* These PRUint64s are ALWAYS big endian, regardless of CPU orientation. */
    PRUint64       t;
    PRUint64       B[2];

#define A B[0]

    /* Check args */
    if (!inputLen || 0 != inputLen % AES_KEY_WRAP_BLOCK_SIZE) {
	PORT_SetError(SEC_ERROR_INPUT_LEN);
	return s;
    }
#ifdef maybe
    if (!output && pOutputLen) {	/* caller is asking for output size */
    	*pOutputLen = outLen;
	return SECSuccess;
    }
#endif
    if (maxOutputLen < outLen) {
	PORT_SetError(SEC_ERROR_OUTPUT_LEN);
	return s;
    }
    if (cx == NULL || output == NULL || input == NULL) {
	PORT_SetError(SEC_ERROR_INVALID_ARGS);
	return s;
    }
    nBlocks = inputLen / AES_KEY_WRAP_BLOCK_SIZE;
    R = PORT_NewArray(PRUint64, nBlocks + 1);
    if (!R)
    	return s;	/* error is already set. */
    /* 
    ** 1) Initialize variables.
    */
    memcpy(&A, cx->iv, AES_KEY_WRAP_IV_BYTES);
    memcpy(&R[1], input, inputLen);
#if BIG_ENDIAN_WITH_64_BIT_REGISTERS
    t = 0;
#else
    memset(&t, 0, sizeof t);
#endif
    /* 
    ** 2) Calculate intermediate values.
    */
    for (j = 0; j < 6; ++j) {
    	for (i = 1; i <= nBlocks; ++i) {
	    B[1] = R[i];
	    s = AES_Encrypt(cx->aescx, (unsigned char *)B, &aesLen, 
	                    sizeof B,  (unsigned char *)B, sizeof B);
	    if (s != SECSuccess) 
	        break;
	    R[i] = B[1];
	    /* here, increment t and XOR A with t (in big endian order); */
#if BIG_ENDIAN_WITH_64_BIT_REGISTERS
   	    A ^= ++t; 
#else
	    increment_and_xor((unsigned char *)&A, (unsigned char *)&t);
#endif
	}
    }
    /* 
    ** 3) Output the results.
    */
    if (s == SECSuccess) {
    	R[0] =  A;
	memcpy(output, &R[0], outLen);
	if (pOutputLen)
	    *pOutputLen = outLen;
    } else if (pOutputLen) {
    	*pOutputLen = 0;
    }
    PORT_ZFree(R, outLen);
    return s;
}
#undef A

/*
** Perform AES key unwrap.
**	"cx" the context
**	"output" the output buffer to store the decrypted data.
**	"outputLen" how much data is stored in "output". Set by the routine
**	   after some data is stored in output.
**	"maxOutputLen" the maximum amount of data that can ever be
**	   stored in "output"
**	"input" the input data
**	"inputLen" the amount of input data
*/
extern SECStatus 
AESKeyWrap_Decrypt(AESKeyWrapContext *cx, unsigned char *output,
            unsigned int *pOutputLen, unsigned int maxOutputLen,
            const unsigned char *input, unsigned int inputLen)
{
    PRUint64 *     R          = NULL;
    unsigned int   nBlocks;
    unsigned int   i, j;
    unsigned int   aesLen     = AES_BLOCK_SIZE;
    unsigned int   outLen;
    SECStatus      s          = SECFailure;
    /* These PRUint64s are ALWAYS big endian, regardless of CPU orientation. */
    PRUint64       t;
    PRUint64       B[2];

#define A B[0]

    /* Check args */
    if (inputLen < 3 * AES_KEY_WRAP_BLOCK_SIZE || 
        0 != inputLen % AES_KEY_WRAP_BLOCK_SIZE) {
	PORT_SetError(SEC_ERROR_INPUT_LEN);
	return s;
    }
    outLen = inputLen - AES_KEY_WRAP_BLOCK_SIZE;
#ifdef maybe
    if (!output && pOutputLen) {	/* caller is asking for output size */
    	*pOutputLen = outLen;
	return SECSuccess;
    }
#endif
    if (maxOutputLen < outLen) {
	PORT_SetError(SEC_ERROR_OUTPUT_LEN);
	return s;
    }
    if (cx == NULL || output == NULL || input == NULL) {
	PORT_SetError(SEC_ERROR_INVALID_ARGS);
	return s;
    }
    nBlocks = inputLen / AES_KEY_WRAP_BLOCK_SIZE;
    R = PORT_NewArray(PRUint64, nBlocks);
    if (!R)
    	return s;	/* error is already set. */
    nBlocks--;
    /* 
    ** 1) Initialize variables.
    */
    memcpy(&R[0], input, inputLen);
    A = R[0];
#if BIG_ENDIAN_WITH_64_BIT_REGISTERS
    t = 6UL * nBlocks;
#else
    set_t((unsigned char *)&t, 6UL * nBlocks);
#endif
    /* 
    ** 2) Calculate intermediate values.
    */
    for (j = 0; j < 6; ++j) {
    	for (i = nBlocks; i; --i) {
	    /* here, XOR A with t (in big endian order) and decrement t; */
#if BIG_ENDIAN_WITH_64_BIT_REGISTERS
   	    A ^= t--; 
#else
	    xor_and_decrement((unsigned char *)&A, (unsigned char *)&t);
#endif
	    B[1] = R[i];
	    s = AES_Decrypt(cx->aescx, (unsigned char *)B, &aesLen, 
	                    sizeof B,  (unsigned char *)B, sizeof B);
	    if (s != SECSuccess) 
	        break;
	    R[i] = B[1];
	}
    }
    /* 
    ** 3) Output the results.
    */
    if (s == SECSuccess) {
	int bad = memcmp(&A, cx->iv, AES_KEY_WRAP_IV_BYTES);
	if (!bad) {
	    memcpy(output, &R[1], outLen);
	    if (pOutputLen)
		*pOutputLen = outLen;
	} else {
	    PORT_SetError(SEC_ERROR_BAD_DATA);
	    if (pOutputLen) 
		*pOutputLen = 0;
    	}
    } else if (pOutputLen) {
    	*pOutputLen = 0;
    }
    PORT_ZFree(R, inputLen);
    return s;
}
#undef A