Fixing blockSize errors (to allow CFB to process byte per byte)

Sources/CryptoSwift/AES.Cryptors.swift

@@ -17,19 +17,21 @@
 
 extension AES: Cryptors {
   public func makeEncryptor() throws -> Cryptor & Updatable {
-    let worker = try blockMode.worker(blockSize: AES.blockSize, cipherOperation: encrypt, encryptionOperation: encrypt)
+    let blockSize = blockMode.customBlockSize ?? AES.blockSize
+    let worker = try blockMode.worker(blockSize: blockSize, cipherOperation: encrypt, encryptionOperation: encrypt)
     if worker is StreamModeWorker {
-      return try StreamEncryptor(blockSize: AES.blockSize, padding: padding, worker)
+      return try StreamEncryptor(blockSize: blockSize, padding: padding, worker)
     }
-    return try BlockEncryptor(blockSize: AES.blockSize, padding: padding, worker)
+    return try BlockEncryptor(blockSize: blockSize, padding: padding, worker)
   }
 
   public func makeDecryptor() throws -> Cryptor & Updatable {
+    let blockSize = blockMode.customBlockSize ?? AES.blockSize
     let cipherOperation: CipherOperationOnBlock = blockMode.options.contains(.useEncryptToDecrypt) == true ? encrypt : decrypt
-    let worker = try blockMode.worker(blockSize: AES.blockSize, cipherOperation: cipherOperation, encryptionOperation: encrypt)
+    let worker = try blockMode.worker(blockSize: blockSize, cipherOperation: cipherOperation, encryptionOperation: encrypt)
     if worker is StreamModeWorker {
-      return try StreamDecryptor(blockSize: AES.blockSize, padding: padding, worker)
+      return try StreamDecryptor(blockSize: blockSize, padding: padding, worker)
     }
-    return try BlockDecryptor(blockSize: AES.blockSize, padding: padding, worker)
+    return try BlockDecryptor(blockSize: blockSize, padding: padding, worker)
   }
 }

Sources/CryptoSwift/AES.swift

@@ -496,7 +496,8 @@ private extension AES {
 
 extension AES: Cipher {
   public func encrypt(_ bytes: ArraySlice<UInt8>) throws -> Array<UInt8> {
-    let chunks = bytes.batched(by: AES.blockSize)
+    let blockSize = self.blockMode.customBlockSize ?? AES.blockSize
+    let chunks = bytes.batched(by: blockSize)
 
     var oneTimeCryptor = try makeEncryptor()
     var out = Array<UInt8>(reserveCapacity: bytes.count)

Sources/CryptoSwift/BlockMode/BlockMode.swift

@@ -19,6 +19,8 @@ public protocol BlockMode {
   var options: BlockModeOption { get }
   //TODO: doesn't have to be public
   func worker(blockSize: Int, cipherOperation: @escaping CipherOperationOnBlock, encryptionOperation: @escaping CipherOperationOnBlock) throws -> CipherModeWorker
+
+  var customBlockSize: Int? { get }
 }
 
 typealias StreamMode = BlockMode

Sources/CryptoSwift/BlockMode/CBC.swift

@@ -24,6 +24,7 @@ public struct CBC: BlockMode {
 
   public let options: BlockModeOption = [.initializationVectorRequired, .paddingRequired]
   private let iv: Array<UInt8>
+  public let customBlockSize: Int? = nil
 
   public init(iv: Array<UInt8>) {
     self.iv = iv

Sources/CryptoSwift/BlockMode/CCM.swift

@@ -38,6 +38,7 @@ public struct CCM: StreamMode {
   private let additionalAuthenticatedData: Array<UInt8>?
   private let tagLength: Int
   private let messageLength: Int // total message length. need to know in advance
+  public let customBlockSize: Int? = nil
 
   // `authenticationTag` nil for encryption, known tag for decryption
   /// For encryption, the value is set at the end of the encryption.

Sources/CryptoSwift/BlockMode/CFB.swift

@@ -23,21 +23,23 @@ public struct CFB: BlockMode {
   }
     
   public enum SegmentSize: Int {
-    case cfb8 = 8 // Encrypt byte per byte
-    case cfb128 = 128 // Encrypt 16 bytes per 16 bytes (default)
+    case cfb8 = 1 // Encrypt byte per byte
+    case cfb128 = 16 // Encrypt 16 bytes per 16 bytes (default)
   }
 
   public let options: BlockModeOption = [.initializationVectorRequired, .useEncryptToDecrypt]
   private let iv: Array<UInt8>
   private let segmentSize: SegmentSize
+  public let customBlockSize: Int?
 
   public init(iv: Array<UInt8>, segmentSize: SegmentSize = .cfb128) {
     self.iv = iv
     self.segmentSize = segmentSize
+    self.customBlockSize = segmentSize.rawValue
   }
 
   public func worker(blockSize: Int, cipherOperation: @escaping CipherOperationOnBlock, encryptionOperation: @escaping CipherOperationOnBlock) throws -> CipherModeWorker {
-    if self.iv.count != blockSize {
+    if self.iv.count != AES.blockSize {
       throw Error.invalidInitializationVector
     }
 
@@ -71,13 +73,12 @@ struct CFBModeWorker: BlockModeWorker {
     }
     // CFB8
     else if segmentSize == .cfb8 {
-      for i in 0 ..< plaintext.count {
-        guard let ciphertext = cipherOperation(prev ?? iv) else {
-          return Array(plaintext)
-        }
-        self.prev = Array((prev ?? iv).dropFirst()) + [plaintext[i] ^ ciphertext[0]]
+      guard let ciphertext = cipherOperation(prev ?? iv) else {
+        return Array(plaintext)
       }
-      return Array(self.prev ?? [])
+      let result = [plaintext[0] ^ ciphertext[0]]
+      self.prev = (prev ?? iv).dropFirst() + result
+      return result
     }
     return Array(plaintext) // Unsupported segment size
   }
@@ -94,15 +95,11 @@ struct CFBModeWorker: BlockModeWorker {
     }
     // CFB8
     else if segmentSize == .cfb8 {
-      var result: Array<UInt8> = []
-      for i in 0 ..< ciphertext.count {
-        guard let plaintext = cipherOperation(prev ?? iv) else {
-          return Array(ciphertext)
-        }
-        self.prev = Array((prev ?? iv).dropFirst()) + [ciphertext[i]]
-        result.append(ciphertext[i] ^ plaintext[0])
+      guard let plaintext = cipherOperation(prev ?? iv) else {
+        return Array(ciphertext)
       }
-      return result
+      self.prev = (prev ?? iv).dropFirst() + [ciphertext[0]]
+      return [ciphertext[0] ^ plaintext[0]]
     }
     return Array(ciphertext) // Unsupported segment size
   }

Sources/CryptoSwift/BlockMode/CTR.swift

@@ -24,6 +24,7 @@ public struct CTR: StreamMode {
   public let options: BlockModeOption = [.initializationVectorRequired, .useEncryptToDecrypt]
   private let iv: Array<UInt8>
   private let counter: Int
+  public let customBlockSize: Int? = nil
 
   public init(iv: Array<UInt8>, counter: Int = 0) {
     self.iv = iv

Sources/CryptoSwift/BlockMode/ECB.swift

@@ -18,6 +18,7 @@
 
 public struct ECB: BlockMode {
   public let options: BlockModeOption = .paddingRequired
+  public let customBlockSize: Int? = nil
 
   public init() {
   }

Sources/CryptoSwift/BlockMode/GCM.swift

@@ -38,6 +38,7 @@ public final class GCM: BlockMode {
   private let iv: Array<UInt8>
   private let additionalAuthenticatedData: Array<UInt8>?
   private let mode: Mode
+  public let customBlockSize: Int? = nil
 
   /// Length of authentication tag, in bytes.
   /// For encryption, the value is given as init parameter.

Sources/CryptoSwift/BlockMode/OCB.swift

@@ -36,6 +36,7 @@ public final class OCB: BlockMode {
   private let N: Array<UInt8>
   private let additionalAuthenticatedData: Array<UInt8>?
   private let mode: Mode
+  public let customBlockSize: Int? = nil
 
   /// Length of authentication tag, in bytes.
   /// For encryption, the value is given as init parameter.

Sources/CryptoSwift/BlockMode/OFB.swift

@@ -24,6 +24,7 @@ public struct OFB: BlockMode {
 
   public let options: BlockModeOption = [.initializationVectorRequired, .useEncryptToDecrypt]
   private let iv: Array<UInt8>
+  public let customBlockSize: Int? = nil
 
   public init(iv: Array<UInt8>) {
     self.iv = iv

Sources/CryptoSwift/BlockMode/PCBC.swift

@@ -24,6 +24,7 @@ public struct PCBC: BlockMode {
 
   public let options: BlockModeOption = [.initializationVectorRequired, .paddingRequired]
   private let iv: Array<UInt8>
+  public let customBlockSize: Int? = nil
 
   public init(iv: Array<UInt8>) {
     self.iv = iv