From 4ef971f63cec414f19c9f5cf4c682edbd847c7ea Mon Sep 17 00:00:00 2001
From: orignal <romakoshelkin@yandex.ru>
Date: Thu, 8 May 2014 15:11:38 -0400
Subject: [PATCH] invert expanded key schedule for decryption

---
 aes.cpp | 42 +++++++++++++++++++++++++++++++++++-------
 aes.h   | 24 +++++++++++++++++++-----
 2 files changed, 54 insertions(+), 12 deletions(-)

diff --git a/aes.cpp b/aes.cpp
index 0c3d723f..1529487b 100644
--- a/aes.cpp
+++ b/aes.cpp
@@ -33,7 +33,7 @@ namespace crypto
 		"add $32, %%rcx \n"	
 		
 
-	void ECBCryptoAESNI::SetKey (const uint8_t * key)
+	void ECBCryptoAESNI::ExpandKey (const uint8_t * key)
 	{
 		__asm__
 		(
@@ -68,11 +68,11 @@ namespace crypto
 			"movups	%%xmm1, (%%rcx) \n"
 			: // output
 			: "S" (key), "D" (m_KeySchedule) // input
-			: "%rcx" // modified
+			: "%rcx" // clogged
 		);
 	}
 
-	void ECBCryptoAESNI::Encrypt (const ChipherBlock * in, ChipherBlock * out)
+	void ECBEncryptionAESNI::Encrypt (const ChipherBlock * in, ChipherBlock * out)
 	{
 		__asm__
 		(
@@ -92,11 +92,11 @@ namespace crypto
 			"aesenc	208(%%rdx), %%xmm0 \n"
 			"aesenclast	224(%%rdx), %%xmm0 \n"
 			"movups	%%xmm0, (%%rdi) \n"	
-			: : "d" ((uint64_t)m_KeySchedule), "S" (in), "D" (out)
+			: : "d" (m_KeySchedule), "S" (in), "D" (out)
 		);
-	}	
+	}		
 
-	void ECBCryptoAESNI::Decrypt (const ChipherBlock * in, ChipherBlock * out)
+	void ECBDecryptionAESNI::Decrypt (const ChipherBlock * in, ChipherBlock * out)
 	{
 		__asm__
 		(
@@ -116,10 +116,38 @@ namespace crypto
 			"aesdec	16(%%rdx), %%xmm0 \n"
 			"aesdeclast (%%rdx), %%xmm0 \n"
 			"movups	%%xmm0, (%%rdi) \n"	
-			: : "d" ((uint64_t)m_KeySchedule), "S" (in), "D" (out)
+			: : "d" (m_KeySchedule), "S" (in), "D" (out)
 		);		
 	}
 
+	#define CallAESIMC(offset) \
+		"movups "#offset"(%%rdx), %%xmm0 \n"	\
+		"aesimc %%xmm0, %%xmm0 \n" \
+		"movups %%xmm0, "#offset"(%%rdx) \n" 
+
+	void ECBDecryptionAESNI::SetKey (const uint8_t * key)
+	{
+		ExpandKey (key); // expand encryption key first
+		// then  invert it using aesimc
+		__asm__
+		(
+			CallAESIMC(16)
+			CallAESIMC(32)
+			CallAESIMC(48)
+			CallAESIMC(64)
+			CallAESIMC(80)
+			CallAESIMC(96)
+			CallAESIMC(112)
+			CallAESIMC(128)
+			CallAESIMC(144)
+			CallAESIMC(160)
+			CallAESIMC(176)
+			CallAESIMC(192)
+			CallAESIMC(208)
+			: : "d" (m_KeySchedule)
+		);
+	}
+
 #endif		
 
 
diff --git a/aes.h b/aes.h
index 2fbe5b0d..0b7866ab 100644
--- a/aes.h
+++ b/aes.h
@@ -18,16 +18,30 @@ namespace crypto
 #ifdef __x86_64__
 	// AES-NI assumed
 	class ECBCryptoAESNI
+	{	
+		protected:
+
+			void ExpandKey (const uint8_t * key);
+		
+		protected:
+
+			uint32_t m_KeySchedule[4*(14+1)]; // 14 rounds for AES-256
+	};	
+
+	class ECBEncryptionAESNI: public ECBCryptoAESNI
 	{
 		public:
 		
-			void SetKey (const uint8_t * key);
+			void SetKey (const uint8_t * key) { ExpandKey (key); };
 			void Encrypt (const ChipherBlock * in, ChipherBlock * out);	
-			void Decrypt (const ChipherBlock * in, ChipherBlock * out);		
-
-		private:
+	};	
 
-			uint32_t m_KeySchedule[4*(14+1)]; // 14 rounds for AES-256
+	class ECBDecryptionAESNI: public ECBCryptoAESNI
+	{
+		public:
+		
+			void SetKey (const uint8_t * key);
+			void Decrypt (const ChipherBlock * in, ChipherBlock * out);		
 	};	
 
 #endif