From dd1a79812815d7283a6050ea17d2ecda6794ccb1 Mon Sep 17 00:00:00 2001
From: orignal <i2porignal@yandex.ru>
Date: Tue, 18 Nov 2014 12:11:45 -0500
Subject: [PATCH] AES buffer aligned to 16 bytes

---
 aes.cpp | 16 ++++------------
 aes.h   | 31 ++++++++++++++++++++++++++-----
 2 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/aes.cpp b/aes.cpp
index b03a46fa..1626682c 100644
--- a/aes.cpp
+++ b/aes.cpp
@@ -8,14 +8,6 @@ namespace crypto
 {
 
 #ifdef AESNI
-
-	ECBCryptoAESNI::ECBCryptoAESNI ()
-	{
-		m_KeySchedule = m_UnalignedBuffer;
-		uint8_t rem = ((uint64_t)m_KeySchedule) & 0x0f;
-		if (rem)
-			m_KeySchedule += (16 - rem);
-	}	
 	
 	#define KeyExpansion256(round0,round1) \
 		"pshufd	$0xff, %%xmm2, %%xmm2 \n" \
@@ -73,7 +65,7 @@ namespace crypto
 			"pxor %%xmm2, %%xmm1 \n"
 			"movups	%%xmm1, 224(%[sched]) \n"
 			: // output
-			: [key]"r"((const uint8_t *)key), [sched]"r"(m_KeySchedule) // input
+			: [key]"r"((const uint8_t *)key), [sched]"r"(GetKeySchedule ()) // input
 			: "%xmm1", "%xmm2", "%xmm3", "%xmm4" // clogged
 		);
 	}
@@ -102,7 +94,7 @@ namespace crypto
 			"movups	(%[in]), %%xmm0 \n"
 			EncryptAES256(sched)
 			"movups	%%xmm0, (%[out]) \n"	
-			: : [sched]"r"(m_KeySchedule), [in]"r"(in), [out]"r"(out) : "%xmm0"
+			: : [sched]"r"(GetKeySchedule ()), [in]"r"(in), [out]"r"(out) : "%xmm0"
 		);
 	}		
 
@@ -130,7 +122,7 @@ namespace crypto
 			"movups	(%[in]), %%xmm0 \n"
 			DecryptAES256(sched)
 			"movups	%%xmm0, (%[out]) \n"	
-			: : [sched]"r"(m_KeySchedule), [in]"r"(in), [out]"r"(out) : "%xmm0"
+			: : [sched]"r"(GetKeySchedule ()), [in]"r"(in), [out]"r"(out) : "%xmm0"
 		);		
 	}
 
@@ -158,7 +150,7 @@ namespace crypto
 			CallAESIMC(176)
 			CallAESIMC(192)
 			CallAESIMC(208)
-			: : [shed]"r"(m_KeySchedule) : "%xmm0"
+			: : [shed]"r"(GetKeySchedule ()) : "%xmm0"
 		);
 	}
 
diff --git a/aes.h b/aes.h
index 3d7e9c7e..866c39ec 100644
--- a/aes.h
+++ b/aes.h
@@ -24,22 +24,43 @@ namespace crypto
 
 	typedef i2p::data::Tag<32> AESKey;
 	
+	template<size_t sz>
+	class AESAlignedBuffer // 16 bytes alignment
+	{
+		public:
+		
+			AESAlignedBuffer ()
+			{
+				m_Buf = m_UnalignedBuffer;
+				uint8_t rem = ((uint64_t)m_Buf) & 0x0f;
+				if (rem)
+					m_Buf += (16 - rem);
+			}
+		
+			operator uint8_t * () { return m_Buf; };
+			operator const uint8_t * () const { return m_Buf; };
+
+		private:
+
+			uint8_t m_UnalignedBuffer[sz + 15]; // up to 15 bytes alignment
+			uint8_t * m_Buf;
+	};			
+
+
 #ifdef AESNI
 	class ECBCryptoAESNI
 	{	
 		public:
 
-			ECBCryptoAESNI ();
 			uint8_t * GetKeySchedule () { return m_KeySchedule; };
-			
+
 		protected:
 
 			void ExpandKey (const AESKey& key);
 		
-		protected:
+		private:
 
-			uint8_t * m_KeySchedule; // start of 16 bytes boundary of m_UnalignedBuffer
-			uint8_t m_UnalignedBuffer[256]; // 14 rounds for AES-256, 240 + 16 bytes
+			AESAlignedBuffer<240> m_KeySchedule;  // 14 rounds for AES-256, 240 bytes
 	};	
 
 	class ECBEncryptionAESNI: public ECBCryptoAESNI