From 0734640e14db198f343144e268ba2209a83eaa2a Mon Sep 17 00:00:00 2001
From: Keegan McAllister <mcallister.keegan@gmail.com>
Date: Sat, 24 Mar 2012 21:37:43 -0400
Subject: [PATCH] Add a class for aligned buffers

This simplifies the core crypto routines, especially the error handling.  In
fact there was already one error path where we were failing to call free().
---
 src/crypto/crypto.cc | 47 ++++++++++++++++++--------------------------
 src/crypto/crypto.h  | 23 ++++++++++++++++++++++
 2 files changed, 42 insertions(+), 28 deletions(-)

diff --git a/src/crypto/crypto.cc b/src/crypto/crypto.cc
index e4f928f..8179b1b 100644
--- a/src/crypto/crypto.cc
+++ b/src/crypto/crypto.cc
@@ -46,7 +46,8 @@ long int myatoi( const char *str )
   return ret;
 }
 
-static void * aligned_alloc( int len )
+AlignedBuffer::AlignedBuffer( size_t len, const char *data )
+  : m_len( len ), m_data( NULL )
 {
   void *ptr = NULL;
 
@@ -72,7 +73,11 @@ static void * aligned_alloc( int len )
   }
 #endif /* !defined(HAVE_POSIX_MEMALIGN) */
 
-  return ptr;
+  m_data = (char *) ptr;
+
+  if ( data ) {
+    memcpy( m_data, data, len );
+  }
 }
 
 Base64Key::Base64Key( string printable_key )
@@ -193,10 +198,8 @@ string Session::encrypt( Message plaintext )
   const size_t pt_len = plaintext.text.size();
   const int ciphertext_len = pt_len + 16;
 
-  char *ciphertext = (char *)aligned_alloc( ciphertext_len );
-  char *pt = (char *)aligned_alloc( pt_len );
-
-  memcpy( pt, plaintext.text.data(), plaintext.text.size() );
+  AlignedBuffer ciphertext( ciphertext_len );
+  AlignedBuffer pt( pt_len, plaintext.text.data() );
 
   if ( (uint64_t( plaintext.nonce.data() ) & 0xf) != 0 ) {
     throw CryptoException( "Bad alignment." );
@@ -204,15 +207,13 @@ string Session::encrypt( Message plaintext )
 
   if ( ciphertext_len != ae_encrypt( ctx,                                     /* ctx */
 				     plaintext.nonce.data(),                  /* nonce */
-				     pt,                                      /* pt */
-				     pt_len,                                  /* pt_len */
+				     pt.data(),                               /* pt */
+				     pt.len(),                                /* pt_len */
 				     NULL,                                    /* ad */
 				     0,                                       /* ad_len */
-				     ciphertext,                              /* ct */
+				     ciphertext.data(),                       /* ct */
 				     NULL,                                    /* tag */
 				     AE_FINALIZE ) ) {                        /* final */
-    free( pt );
-    free( ciphertext );
     throw CryptoException( "ae_encrypt() returned error." );
   }
 
@@ -235,14 +236,10 @@ string Session::encrypt( Message plaintext )
      client use the same key, so we actually need to die after 2^47 blocks.
   */
   if ( blocks_encrypted >> 47 ) {
-    free( pt );
-    free( ciphertext );
     throw CryptoException( "Encrypted 2^47 blocks.", true );
   }
 
-  string text( (char *)ciphertext, ciphertext_len );
-  free( pt );
-  free( ciphertext );
+  string text( ciphertext.data(), ciphertext.len() );
 
   return plaintext.nonce.cc_str() + text;
 }
@@ -264,28 +261,22 @@ Message Session::decrypt( string ciphertext )
   }
 
   Nonce __attribute__((__aligned__ (16))) nonce( str, 8 );
-  char *body = (char *)aligned_alloc( body_len );
-  memcpy( body, str + 8, body_len );
-
-  char *plaintext = (char *)aligned_alloc( pt_len );
+  AlignedBuffer body( body_len, str + 8 );
+  AlignedBuffer plaintext( pt_len );
 
   if ( pt_len != ae_decrypt( ctx,               /* ctx */
 			     nonce.data(),      /* nonce */
-			     body,              /* ct */
-			     body_len,          /* ct_len */
+			     body.data(),       /* ct */
+			     body.len(),        /* ct_len */
 			     NULL,              /* ad */
 			     0,                 /* ad_len */
-			     plaintext,         /* pt */
+			     plaintext.data(),  /* pt */
 			     NULL,              /* tag */
 			     AE_FINALIZE ) ) {  /* final */
-    free( plaintext );
-    free( body );
     throw CryptoException( "Packet failed integrity check." );
   }
 
-  Message ret( nonce, string( plaintext, pt_len ) );
-  free( plaintext );
-  free( body );
+  Message ret( nonce, string( plaintext.data(), plaintext.len() ) );
 
   return ret;
 }
diff --git a/src/crypto/crypto.h b/src/crypto/crypto.h
index 1a2d897..ddbf12f 100644
--- a/src/crypto/crypto.h
+++ b/src/crypto/crypto.h
@@ -23,6 +23,7 @@
 #include <string>
 #include <string.h>
 #include <stdint.h>
+#include <stdlib.h>
 
 using std::string;
 
@@ -37,6 +38,28 @@ namespace Crypto {
       : text( s_text ), fatal( s_fatal ) {};
   };
 
+  /* 16-byte-aligned buffer, with length. */
+  class AlignedBuffer {
+  private:
+    size_t m_len;
+    char *m_data;
+
+  public:
+    AlignedBuffer( size_t len, const char *data = NULL );
+
+    ~AlignedBuffer() {
+      free( m_data );
+    }
+
+    char * data( void ) const { return m_data; }
+    size_t len( void )  const { return m_len;  }
+
+  private:
+    /* Not implemented */
+    AlignedBuffer( const AlignedBuffer& );
+    AlignedBuffer& operator=( const AlignedBuffer& );
+  };
+
   class Base64Key {
   private:
     unsigned char key[ 16 ];