cbcutil: optimize memory usage when adding padding and mac

Author: tulir

util/cbcutil/cbc.go

@@ -108,10 +108,16 @@ func DecryptFile(key, iv []byte, file File) error {
 Encrypt is a function that encrypts plaintext with a given key and an optional initialization vector(iv).
 */
 func Encrypt(key, iv, plaintext []byte) ([]byte, error) {
-	plaintext = pad(plaintext, aes.BlockSize)
+	sizeOfLastBlock := len(plaintext) % aes.BlockSize
+	paddingLen := aes.BlockSize - sizeOfLastBlock
+	plaintextStart := plaintext[:len(plaintext)-sizeOfLastBlock]
+	lastBlock := append(plaintext[len(plaintext)-sizeOfLastBlock:], bytes.Repeat([]byte{byte(paddingLen)}, paddingLen)...)
 
-	if len(plaintext)%aes.BlockSize != 0 {
-		return nil, fmt.Errorf("plaintext is not a multiple of the block size: %d / %d", len(plaintext), aes.BlockSize)
+	if len(plaintextStart)%aes.BlockSize != 0 {
+		panic(fmt.Errorf("plaintext is not the correct size: %d %% %d != 0", len(plaintextStart), aes.BlockSize))
+	}
+	if len(lastBlock) != aes.BlockSize {
+		panic(fmt.Errorf("last block is not the correct size: %d != %d", len(lastBlock), aes.BlockSize))
 	}
 
 	block, err := aes.NewCipher(key)
@@ -121,30 +127,26 @@ func Encrypt(key, iv, plaintext []byte) ([]byte, error) {
 
 	var ciphertext []byte
 	if iv == nil {
-		ciphertext = make([]byte, aes.BlockSize+len(plaintext))
+		ciphertext = make([]byte, aes.BlockSize+len(plaintext)+paddingLen)
 		iv := ciphertext[:aes.BlockSize]
 		if _, err := io.ReadFull(rand.Reader, iv); err != nil {
 			return nil, err
 		}
 
 		cbc := cipher.NewCBCEncrypter(block, iv)
-		cbc.CryptBlocks(ciphertext[aes.BlockSize:], plaintext)
+		cbc.CryptBlocks(ciphertext[aes.BlockSize:], plaintextStart)
+		cbc.CryptBlocks(ciphertext[aes.BlockSize+len(plaintextStart):], lastBlock)
 	} else {
-		ciphertext = make([]byte, len(plaintext))
+		ciphertext = make([]byte, len(plaintext)+paddingLen, len(plaintext)+paddingLen+10)
 
 		cbc := cipher.NewCBCEncrypter(block, iv)
-		cbc.CryptBlocks(ciphertext, plaintext)
+		cbc.CryptBlocks(ciphertext, plaintextStart)
+		cbc.CryptBlocks(ciphertext[len(plaintextStart):], lastBlock)
 	}
 
 	return ciphertext, nil
 }
 
-func pad(ciphertext []byte, blockSize int) []byte {
-	padding := blockSize - len(ciphertext)%blockSize
-	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
-	return append(ciphertext, padtext...)
-}
-
 func unpad(src []byte) ([]byte, error) {
 	length := len(src)
 	padLen := int(src[length-1])