Update test for elimination of ae_allocate() / ae_free().

1 files changed, 18 insertions, 14 deletions

diff --git a/src/tests/ocb-aes.cc b/src/tests/ocb-aes.cc
index 4f8ded0..63418d2 100644
--- a/src/tests/ocb-aes.cc
+++ b/src/tests/ocb-aes.cc
@@ -46,22 +46,23 @@ bool equal( const AlignedBuffer &a, const AlignedBuffer &b ) {
     && !memcmp( a.data(), b.data(), a.len() );
 }
 
-ae_ctx *get_ctx( const AlignedBuffer &key ) {
-  ae_ctx *ctx = ae_allocate( NULL );
-  fatal_assert( ctx );
-  fatal_assert( AE_SUCCESS == ae_init( ctx, key.data(), key.len(), NONCE_LEN, TAG_LEN ) );
-  return ctx;
+AlignedBuffer *get_ctx( const AlignedBuffer &key ) {
+  AlignedBuffer *ctx_buf = new AlignedBuffer( ae_ctx_sizeof() );
+  fatal_assert( ctx_buf );
+  fatal_assert( AE_SUCCESS == ae_init( (ae_ctx *)ctx_buf->data(), key.data(), key.len(), NONCE_LEN, TAG_LEN ) );
+  return ctx_buf;
 }
 
-void scrap_ctx( ae_ctx *ctx ) {
-  fatal_assert( AE_SUCCESS == ae_clear( ctx ) );
-  ae_free( ctx );
+void scrap_ctx( AlignedBuffer *ctx_buf ) {
+  fatal_assert( AE_SUCCESS == ae_clear( (ae_ctx *)ctx_buf->data() ) );
+  delete ctx_buf;
 }
 
 void test_encrypt( const AlignedBuffer &key, const AlignedBuffer &nonce,
                    const AlignedBuffer &plaintext, const AlignedBuffer &assoc,
                    const AlignedBuffer &expected_ciphertext ) {
-  ae_ctx *ctx = get_ctx( key );
+  AlignedBuffer *ctx_buf = get_ctx( key );
+  ae_ctx *ctx = (ae_ctx *)ctx_buf->data();
 
   AlignedBuffer observed_ciphertext( plaintext.len() + TAG_LEN );
 
@@ -79,14 +80,15 @@ void test_encrypt( const AlignedBuffer &key, const AlignedBuffer &nonce,
   fatal_assert( ret == int( expected_ciphertext.len() ) );
   fatal_assert( equal( expected_ciphertext, observed_ciphertext ) );
 
-  scrap_ctx( ctx );
+  scrap_ctx( ctx_buf );
 }
 
 void test_decrypt( const AlignedBuffer &key, const AlignedBuffer &nonce,
                    const AlignedBuffer &ciphertext, const AlignedBuffer &assoc,
                    const AlignedBuffer &expected_plaintext,
                    bool valid ) {
-  ae_ctx *ctx = get_ctx( key );
+  AlignedBuffer *ctx_buf = get_ctx( key );
+  ae_ctx *ctx = (ae_ctx *)ctx_buf->data();
 
   AlignedBuffer observed_plaintext( ciphertext.len() - TAG_LEN );
 
@@ -110,7 +112,7 @@ void test_decrypt( const AlignedBuffer &key, const AlignedBuffer &nonce,
     fatal_assert( ret == AE_INVALID );
   }
 
-  scrap_ctx( ctx );
+  scrap_ctx( ctx_buf );
 }
 
 void test_vector( const char *key_p, const char *nonce_p,
@@ -451,7 +453,9 @@ void test_iterative( void ) {
   /* Key is always all zeros */
   AlignedBuffer key( KEY_LEN );
   memset( key.data(), 0, KEY_LEN );
-  ae_ctx *ctx = get_ctx( key );
+
+  AlignedBuffer *ctx_buf = get_ctx( key );
+  ae_ctx *ctx = (ae_ctx *)ctx_buf->data();
 
   AlignedBuffer nonce( NONCE_LEN );
   memset( nonce.data(), 0, NONCE_LEN );
@@ -513,7 +517,7 @@ void test_iterative( void ) {
   AlignedBuffer correct( TAG_LEN, "\xB2\xB4\x1C\xBF\x9B\x05\x03\x7D\xA7\xF1\x6C\x24\xA3\x5C\x1C\x94" );
   fatal_assert( equal( out, correct ) );
 
-  scrap_ctx( ctx );
+  scrap_ctx( ctx_buf );
 
   if ( verbose ) {
     printf( "iterative PASSED\n\n" );