1 files changed, 42 insertions, 26 deletions

diff --git a/security/nss/lib/freebl/rijndael.c b/security/nss/lib/freebl/rijndael.c
index 569500408..58f5e5cdf 100644
--- a/security/nss/lib/freebl/rijndael.c
+++ b/security/nss/lib/freebl/rijndael.c
@@ -969,16 +969,16 @@ rijndael_decryptCBC(AESContext *cx, unsigned char *output,
  *
  ***********************************************************************/
 
-/* AES_CreateContext
- *
- * create a new context for Rijndael operations
- */
-AESContext *
-AES_CreateContext(const unsigned char *key, const unsigned char *iv, 
-                  int mode, int encrypt,
-                  unsigned int keysize, unsigned int blocksize)
+AESContext * AES_AllocateContext(void)
+{
+    return PORT_ZNew(AESContext);
+}
+
+SECStatus   
+AES_InitContext(AESContext *cx, const unsigned char *key, unsigned int keysize, 
+	        const unsigned char *iv, int mode, unsigned int encrypt,
+	        unsigned int blocksize)
 {
-    AESContext *cx;
     unsigned int Nk;
     /* According to Rijndael AES Proposal, section 12.1, block and key
      * lengths between 128 and 256 bits are supported, as long as the
@@ -992,20 +992,19 @@ AES_CreateContext(const unsigned char *key, const unsigned char *iv,
 	blocksize > RIJNDAEL_MAX_BLOCKSIZE || 
 	blocksize % 4 != 0) {
 	PORT_SetError(SEC_ERROR_INVALID_ARGS);
-	return NULL;
+	return SECFailure;
     }
     if (mode != NSS_AES && mode != NSS_AES_CBC) {
 	PORT_SetError(SEC_ERROR_INVALID_ARGS);
-	return NULL;
+	return SECFailure;
     }
     if (mode == NSS_AES_CBC && iv == NULL) {
 	PORT_SetError(SEC_ERROR_INVALID_ARGS);
-	return NULL;
+	return SECFailure;
     }
-    cx = PORT_ZNew(AESContext);
     if (!cx) {
-	PORT_SetError(SEC_ERROR_NO_MEMORY);
-    	return NULL;
+	PORT_SetError(SEC_ERROR_INVALID_ARGS);
+    	return SECFailure;
     }
     /* Nb = (block size in bits) / 32 */
     cx->Nb = blocksize / 4;
@@ -1020,10 +1019,9 @@ AES_CreateContext(const unsigned char *key, const unsigned char *iv,
     } else {
 	cx->worker = (encrypt) ? &rijndael_encryptECB : &rijndael_decryptECB;
     }
-    /* Allocate memory for the expanded key */
-    cx->expandedKey = PORT_ZNewArray(PRUint32, cx->Nb * (cx->Nr + 1));
-    if (!cx->expandedKey) {
-	PORT_SetError(SEC_ERROR_NO_MEMORY);
+    PORT_Assert((cx->Nb * (cx->Nr + 1)) <= RIJNDAEL_MAX_EXP_KEY_SIZE);
+    if ((cx->Nb * (cx->Nr + 1)) > RIJNDAEL_MAX_EXP_KEY_SIZE) {
+	PORT_SetError(SEC_ERROR_LIBRARY_FAILURE);
 	goto cleanup;
     }
     /* Generate expanded key */
@@ -1034,12 +1032,31 @@ AES_CreateContext(const unsigned char *key, const unsigned char *iv,
 	if (rijndael_invkey_expansion(cx, key, Nk) != SECSuccess)
 	    goto cleanup;
     }
-    return cx;
+    return SECSuccess;
 cleanup:
-    if (cx->expandedKey)
-	PORT_ZFree(cx->expandedKey, cx->Nb * (cx->Nr + 1));
-    PORT_ZFree(cx, sizeof *cx);
-    return NULL;
+    return SECFailure;
+}
+
+
+/* AES_CreateContext
+ *
+ * create a new context for Rijndael operations
+ */
+AESContext *
+AES_CreateContext(const unsigned char *key, const unsigned char *iv, 
+                  int mode, int encrypt,
+                  unsigned int keysize, unsigned int blocksize)
+{
+    AESContext *cx = AES_AllocateContext();
+    if (cx) {
+	SECStatus rv = AES_InitContext(cx, key, keysize, iv, mode, encrypt,
+				       blocksize);
+	if (rv != SECSuccess) {
+	    AES_DestroyContext(cx, PR_TRUE);
+	    cx = NULL;
+	}
+    }
+    return cx;
 }
 
 /*
@@ -1051,8 +1068,7 @@ cleanup:
 void 
 AES_DestroyContext(AESContext *cx, PRBool freeit)
 {
-    PORT_ZFree(cx->expandedKey, cx->Nb * (cx->Nr + 1));
-    memset(cx, 0, sizeof *cx);
+/*  memset(cx, 0, sizeof *cx); */
     if (freeit)
 	PORT_Free(cx);
 }