Validate AES key length

CHANGELOG

@@ -2,6 +2,7 @@
 - API: BlockMode is no longer an enum. Please migrate to eg. CBC() etc...
 - Adds AES-GCM support.
 - Adds CRC32c support.
+- Improves AES variant validation.
 - Fixes empty password in PBKDF2.
 
 0.9.0

Sources/CryptoSwift/AES.swift

@@ -19,6 +19,8 @@
 /// The Advanced Encryption Standard (AES)
 public final class AES: BlockCipher {
     public enum Error: Swift.Error {
+        /// Invalid key
+        case invalidKeySize
         /// Data padding is required
         case dataPaddingRequired
         /// Invalid Data
@@ -41,25 +43,15 @@ public final class AES: BlockCipher {
         }
     }
 
-    private lazy var variantNr: Int = self.variant.Nr
-    private lazy var variantNb: Int = self.variant.Nb
-    private lazy var variantNk: Int = self.variant.Nk
+    private let variantNr: Int
+    private let variantNb: Int
+    private let variantNk: Int
 
     public static let blockSize: Int = 16 // 128 /8
     public let keySize: Int
 
-    public var variant: Variant {
-        switch keySize * 8 {
-        case 128:
-            return .aes128
-        case 192:
-            return .aes192
-        case 256:
-            return .aes256
-        default:
-            preconditionFailure("Unknown AES variant for given key.")
-        }
-    }
+    /// AES Variant
+    public let variant: Variant
 
     // Parameters
     let key: Key
@@ -124,7 +116,23 @@ public final class AES: BlockCipher {
         self.key = Key(bytes: key)
         self.blockMode = blockMode
         self.padding = padding
-        keySize = self.key.count
+        self.keySize = self.key.count
+
+        // Validate key size
+        switch keySize * 8 {
+        case 128:
+            variant = .aes128
+        case 192:
+            variant = .aes192
+        case 256:
+            variant = .aes256
+        default:
+            throw Error.invalidKeySize
+        }
+
+        variantNb = variant.Nb
+        variantNk = variant.Nk
+        variantNr = variant.Nr
     }
 
     internal func encrypt(block: ArraySlice<UInt8>) -> Array<UInt8>? {
@@ -356,7 +364,7 @@ private extension AES {
 
     private func expandKey(_ key: Key, variant _: Variant) -> Array<Array<UInt32>> {
         func convertExpandedKey(_ expanded: Array<UInt8>) -> Array<Array<UInt32>> {
-            return expanded.batched(by: 4).map({ UInt32(bytes: $0.reversed()) }).batched(by: 4).map({ Array($0) })
+            return expanded.batched(by: 4).map({ UInt32(bytes: $0.reversed()) }).batched(by: 4).map { Array($0) }
         }
 
         /*