AES.Encryptor and AES.Decryptor allows encrypt/decrypt in chunks. Revert AnyGenerators to workers.

CryptoSwift.xcodeproj/project.pbxproj

@@ -227,6 +227,10 @@
 		75CB934F1C8F609D0087740D /* BlockModeOptions.swift in Sources */ = {isa = PBXBuildFile; fileRef = 75CB934D1C8F609D0087740D /* BlockModeOptions.swift */; };
 		75CB93501C8F609D0087740D /* BlockModeOptions.swift in Sources */ = {isa = PBXBuildFile; fileRef = 75CB934D1C8F609D0087740D /* BlockModeOptions.swift */; };
 		75CB93511C8F609D0087740D /* BlockModeOptions.swift in Sources */ = {isa = PBXBuildFile; fileRef = 75CB934D1C8F609D0087740D /* BlockModeOptions.swift */; };
+		75DB81A31CDBFDC700ED181A /* BlockModeWorker.swift in Sources */ = {isa = PBXBuildFile; fileRef = 75DB81A21CDBFDC700ED181A /* BlockModeWorker.swift */; };
+		75DB81A41CDBFDC700ED181A /* BlockModeWorker.swift in Sources */ = {isa = PBXBuildFile; fileRef = 75DB81A21CDBFDC700ED181A /* BlockModeWorker.swift */; };
+		75DB81A51CDBFDC700ED181A /* BlockModeWorker.swift in Sources */ = {isa = PBXBuildFile; fileRef = 75DB81A21CDBFDC700ED181A /* BlockModeWorker.swift */; };
+		75DB81A61CDBFDC700ED181A /* BlockModeWorker.swift in Sources */ = {isa = PBXBuildFile; fileRef = 75DB81A21CDBFDC700ED181A /* BlockModeWorker.swift */; };
 		80545D131CA9FECD00474A99 /* Bit.swift in Sources */ = {isa = PBXBuildFile; fileRef = 80545D121CA9FECD00474A99 /* Bit.swift */; };
 		80545D141CA9FECD00474A99 /* Bit.swift in Sources */ = {isa = PBXBuildFile; fileRef = 80545D121CA9FECD00474A99 /* Bit.swift */; };
 		80545D151CA9FECD00474A99 /* Bit.swift in Sources */ = {isa = PBXBuildFile; fileRef = 80545D121CA9FECD00474A99 /* Bit.swift */; };
@@ -385,6 +389,7 @@
 		75CB93481C8F60700087740D /* ECB.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = ECB.swift; path = Sources/CryptoSwift/BlockMode/ECB.swift; sourceTree = SOURCE_ROOT; };
 		75CB934D1C8F609D0087740D /* BlockModeOptions.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = BlockModeOptions.swift; path = Sources/CryptoSwift/BlockMode/BlockModeOptions.swift; sourceTree = SOURCE_ROOT; };
 		75D614BF1BD844F2001358B2 /* CryptoSwift.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = CryptoSwift.framework; sourceTree = BUILT_PRODUCTS_DIR; };
+		75DB81A21CDBFDC700ED181A /* BlockModeWorker.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = BlockModeWorker.swift; path = Sources/CryptoSwift/BlockMode/BlockModeWorker.swift; sourceTree = SOURCE_ROOT; };
 		75DF77721BC8EB59006E9520 /* CryptoSwift.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = CryptoSwift.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		80545D121CA9FECD00474A99 /* Bit.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = Bit.swift; path = Sources/CryptoSwift/Bit.swift; sourceTree = SOURCE_ROOT; };
 /* End PBXFileReference section */
@@ -569,6 +574,7 @@
 				7588034B1C8F8C33008C1576 /* BlockModeGenerator.swift */,
 				75CB934D1C8F609D0087740D /* BlockModeOptions.swift */,
 				75CB932F1C8F5F580087740D /* CipherBlockMode.swift */,
+				75DB81A21CDBFDC700ED181A /* BlockModeWorker.swift */,
 				75CB93241C8F5EC10087740D /* CBC.swift */,
 				75CB93341C8F5FCE0087740D /* PCBC.swift */,
 				75CB93391C8F5FFD0087740D /* CFB.swift */,
@@ -806,6 +812,7 @@
 				754C9D1C1C30800A0003D5FE /* String+LinuxFoundation.swift in Sources */,
 				757BC99D1C1CA5790093AAA9 /* PKCS7.swift in Sources */,
 				757BC9AD1C1CA5790093AAA9 /* SHA2.swift in Sources */,
+				75DB81A41CDBFDC700ED181A /* BlockModeWorker.swift in Sources */,
 				757BC8FF1C1CA56A0093AAA9 /* AES+Foundation.swift in Sources */,
 				7574E5FF1CD02C9300E96346 /* CSArrayType+Foundation.swift in Sources */,
 				757BC9491C1CA5790093AAA9 /* Authenticator.swift in Sources */,
@@ -865,6 +872,7 @@
 				754C9D1B1C30800A0003D5FE /* String+LinuxFoundation.swift in Sources */,
 				757BC99C1C1CA5790093AAA9 /* PKCS7.swift in Sources */,
 				757BC9AC1C1CA5790093AAA9 /* SHA2.swift in Sources */,
+				75DB81A31CDBFDC700ED181A /* BlockModeWorker.swift in Sources */,
 				757BC9161C1CA56A0093AAA9 /* Utils+Foundation.swift in Sources */,
 				7574E5FE1CD02C8B00E96346 /* CSArrayType+Foundation.swift in Sources */,
 				757BC9481C1CA5790093AAA9 /* Authenticator.swift in Sources */,
@@ -941,6 +949,7 @@
 				754C9D1D1C30800A0003D5FE /* String+LinuxFoundation.swift in Sources */,
 				757BC99E1C1CA5790093AAA9 /* PKCS7.swift in Sources */,
 				757BC9AE1C1CA5790093AAA9 /* SHA2.swift in Sources */,
+				75DB81A61CDBFDC700ED181A /* BlockModeWorker.swift in Sources */,
 				757BC9181C1CA56A0093AAA9 /* Utils+Foundation.swift in Sources */,
 				7574E6011CD02C9400E96346 /* CSArrayType+Foundation.swift in Sources */,
 				757BC94A1C1CA5790093AAA9 /* Authenticator.swift in Sources */,
@@ -1000,6 +1009,7 @@
 				754C9D1E1C30800A0003D5FE /* String+LinuxFoundation.swift in Sources */,
 				757BC99F1C1CA5790093AAA9 /* PKCS7.swift in Sources */,
 				757BC9AF1C1CA5790093AAA9 /* SHA2.swift in Sources */,
+				75DB81A51CDBFDC700ED181A /* BlockModeWorker.swift in Sources */,
 				757BC9011C1CA56A0093AAA9 /* AES+Foundation.swift in Sources */,
 				7574E6001CD02C9300E96346 /* CSArrayType+Foundation.swift in Sources */,
 				757BC94B1C1CA5790093AAA9 /* Authenticator.swift in Sources */,

CryptoSwiftTests/AESTests.swift

@@ -30,7 +30,6 @@ final class AESTests: XCTestCase {
         let key = "679fb1ddf7d81bee"
         let iv = "kdf67398DF7383fd"
         let input:[UInt8] = [0x62, 0x72, 0x61, 0x64, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
-        
         let expected:[UInt8] = [0xae,0x8c,0x59,0x95,0xb2,0x6f,0x8e,0x3d,0xb0,0x6f,0x0a,0xa5,0xfe,0xc4,0xf0,0xc2];
         
         let aes = try! AES(key: key, iv: iv, blockMode: .CBC, padding: NoPadding())
@@ -41,15 +40,8 @@ final class AESTests: XCTestCase {
     }
     
     func testAES_encrypt() {
-        let input:[UInt8] = [0x00, 0x11, 0x22, 0x33,
-            0x44, 0x55, 0x66, 0x77,
-            0x88, 0x99, 0xaa, 0xbb,
-            0xcc, 0xdd, 0xee, 0xff];
-        
-        let expected:[UInt8] = [0x69, 0xc4, 0xe0, 0xd8,
-            0x6a, 0x7b, 0x4, 0x30,
-            0xd8, 0xcd, 0xb7, 0x80,
-            0x70, 0xb4, 0xc5, 0x5a];
+        let input:[UInt8] = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff];
+        let expected:[UInt8] = [0x69, 0xc4, 0xe0, 0xd8, 0x6a, 0x7b, 0x4, 0x30, 0xd8, 0xcd, 0xb7, 0x80, 0x70, 0xb4, 0xc5, 0x5a];
         
         let aes = try! AES(key: aesKey, blockMode: .ECB, padding: NoPadding())
         let encrypted = try! aes.encrypt(input)
@@ -65,7 +57,6 @@ final class AESTests: XCTestCase {
         let expected:[UInt8] = [0x76,0x49,0xab,0xac,0x81,0x19,0xb2,0x46,0xce,0xe9,0x8e,0x9b,0x12,0xe9,0x19,0x7d];
 
         let aes = try! AES(key: key, iv:iv, blockMode: .CBC, padding: NoPadding())
-        XCTAssertTrue(aes.blockMode == .CBC, "Invalid block mode")
         let encrypted = try! aes.encrypt(plaintext)
         XCTAssertEqual(encrypted, expected, "encryption failed")
         let decrypted = try! aes.decrypt(encrypted)
@@ -79,13 +70,27 @@ final class AESTests: XCTestCase {
         let expected:[UInt8] = [0x76,0x49,0xab,0xac,0x81,0x19,0xb2,0x46,0xce,0xe9,0x8e,0x9b,0x12,0xe9,0x19,0x7d,0x89,0x64,0xe0,0xb1,0x49,0xc1,0x0b,0x7b,0x68,0x2e,0x6e,0x39,0xaa,0xeb,0x73,0x1c]
         
         let aes = try! AES(key: key, iv:iv, blockMode: .CBC, padding: PKCS7())
-        XCTAssertTrue(aes.blockMode == .CBC, "Invalid block mode")
         let encrypted = try! aes.encrypt(plaintext)
         XCTAssertEqual(encrypted, expected, "encryption failed")
         let decrypted = try! aes.decrypt(encrypted)
         XCTAssertEqual(decrypted, plaintext, "decryption failed")
     }
-    
+
+    func testAES_encrypt_cbc_with_padding_partial() {
+        let key:[UInt8] = [0x2b,0x7e,0x15,0x16,0x28,0xae,0xd2,0xa6,0xab,0xf7,0x15,0x88,0x09,0xcf,0x4f,0x3c];
+        let iv:[UInt8] = [0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07,0x08,0x09,0x0A,0x0B,0x0C,0x0D,0x0E,0x0F]
+        let plaintext:[UInt8] = [0x6b,0xc1,0xbe,0xe2,0x2e,0x40,0x9f,0x96,0xe9,0x3d,0x7e,0x11,0x73,0x93,0x17,0x2a,0x6b,0xc1,0xbe,0xe2,0x2e,0x40,0x9f,0x96,0xe9,0x3d,0x7e,0x11,0x73,0x93,0x17,0x2a]
+
+        let aes = try! AES(key: key, iv:iv, blockMode: .CBC, padding: PKCS7())
+        let encrypted = try! aes.encrypt(plaintext)
+
+        var partialEncrypted = [UInt8]()
+        var encryptor = aes.encryptor()
+        partialEncrypted.appendContentsOf(try! encryptor.update(Array(plaintext[0..<16]), isLast: false))
+        partialEncrypted.appendContentsOf(try! encryptor.update(Array(plaintext[16..<32]), isLast: true))
+        XCTAssertEqual(encrypted, partialEncrypted, "encryption failed")
+    }
+
     func testAES_encrypt_cfb() {
         let key:[UInt8] = [0x2b,0x7e,0x15,0x16,0x28,0xae,0xd2,0xa6,0xab,0xf7,0x15,0x88,0x09,0xcf,0x4f,0x3c];
         let iv:[UInt8] = [0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07,0x08,0x09,0x0A,0x0B,0x0C,0x0D,0x0E,0x0F]
@@ -93,7 +98,6 @@ final class AESTests: XCTestCase {
         let expected:[UInt8] = [0x3b,0x3f,0xd9,0x2e,0xb7,0x2d,0xad,0x20,0x33,0x34,0x49,0xf8,0xe8,0x3c,0xfb,0x4a];
         
         let aes = try! AES(key: key, iv:iv, blockMode: .CFB, padding: NoPadding())
-        XCTAssertTrue(aes.blockMode == .CFB, "Invalid block mode")
         let encrypted = try! aes.encrypt(plaintext)
         XCTAssertEqual(encrypted, expected, "encryption failed")
         let decrypted = try! aes.decrypt(encrypted)
@@ -117,7 +121,6 @@ final class AESTests: XCTestCase {
         let expected:[UInt8] = [0x3b,0x3f,0xd9,0x2e,0xb7,0x2d,0xad,0x20,0x33,0x34,0x49,0xf8,0xe8,0x3c,0xfb,0x4a];
 
         let aes = try! AES(key: key, iv:iv, blockMode: .OFB, padding: NoPadding())
-        XCTAssertTrue(aes.blockMode == .OFB, "Invalid block mode")
         let encrypted = try! aes.encrypt(plaintext)
         XCTAssertEqual(encrypted, expected, "encryption failed")
         let decrypted = try! aes.decrypt(encrypted)
@@ -131,7 +134,6 @@ final class AESTests: XCTestCase {
         let expected:[UInt8] = [0xdc,0x7e,0x84,0xbf,0xda,0x79,0x16,0x4b,0x7e,0xcd,0x84,0x86,0x98,0x5d,0x38,0x60];
 
         let aes = try! AES(key: key, iv:iv, blockMode: .OFB, padding: NoPadding())
-        XCTAssertTrue(aes.blockMode == .OFB, "Invalid block mode")
         let encrypted = try! aes.encrypt(plaintext)
         XCTAssertEqual(encrypted, expected, "encryption failed")
         let decrypted = try! aes.decrypt(encrypted)
@@ -145,7 +147,6 @@ final class AESTests: XCTestCase {
         let expected:[UInt8] = [0xf5,0x8c,0x4c,0x04,0xd6,0xe5,0xf1,0xba,0x77,0x9e,0xab,0xfb,0x5f,0x7b,0xfb,0xd6];
 
         let aes = try! AES(key: key, iv:iv, blockMode: .PCBC, padding: NoPadding())
-        XCTAssertTrue(aes.blockMode == .PCBC, "Invalid block mode")
         let encrypted = try! aes.encrypt(plaintext)
         print(encrypted.toHexString())
         XCTAssertEqual(encrypted, expected, "encryption failed")
@@ -160,7 +161,6 @@ final class AESTests: XCTestCase {
         let expected:[UInt8] = [0x87,0x4d,0x61,0x91,0xb6,0x20,0xe3,0x26,0x1b,0xef,0x68,0x64,0x99,0x0d,0xb6,0xce]
         
         let aes = try! AES(key: key, iv:iv, blockMode: .CTR, padding: NoPadding())
-        XCTAssertTrue(aes.blockMode == .CTR, "Invalid block mode")
         let encrypted = try! aes.encrypt(plaintext)
         XCTAssertEqual(encrypted, expected, "encryption failed")
         let decrypted = try! aes.decrypt(encrypted)
@@ -174,7 +174,6 @@ final class AESTests: XCTestCase {
         let expected:[UInt8] = [0x87,0x4d,0x61,0x91,0xb6,0x20,0xe3,0x26,0x1b,0xef,0x68,0x64,0x99,0x0d,0xb6,0xce,0x37]
 
         let aes = try! AES(key: key, iv:iv, blockMode: .CTR, padding: NoPadding())
-        XCTAssertTrue(aes.blockMode == .CTR, "Invalid block mode")
         let encrypted = try! aes.encrypt(plaintext)
         XCTAssertEqual(encrypted, expected, "encryption failed")
         let decrypted = try! aes.decrypt(encrypted)
@@ -248,7 +247,6 @@ final class AESTests: XCTestCase {
 
         let aes = try! AES(key: key, iv:iv, blockMode: .CBC, padding: PKCS7())
         let aes2 = try! AES(key: key2, iv:iv, blockMode: .CBC, padding: PKCS7())
-        XCTAssertTrue(aes.blockMode == .CBC, "Invalid block mode")
         let encrypted = try! aes.encrypt(plaintext)
         let decrypted = try? aes2.decrypt(encrypted)
         XCTAssertTrue(decrypted! != plaintext, "failed")

Sources/CryptoSwift/AES.swift

@@ -12,12 +12,12 @@ private typealias Key = SecureBytes
 
 final public class AES: BlockCipher {
     public enum Error: ErrorType {
-        case BlockSizeExceeded
+        case DataPaddingRequired
         case InvalidKeyOrInitializationVector
         case InvalidInitializationVector
     }
     
-    public enum AESVariant:Int {
+    public enum AESVariant: Int {
         case aes128 = 1, aes192, aes256
         
         var Nk:Int { // Nk words
@@ -32,8 +32,56 @@ final public class AES: BlockCipher {
             return Nk + 6
         }
     }
-    
-    public let blockMode:CipherBlockMode
+
+    public struct Encryptor {
+        private var worker: BlockModeWorker
+        let padding: Padding
+
+        init(aes: AES) {
+            self.padding = aes.padding;
+            self.worker = aes.cipherBlockMode.worker(aes.iv, cipherOperation: aes.encryptBlock)
+        }
+
+        mutating func update(bytes:[UInt8], isLast: Bool) throws -> [UInt8] {
+            if isLast {
+                let paddedBytes = padding.add(bytes, blockSize: AES.blockSize)
+                var result = [UInt8]()
+                for chunk in paddedBytes.chunks(AES.blockSize) ?? [] {
+                    result.appendContentsOf(worker.encrypt(chunk))
+                }
+                return result
+            }
+
+            return worker.encrypt(bytes)
+        }
+    }
+
+    public struct Decryptor {
+        private var worker: BlockModeWorker
+        let padding: Padding
+
+        init(aes: AES) {
+            self.padding = aes.padding;
+
+            switch (aes.cipherBlockMode) {
+                case .CFB, .OFB, .CTR:
+                    // CFB, OFB, CTR uses encryptBlock to decrypt
+                    self.worker = aes.cipherBlockMode.worker(aes.iv, cipherOperation: aes.encryptBlock)
+                default:
+                    self.worker = aes.cipherBlockMode.worker(aes.iv, cipherOperation: aes.decryptBlock)
+            }
+        }
+
+        mutating func update(bytes:[UInt8], isLast: Bool) throws -> [UInt8] {
+            let plaintext = worker.decrypt(bytes)
+            if isLast {
+                return padding.remove(plaintext, blockSize: AES.blockSize)
+            }
+            return plaintext
+        }
+    }
+
+    private let cipherBlockMode:CipherBlockMode
     public static let blockSize:Int = 16 // 128 /8
     
     public var variant:AESVariant {
@@ -92,7 +140,7 @@ final public class AES: BlockCipher {
     
     public init(key:[UInt8], iv:[UInt8]? = nil, blockMode:CipherBlockMode = .CBC, padding: Padding = PKCS7()) throws {
         self.key = Key(bytes: key)
-        self.blockMode = blockMode
+        self.cipherBlockMode = blockMode
         self.padding = padding
 
         if let iv = iv where iv.count > 0 {
@@ -108,30 +156,32 @@ final public class AES: BlockCipher {
         }
     }
 
-    /**
-    Encrypt message. If padding is necessary, then PKCS7 padding is added and needs to be removed after decryption.
-    
-    - parameter message: Plaintext data
-    - parameter padding: Optional padding
-    
-    - returns: Encrypted data
-    */
-    
-    public func encrypt(bytes:[UInt8]) throws -> [UInt8] {
-        let finalBytes = self.padding.add(bytes, blockSize: AES.blockSize)
+    public func encryptor() -> Encryptor {
+        return Encryptor(aes: self)
+    }
 
-        if blockMode.options.contains(.PaddingRequired) && (finalBytes.count % AES.blockSize != 0) {
-            throw Error.BlockSizeExceeded
-        }
+    public func decryptor() -> Decryptor {
+        return Decryptor(aes: self)
+    }
 
-        let blocks = finalBytes.chunks(AES.blockSize)
-        let encryptGenerator = blockMode.encryptGenerator(iv, cipherOperation: encryptBlock, inputGenerator: AnyGenerator<Array<UInt8>>(blocks.generate()))
+    /// Encrypt given bytes at once
+    ///
+    /// - parameter bytes: Plaintext data
+    /// - returns: Encrypted data
+    public func encrypt(bytes:[UInt8]) throws -> [UInt8] {
+        let chunks = bytes.chunks(AES.blockSize)
 
+        var oneTimeCryptor = Encryptor(aes: self)
         var out = [UInt8]()
         out.reserveCapacity(bytes.count)
-        while let processedBlock = encryptGenerator.next() {
-            out.appendContentsOf(processedBlock)
+        for (idx, block) in chunks.enumerate() {
+            out.appendContentsOf(try oneTimeCryptor.update(block, isLast: idx == max(0,chunks.count - 1)))
         }
+
+        if cipherBlockMode.options.contains(.PaddingRequired) && (out.count % AES.blockSize != 0) {
+            throw Error.DataPaddingRequired
+        }
+
         return out
     }
 
@@ -199,29 +249,23 @@ final public class AES: BlockCipher {
         return out
     }
     
+    /// Decrypt given bytes at once
+    ///
+    /// - parameter bytes: Ciphertext data
+    /// - returns: Plaintext data
     public func decrypt(bytes:[UInt8]) throws -> [UInt8] {
-        if blockMode.options.contains(.PaddingRequired) && (bytes.count % AES.blockSize != 0) {
-            throw Error.BlockSizeExceeded
+        if cipherBlockMode.options.contains(.PaddingRequired) && (bytes.count % AES.blockSize != 0) {
+            throw Error.DataPaddingRequired
         }
 
-        let blocks = bytes.chunks(AES.blockSize)
+        var oneTimeCryptor = Decryptor(aes: self)
+        let chunks = bytes.chunks(AES.blockSize)
         var out = [UInt8]()
         out.reserveCapacity(bytes.count)
-        switch (blockMode) {
-        case .CFB, .OFB, .CTR:
-            // CFB, OFB, CTR uses encryptBlock to decrypt
-            let decryptGenerator = blockMode.decryptGenerator(iv, cipherOperation: encryptBlock, inputGenerator: AnyGenerator<Array<UInt8>>(blocks.generate()))
-            for processedBlock in AnySequence<Array<UInt8>>({ decryptGenerator }) {
-                out.appendContentsOf(processedBlock)
-            }
-        default:
-            let decryptGenerator = blockMode.decryptGenerator(iv, cipherOperation: decryptBlock, inputGenerator: AnyGenerator<Array<UInt8>>(blocks.generate()))
-            for processedBlock in AnySequence<Array<UInt8>>({ decryptGenerator }) {
-                out.appendContentsOf(processedBlock)
-            }
+        for (idx,chunk) in chunks.enumerate() {
+            out.appendContentsOf(try oneTimeCryptor.update(chunk, isLast: idx == max(0,chunks.count - 1)))
         }
-        
-        return self.padding.remove(out, blockSize: AES.blockSize)
+        return out
     }
     
     private func decryptBlock(block:[UInt8]) -> [UInt8]? {

Sources/CryptoSwift/BlockMode/BlockModeWorker.swift

@@ -0,0 +1,13 @@
+//
+//  BlockModeWorker.swift
+//  CryptoSwift
+//
+//  Created by Marcin Krzyzanowski on 06/05/16.
+//  Copyright © 2016 Marcin Krzyzanowski. All rights reserved.
+//
+
+protocol BlockModeWorker {
+    var cipherOperation: CipherOperationOnBlock { get }
+    mutating func encrypt(plaintext: Array<UInt8>) -> [UInt8]
+    mutating func decrypt(ciphertext: Array<UInt8>) -> [UInt8]
+}

Sources/CryptoSwift/BlockMode/CBC.swift

@@ -8,57 +8,32 @@
 //  Cipher-block chaining (CBC)
 //
 
-struct CBCModeEncryptGenerator: BlockModeGenerator {
+struct CBCModeWorker: BlockModeWorker {
     typealias Element = Array<UInt8>
 
+    let cipherOperation: CipherOperationOnBlock
     private let iv: Element
-    private let inputGenerator: AnyGenerator<Element>
+    private var prev: Element?
 
-    private let cipherOperation: CipherOperationOnBlock
-    private var prevCiphertext: Element?
-
-    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock, inputGenerator: AnyGenerator<Element>) {
+    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock) {
         self.iv = iv
         self.cipherOperation = cipherOperation
-        self.inputGenerator = inputGenerator
     }
 
-    mutating func next() -> Element? {
-        guard let plaintext = inputGenerator.next(),
-              let encrypted = cipherOperation(block: xor(prevCiphertext ?? iv, plaintext))
-        else {
-            return nil
+    mutating func encrypt(plaintext: Array<UInt8>) -> [UInt8] {
+        guard let ciphertext = cipherOperation(block: xor(prev ?? iv, plaintext)) else {
+            return plaintext
         }
-
-        self.prevCiphertext = encrypted
-        return encrypted
+        prev = ciphertext
+        return ciphertext ?? []
     }
-}
-
-struct CBCModeDecryptGenerator: BlockModeGenerator {
-    typealias Element = Array<UInt8>
-
-    private let iv: Element
-    private let inputGenerator: AnyGenerator<Element>
 
-    private let cipherOperation: CipherOperationOnBlock
-    private var prevCiphertext: Element?
-
-    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock, inputGenerator: AnyGenerator<Element>) {
-        self.iv = iv
-        self.cipherOperation = cipherOperation
-        self.inputGenerator = inputGenerator
-    }
-
-    mutating func next() -> Element? {
-        guard let ciphertext = inputGenerator.next(),
-              let decrypted = cipherOperation(block: ciphertext)
-        else {
-            return nil
+    mutating func decrypt(ciphertext: Array<UInt8>) -> [UInt8] {
+        guard let plaintext = cipherOperation(block: ciphertext) else {
+            return ciphertext
         }
-
-        let result = xor(prevCiphertext ?? iv, decrypted)
-        self.prevCiphertext = ciphertext
+        let result = xor(prev ?? iv, plaintext)
+        self.prev = ciphertext
         return result
     }
 }

Sources/CryptoSwift/BlockMode/CFB.swift

@@ -8,57 +8,32 @@
 //  Cipher feedback (CFB)
 //
 
-struct CFBModeEncryptGenerator: BlockModeGenerator {
+struct CFBModeWorker: BlockModeWorker {
     typealias Element = Array<UInt8>
 
+    let cipherOperation: CipherOperationOnBlock
     private let iv: Element
-    private let inputGenerator: AnyGenerator<Element>
+    private var prev: Element?
 
-    private let cipherOperation: CipherOperationOnBlock
-    private var prevCiphertext: Element?
-
-    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock, inputGenerator: AnyGenerator<Array<UInt8>>) {
+    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock) {
         self.iv = iv
         self.cipherOperation = cipherOperation
-        self.inputGenerator = inputGenerator
     }
 
-    mutating func next() -> Element? {
-        guard let plaintext = inputGenerator.next(),
-            let ciphertext = cipherOperation(block: prevCiphertext ?? iv)
-            else {
-                return nil
+    mutating func encrypt(plaintext: Array<UInt8>) -> [UInt8] {
+        guard let ciphertext = cipherOperation(block: prev ?? iv) else {
+            return plaintext
         }
-
-        self.prevCiphertext = xor(plaintext, ciphertext)
-        return self.prevCiphertext
-    }
-}
-
-struct CFBModeDecryptGenerator: BlockModeGenerator {
-    typealias Element = Array<UInt8>
-
-    private let iv: Element
-    private let inputGenerator: AnyGenerator<Element>
-
-    private let cipherOperation: CipherOperationOnBlock
-    private var prevCiphertext: Element?
-
-    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock, inputGenerator: AnyGenerator<Element>) {
-        self.iv = iv
-        self.cipherOperation = cipherOperation
-        self.inputGenerator = inputGenerator
+        prev = xor(plaintext, ciphertext)
+        return prev ?? []
     }
 
-    mutating func next() -> Element? {
-        guard let ciphertext = inputGenerator.next(),
-            let decrypted = cipherOperation(block: self.prevCiphertext ?? iv)
-            else {
-                return nil
+    mutating func decrypt(ciphertext: Array<UInt8>) -> [UInt8] {
+        guard let plaintext = cipherOperation(block: prev ?? iv) else {
+            return ciphertext
         }
-
-        let result = xor(decrypted, ciphertext)
-        self.prevCiphertext = ciphertext
+        let result = xor(plaintext, ciphertext)
+        self.prev = ciphertext
         return result
     }
-}
+}

Sources/CryptoSwift/BlockMode/CTR.swift

@@ -8,64 +8,37 @@
 //  Counter (CTR)
 //
 
-struct CTRModeEncryptGenerator: BlockModeGenerator {
+struct CTRModeWorker: BlockModeWorker {
     typealias Element = Array<UInt8>
 
+    let cipherOperation: CipherOperationOnBlock
     private let iv: Element
-    private let inputGenerator: AnyGenerator<Element>
-
-    private let cipherOperation: CipherOperationOnBlock
     private var counter: UInt = 0
 
-    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock, inputGenerator: AnyGenerator<Array<UInt8>>) {
+    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock) {
         self.iv = iv
         self.cipherOperation = cipherOperation
-        self.inputGenerator = inputGenerator
     }
 
-    mutating func next() -> Element? {
-        guard let plaintext = inputGenerator.next() else {
-            return nil
-        }
-
+    mutating func encrypt(plaintext: Array<UInt8>) -> [UInt8] {
         let nonce = buildNonce(iv, counter: UInt64(counter))
         counter = counter + 1
-        if let encrypted = cipherOperation(block: nonce) {
-            return xor(plaintext, encrypted)
-        }
-
-        return nil
-    }
-}
-
-struct CTRModeDecryptGenerator: BlockModeGenerator {
-    typealias Element = Array<UInt8>
 
-    private let iv: Element
-    private let inputGenerator: AnyGenerator<Element>
-
-    private let cipherOperation: CipherOperationOnBlock
-    private var counter: UInt = 0
+        guard let ciphertext = cipherOperation(block: nonce) else {
+            return plaintext
+        }
 
-    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock, inputGenerator: AnyGenerator<Element>) {
-        self.iv = iv
-        self.cipherOperation = cipherOperation
-        self.inputGenerator = inputGenerator
+        return xor(plaintext, ciphertext)
     }
 
-    mutating func next() -> Element? {
-        guard let ciphertext = inputGenerator.next() else {
-            return nil
-        }
-
+    mutating func decrypt(ciphertext: Array<UInt8>) -> [UInt8] {
         let nonce = buildNonce(iv, counter: UInt64(counter))
         counter = counter + 1
 
-        if let decrypted = cipherOperation(block: nonce) {
-            return xor(decrypted, ciphertext)
+        guard let plaintext = cipherOperation(block: nonce) else {
+            return ciphertext
         }
-
-        return nil
+        return xor(plaintext, ciphertext)
     }
 }
 

Sources/CryptoSwift/BlockMode/CipherBlockMode.swift

@@ -9,84 +9,37 @@
 public enum CipherBlockMode {
     case ECB, CBC, PCBC, CFB, OFB, CTR
 
-    func encryptGenerator(iv: Array<UInt8>?, cipherOperation: CipherOperationOnBlock, inputGenerator: AnyGenerator<Array<UInt8>>) -> AnyGenerator<Array<UInt8>> {
+    func worker(iv: Array<UInt8>?, cipherOperation: CipherOperationOnBlock) -> BlockModeWorker {
         switch (self) {
-        case CBC:
-            return AnyGenerator<Array<UInt8>>(CBCModeEncryptGenerator(iv: iv ?? [], cipherOperation: cipherOperation, inputGenerator: inputGenerator))
-        case CFB:
-            return AnyGenerator<Array<UInt8>>(CFBModeEncryptGenerator(iv: iv ?? [], cipherOperation: cipherOperation, inputGenerator: inputGenerator))
-        case OFB:
-            return AnyGenerator<Array<UInt8>>(OFBModeEncryptGenerator(iv: iv ?? [], cipherOperation: cipherOperation, inputGenerator: inputGenerator))
-        case CTR:
-            return AnyGenerator<Array<UInt8>>(CTRModeEncryptGenerator(iv: iv ?? [], cipherOperation: cipherOperation, inputGenerator: inputGenerator))
-        case PCBC:
-            return AnyGenerator<Array<UInt8>>(PCBCModeEncryptGenerator(iv: iv ?? [], cipherOperation: cipherOperation, inputGenerator: inputGenerator))
         case ECB:
-            return AnyGenerator<Array<UInt8>>(ECBModeEncryptGenerator(iv: iv ?? [], cipherOperation: cipherOperation, inputGenerator: inputGenerator))
-        }
-    }
-
-    func decryptGenerator(iv: Array<UInt8>?, cipherOperation: CipherOperationOnBlock, inputGenerator: AnyGenerator<Array<UInt8>>) -> AnyGenerator<Array<UInt8>> {
-        switch (self) {
+            return ECBModeWorker(iv: iv ?? [], cipherOperation: cipherOperation)
         case CBC:
-            return AnyGenerator<Array<UInt8>>(CBCModeDecryptGenerator(iv: iv ?? [], cipherOperation: cipherOperation, inputGenerator: inputGenerator))
+            return CBCModeWorker(iv: iv ?? [], cipherOperation: cipherOperation)
+        case PCBC:
+            return PCBCModeWorker(iv: iv ?? [], cipherOperation: cipherOperation)
         case CFB:
-            return AnyGenerator<Array<UInt8>>(CFBModeDecryptGenerator(iv: iv ?? [], cipherOperation: cipherOperation, inputGenerator: inputGenerator))
+            return CFBModeWorker(iv: iv ?? [], cipherOperation: cipherOperation)
         case OFB:
-            return AnyGenerator<Array<UInt8>>(OFBModeDecryptGenerator(iv: iv ?? [], cipherOperation: cipherOperation, inputGenerator: inputGenerator))
+            return OFBModeWorker(iv: iv ?? [], cipherOperation: cipherOperation)
         case CTR:
-            return AnyGenerator<Array<UInt8>>(CTRModeDecryptGenerator(iv: iv ?? [], cipherOperation: cipherOperation, inputGenerator: inputGenerator))
-        case PCBC:
-            return AnyGenerator<Array<UInt8>>(PCBCModeDecryptGenerator(iv: iv ?? [], cipherOperation: cipherOperation, inputGenerator: inputGenerator))
-        case ECB:
-            return AnyGenerator<Array<UInt8>>(ECBModeDecryptGenerator(iv: iv ?? [], cipherOperation: cipherOperation, inputGenerator: inputGenerator))
+            return CTRModeWorker(iv: iv ?? [], cipherOperation: cipherOperation)
         }
     }
 
     var options: BlockModeOptions {
         switch (self) {
+        case .ECB:
+            return .PaddingRequired
         case .CBC:
             return [.InitializationVectorRequired, .PaddingRequired]
         case .CFB:
             return .InitializationVectorRequired
         case .CTR:
             return .InitializationVectorRequired
-        case .ECB:
-            return .PaddingRequired
         case .OFB:
             return .InitializationVectorRequired
         case .PCBC:
             return [.InitializationVectorRequired, .PaddingRequired]
         }
     }
-
-    /**
-     Process input blocks with given block cipher mode. With fallback to plain mode.
-
-     - parameter blocks: cipher block size blocks
-     - parameter iv:     IV
-     - parameter cipher: single block encryption closure
-
-     - returns: encrypted bytes
-     */
-//    func encryptBlocks(blocks:[[UInt8]], iv:[UInt8]?, cipherOperation:CipherOperationOnBlock) throws -> [UInt8] {
-//
-//        // if IV is not available, fallback to plain
-//        var finalBlockMode:CipherBlockMode = self
-//        if (iv == nil) {
-//            finalBlockMode = .ECB
-//        }
-//
-//        return try finalBlockMode.mode.encryptBlocks(blocks, iv: iv, cipherOperation: cipherOperation)
-//    }
-//
-//    func decryptBlocks(blocks:[[UInt8]], iv:[UInt8]?, cipherOperation:CipherOperationOnBlock) throws -> [UInt8] {
-//        // if IV is not available, fallback to plain
-//        var finalBlockMode:CipherBlockMode = self
-//        if (iv == nil) {
-//            finalBlockMode = .ECB
-//        }
-//
-//        return try finalBlockMode.mode.decryptBlocks(blocks, iv: iv, cipherOperation: cipherOperation)
-//    }
 }

Sources/CryptoSwift/BlockMode/ECB.swift

@@ -8,6 +8,27 @@
 //  Electronic codebook (ECB)
 //
 
+struct ECBModeWorker: BlockModeWorker {
+    typealias Element = Array<UInt8>
+    let cipherOperation: CipherOperationOnBlock
+
+    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock) {
+        self.cipherOperation = cipherOperation
+    }
+
+    mutating func encrypt(plaintext: Array<UInt8>) -> [UInt8] {
+        guard let ciphertext = cipherOperation(block: plaintext) else {
+            return plaintext
+        }
+        return ciphertext
+    }
+
+    mutating func decrypt(ciphertext: Array<UInt8>) -> [UInt8] {
+        return encrypt(ciphertext)
+    }
+}
+
+
 struct ECBModeEncryptGenerator: BlockModeGenerator {
     typealias Element = Array<UInt8>
 

Sources/CryptoSwift/BlockMode/OFB.swift

@@ -8,57 +8,32 @@
 // Output Feedback (OFB)
 //
 
-struct OFBModeEncryptGenerator: BlockModeGenerator {
+struct OFBModeWorker: BlockModeWorker {
     typealias Element = Array<UInt8>
 
+    let cipherOperation: CipherOperationOnBlock
     private let iv: Element
-    private let inputGenerator: AnyGenerator<Element>
+    private var prev: Element?
 
-    private let cipherOperation: CipherOperationOnBlock
-    private var prevCiphertext: Element?
-
-    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock, inputGenerator: AnyGenerator<Array<UInt8>>) {
+    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock) {
         self.iv = iv
         self.cipherOperation = cipherOperation
-        self.inputGenerator = inputGenerator
     }
 
-    mutating func next() -> Element? {
-        guard let plaintext = inputGenerator.next(),
-            let ciphertext = cipherOperation(block: prevCiphertext ?? iv)
-            else {
-                return nil
+    mutating func encrypt(plaintext: Array<UInt8>) -> [UInt8] {
+        guard let ciphertext = cipherOperation(block: prev ?? iv) else {
+            return plaintext
         }
-
-        self.prevCiphertext = ciphertext
+        prev = ciphertext
         return xor(plaintext, ciphertext)
     }
-}
-
-struct OFBModeDecryptGenerator: BlockModeGenerator {
-    typealias Element = Array<UInt8>
-
-    private let iv: Element
-    private let inputGenerator: AnyGenerator<Element>
-
-    private let cipherOperation: CipherOperationOnBlock
-    private var prevCiphertext: Element?
-
-    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock, inputGenerator: AnyGenerator<Element>) {
-        self.iv = iv
-        self.cipherOperation = cipherOperation
-        self.inputGenerator = inputGenerator
-    }
 
-    mutating func next() -> Element? {
-        guard let ciphertext = inputGenerator.next(),
-            let decrypted = cipherOperation(block: self.prevCiphertext ?? iv)
-            else {
-                return nil
+    mutating func decrypt(ciphertext: Array<UInt8>) -> [UInt8] {
+        guard let decrypted = cipherOperation(block: prev ?? iv) else {
+            return ciphertext
         }
-
         let plaintext = xor(decrypted, ciphertext)
-        self.prevCiphertext = decrypted
+        self.prev = decrypted
         return plaintext
     }
 }

Sources/CryptoSwift/BlockMode/PCBC.swift

@@ -8,57 +8,32 @@
 //  Propagating Cipher Block Chaining (PCBC)
 //
 
-struct PCBCModeEncryptGenerator: BlockModeGenerator {
+struct PCBCModeWorker: BlockModeWorker {
     typealias Element = Array<UInt8>
 
+    let cipherOperation: CipherOperationOnBlock
     private let iv: Element
-    private let inputGenerator: AnyGenerator<Element>
+    private var prev: Element?
 
-    private let cipherOperation: CipherOperationOnBlock
-    private var prevCiphertext: Element?
-
-    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock, inputGenerator: AnyGenerator<Array<UInt8>>) {
+    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock) {
         self.iv = iv
         self.cipherOperation = cipherOperation
-        self.inputGenerator = inputGenerator
     }
 
-    mutating func next() -> Element? {
-        guard let plaintext = inputGenerator.next(),
-              let encrypted = cipherOperation(block: xor(prevCiphertext ?? iv, plaintext))
-        else {
-            return nil
+    mutating func encrypt(plaintext: Array<UInt8>) -> [UInt8] {
+        guard let ciphertext = cipherOperation(block: xor(prev ?? iv, plaintext)) else {
+            return plaintext
         }
-
-        self.prevCiphertext = xor(plaintext, encrypted)
-        return encrypted
-    }
-}
-
-struct PCBCModeDecryptGenerator: BlockModeGenerator {
-    typealias Element = Array<UInt8>
-
-    private let iv: Element
-    private let inputGenerator: AnyGenerator<Element>
-
-    private let cipherOperation: CipherOperationOnBlock
-    private var prevCiphertext: Element?
-
-    init(iv: Array<UInt8>, cipherOperation: CipherOperationOnBlock, inputGenerator: AnyGenerator<Element>) {
-        self.iv = iv
-        self.cipherOperation = cipherOperation
-        self.inputGenerator = inputGenerator
+        prev = xor(plaintext, ciphertext)
+        return ciphertext ?? []
     }
 
-    mutating func next() -> Element? {
-        guard let ciphertext = inputGenerator.next(),
-              let decrypted = cipherOperation(block: ciphertext)
-        else {
-            return nil
+    mutating func decrypt(ciphertext: Array<UInt8>) -> [UInt8] {
+        guard let plaintext = cipherOperation(block: ciphertext) else {
+            return ciphertext
         }
-
-        let plaintext = xor(prevCiphertext ?? iv, decrypted)
-        self.prevCiphertext = xor(plaintext, ciphertext)
-        return plaintext
+        let result = xor(prev ?? iv, plaintext)
+        self.prev = xor(plaintext, ciphertext)
+        return result
     }
 }