From b8f42d784d6580a973da35268799710dd3008ab5 Mon Sep 17 00:00:00 2001
From: Berendea Nicolae <berendeanicolae@users.noreply.github.com>
Date: Sat, 25 Jun 2016 11:58:16 +0300
Subject: [PATCH] Update modes.cpp

---
 modes.cpp | 81 +++++++++++++++++++++++++++----------------------------
 1 file changed, 39 insertions(+), 42 deletions(-)

diff --git a/modes.cpp b/modes.cpp
index df3d4d86..2935b34c 100644
--- a/modes.cpp
+++ b/modes.cpp
@@ -201,30 +201,32 @@ void CBC_CTS_Encryption::ProcessLastBlock(byte *outString, const byte *inString,
 		// steal from IV
 		memcpy(outString, m_register, length);
 		outString = m_stolenIV;
+
+		// output last full ciphertext block
+		xorbuf(m_register, inString, length);
+		m_cipher->ProcessBlock(m_register);
+		memcpy(outString, m_register, BlockSize());
 	}
 	else
 	{
 		// steal from next to last block
 		xorbuf(m_register, inString, BlockSize());
 		m_cipher->ProcessBlock(m_register);
-		if (inString == outString)
-		{
-			length -= BlockSize();
-			memmove(outString, inString+BlockSize(), length);
-			memcpy(outString+BlockSize(), m_register, length);
-		}
-		else
-		{
-			inString += BlockSize();
-			length -= BlockSize();
-			memcpy(outString+BlockSize(), m_register, length);
-		}
-	}
+		length -= BlockSize();
 
-	// output last full ciphertext block
-	xorbuf(m_register, inString, length);
-	m_cipher->ProcessBlock(m_register);
-	memcpy(outString, m_register, BlockSize());
+		// output last full ciphertext block
+		if (inString == outString) {
+			memcpy(m_buffer, outString+BlockSize(), length);
+			memcpy(outString+BlockSize(), m_register, length);
+			xorbuf(m_register, m_buffer, length);
+		}
+		else {
+			memcpy(outString+BlockSize(), m_register, length);
+			xorbuf(m_register, inString+BlockSize(), length);
+		}
+		m_cipher->ProcessBlock(m_register);
+		memcpy(outString, m_register, BlockSize());
+	}
 }
 
 // Thanks to Zireael, http://github.com/weidai11/cryptopp/pull/46
@@ -259,40 +261,35 @@ void CBC_CTS_Decryption::ProcessLastBlock(byte *outString, const byte *inString,
 	{
 		pn = inString;
 		pn1 = m_register;
-	}
-	else
-	{
-		pn = inString + BlockSize();
-		pn1 = inString;
-		length -= BlockSize();
-	}
 
-	// decrypt last partial plaintext block
-	memcpy(m_temp, pn1, BlockSize());
-	m_cipher->ProcessBlock(m_temp);
-	xorbuf(m_temp, pn, length);
+		// decrypt last partial plaintext block
+		memcpy(m_temp, pn1, BlockSize());
+		m_cipher->ProcessBlock(m_temp);
+		xorbuf(m_temp, pn, length);
 
-	if (stealIV)
 		memcpy(outString, m_temp, length);
+	}
 	else
 	{
-		if (inString == outString)
-		{
-			memmove(outString, inString+BlockSize(), length);
+		length -= BlockSize();
+
+		// decrypt last partial plaintext block
+		memcpy(m_temp, inString, BlockSize());
+		m_cipher->ProcessBlock(m_temp);
+		xorbuf(m_temp, inString+BlockSize(), length);
+
+		// decrypt next to last plaintext block
+		if (inString == outString) {
+			memcpy(m_buffer, outString+BlockSize(), length);
 			memcpy(outString+BlockSize(), m_temp, length);
-			// decrypt next to last plaintext block
-			memcpy(m_temp, pn1, length);
-			m_cipher->ProcessBlock(m_temp);
-			xorbuf(outString, m_temp, m_register, BlockSize());
+			memcpy(m_temp, m_buffer, length);
 		}
-		else
-		{
+		else {
 			memcpy(outString+BlockSize(), m_temp, length);
-			// decrypt next to last plaintext block
-			memcpy(m_temp, pn, length);
-			m_cipher->ProcessBlock(m_temp);
-			xorbuf(outString, m_temp, m_register, BlockSize());
+			memcpy(m_temp, inString+BlockSize(), length);
 		}
+		m_cipher->ProcessBlock(m_temp);
+		xorbuf(outString, m_temp, m_register, BlockSize());
 	}
 }