From f5efdee75b44051bc94cb4714f310609e021093a Mon Sep 17 00:00:00 2001
From: Tim McIver <tim@timmciver.com>
Date: Tue, 6 Dec 2016 22:06:57 -0500
Subject: [PATCH] Add key length validation to several AES Ciphers.

---
 Crypto/Cipher/AES.hs | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/Crypto/Cipher/AES.hs b/Crypto/Cipher/AES.hs
index 7968049..b00a026 100644
--- a/Crypto/Cipher/AES.hs
+++ b/Crypto/Cipher/AES.hs
@@ -18,6 +18,8 @@ import Crypto.Cipher.Types.Block
 import Crypto.Cipher.AES.Primitive
 import Crypto.Internal.Imports
 
+import Data.ByteArray as BA
+
 -- | AES with 128 bit key
 newtype AES128 = AES128 AES
     deriving (NFData)
@@ -33,17 +35,27 @@ newtype AES256 = AES256 AES
 instance Cipher AES128 where
     cipherName    _ = "AES128"
     cipherKeySize _ = KeySizeFixed 16
-    cipherInit k    = AES128 `fmap` initAES k
+    cipherInit k    = AES128 <$> (initAES =<< validateKeySize (undefined :: AES128) k)
 
 instance Cipher AES192 where
     cipherName    _ = "AES192"
     cipherKeySize _ = KeySizeFixed 24
-    cipherInit k    = AES192 `fmap` initAES k
+    cipherInit k    = AES192 <$> (initAES =<< validateKeySize (undefined :: AES192) k)
 
 instance Cipher AES256 where
     cipherName    _ = "AES256"
     cipherKeySize _ = KeySizeFixed 32
-    cipherInit k    = AES256 `fmap` initAES k
+    cipherInit k    = AES256 <$> (initAES =<< validateKeySize (undefined :: AES256) k)
+
+validateKeySize :: (ByteArrayAccess key, Cipher cipher) => cipher -> key -> CryptoFailable key
+validateKeySize c k = if validKeyLength
+                      then CryptoPassed k
+                      else CryptoFailed CryptoError_KeySizeInvalid
+  where keyLength = BA.length k
+        validKeyLength = case cipherKeySize c of
+          KeySizeRange low high -> keyLength >= low && keyLength <= high
+          KeySizeEnum lengths -> keyLength `elem` lengths
+          KeySizeFixed s -> keyLength == s
 
 #define INSTANCE_BLOCKCIPHER(CSTR) \
 instance BlockCipher CSTR where \