Fixed MODE_OFB requiring padding

Closes: https://bugs.launchpad.net/pycrypto/+bug/996193
Closes: https://github.com/dlitz/pycrypto/pull/26
[dlitz: Squashed and fixed whitespace.]

1 files changed, 42 insertions, 20 deletions

diff --git a/src/block_template.c b/src/block_template.c
index 436a4cc..964a6f3 100644
--- a/src/block_template.c
+++ b/src/block_template.c
@@ -248,6 +248,7 @@ ALG_Encrypt(ALGobject *self, PyObject *args)
 	}
 	if ( (len % BLOCK_SIZE) !=0 && 
 	     (self->mode!=MODE_CFB) &&
+	     (self->mode!=MODE_OFB) &&
 	     (self->mode!=MODE_CTR))
 	{
 		PyErr_Format(PyExc_ValueError, 
@@ -322,15 +323,48 @@ ALG_Encrypt(ALGobject *self, PyObject *args)
 		break;
 
 	case(MODE_OFB):
-		for(i=0; i<len; i+=BLOCK_SIZE) 
+		/* OFB mode is a stream cipher whose keystream is generated by encrypting the previous ciphered output.
+		 * - self->IV stores the current keystream block
+		 * - self->count indicates the current offset within the current keystream block
+		 * - str stores the input string
+		 * - buffer stores the output string
+		 * - len indicates the length of the input and output strings
+		 * - i indicates the current offset within the input and output strings
+		 * (len-i) is the number of bytes remaining to encrypt
+		 * (BLOCK_SIZE-self->count) is the number of bytes remaining in the current keystream block
+		 */
+		i = 0;
+		while(i < len)
 		{
+			/* If we don't need more than what remains of the current keystream block, then just XOR it in */
+			if (len-i <= BLOCK_SIZE-self->count) { /* remaining_bytes_to_encrypt <= remaining_bytes_in_IV */
+				/* XOR until the input is used up */
+				for(j=0; j<(len-i); j++) {
+					assert(i+j < len);
+					assert(self->count+j < BLOCK_SIZE);
+					buffer[i+j] = self->IV[self->count+j] ^ str[i+j];
+				}
+				self->count += len-i;
+				i = len;
+				continue;
+			}
+
+			/* Use up the current keystream block */
+			for(j=0; j<BLOCK_SIZE-self->count; j++) {
+				assert(i+j < len);
+				assert(self->count+j < BLOCK_SIZE);
+				buffer[i+j] = self->IV[self->count+j] ^ str[i+j];
+			}
+			i += BLOCK_SIZE-self->count;
+			self->count = BLOCK_SIZE;
+
+			/* Generate a new keystream block */
 			block_encrypt(&(self->st), self->IV, temp);
 			memcpy(self->IV, temp, BLOCK_SIZE);
-			for(j=0; j<BLOCK_SIZE; j++)
-			{
-				buffer[i+j] = str[i+j] ^ temp[j];
-			}
-		}      
+
+			/* Move the pointer to the start of the keystream */
+			self->count = 0;
+		}
 		break;
 
 	case(MODE_CTR):
@@ -464,8 +498,8 @@ ALG_Decrypt(ALGobject *self, PyObject *args)
 	int i, j, len;
 	PyObject *result;
 
-	/* CTR mode decryption is identical to encryption */
-	if (self->mode == MODE_CTR)
+	/* CTR and OFB mode decryption is identical to encryption */
+	if (self->mode == MODE_CTR || self->mode == MODE_OFB)
 		return ALG_Encrypt(self, args);
 
 	if (!PyArg_Parse(args, "s#", &str, &len))
@@ -547,18 +581,6 @@ ALG_Decrypt(ALGobject *self, PyObject *args)
 		}
 		break;
 
-	case (MODE_OFB):
-		for(i=0; i<len; i+=BLOCK_SIZE) 
-		{
-			block_encrypt(&(self->st), self->IV, temp);
-			memcpy(self->IV, temp, BLOCK_SIZE);
-			for(j=0; j<BLOCK_SIZE; j++)
-			{
-				buffer[i+j] = str[i+j] ^ self->IV[j];
-			}
-		}      
-		break;
-
 	default:
 		Py_BLOCK_THREADS;
 		PyErr_Format(PyExc_SystemError,