From f6d9fb0cf1980e71a5787e57b1fa02591d05fce1 Mon Sep 17 00:00:00 2001
From: Baojun Wang <wangbj@gmail.com>
Date: Mon, 22 Jan 2018 12:24:22 -0800
Subject: [PATCH] aeadInit (ccm) returns CryptoError_IvSizeInvalid when iv size
 is wrong

---
 Crypto/Cipher/AES.hs |  9 ++++++++-
 tests/BlockCipher.hs | 12 ++++++------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/Crypto/Cipher/AES.hs b/Crypto/Cipher/AES.hs
index 5ce10c4..5e2c76c 100644
--- a/Crypto/Cipher/AES.hs
+++ b/Crypto/Cipher/AES.hs
@@ -48,6 +48,13 @@ instance Cipher AES256 where
     cipherKeySize _ = KeySizeFixed 32
     cipherInit k    = AES256 <$> (initAES =<< validateKeySize (undefined :: AES256) k)
 
+aeadInitCcm :: ByteArrayAccess iv => Int -> CCM_M -> CCM_L -> AES -> iv -> CryptoFailable (AEAD cihper)
+aeadInitCcm n m l aes iv = if BA.length iv /= 15 - ln then CryptoFailed CryptoError_IvSizeInvalid else CryptoPassed $ AEAD (ccmMode aes) (ccmInit aes iv n m l)
+  where
+    ln = case l of
+      CCM_L2 -> 2
+      CCM_L3 -> 3
+      CCM_L4 -> 4
 
 #define INSTANCE_BLOCKCIPHER(CSTR) \
 instance BlockCipher CSTR where \
@@ -59,7 +66,7 @@ instance BlockCipher CSTR where \
     ; ctrCombine (CSTR aes) (IV iv) = encryptCTR aes (IV iv) \
     ; aeadInit AEAD_GCM (CSTR aes) iv = CryptoPassed $ AEAD (gcmMode aes) (gcmInit aes iv) \
     ; aeadInit AEAD_OCB (CSTR aes) iv = CryptoPassed $ AEAD (ocbMode aes) (ocbInit aes iv) \
-    ; aeadInit (AEAD_CCM n m l) (CSTR aes) iv = CryptoPassed $ AEAD (ccmMode aes) (ccmInit aes iv n m l) \
+    ; aeadInit (AEAD_CCM n m l) (CSTR aes) iv = aeadInitCcm n m l aes iv \
     ; aeadInit _        _          _  = CryptoFailed CryptoError_AEADModeNotSupported \
     }; \
 instance BlockCipher128 CSTR where \
diff --git a/tests/BlockCipher.hs b/tests/BlockCipher.hs
index 44b571b..6e8f8af 100644
--- a/tests/BlockCipher.hs
+++ b/tests/BlockCipher.hs
@@ -17,7 +17,7 @@ import           Data.Maybe
 import           Crypto.Error
 import           Crypto.Cipher.Types
 import           Data.ByteArray as B hiding (pack, null, length)
-import qualified Data.ByteString as B hiding (all)
+import qualified Data.ByteString as B hiding (all, take, replicate)
 
 ------------------------------------------------------------------------
 -- KAT
@@ -402,7 +402,7 @@ testBlockCipherAEAD cipher =
         toTests :: BlockCipher a => a -> (AEADMode -> AEADUnit a -> Bool)
         toTests _ = testProperty_AEAD
         testProperty_AEAD mode (AEADUnit key testIV (unPlaintext -> aad) (unPlaintext -> plaintext)) = withCtx key $ \ctx ->
-            case aeadInit mode' ctx testIV of
+            case aeadInit mode' ctx iv' of
                 CryptoPassed iniAead ->
                     let aead           = aeadAppendHeader iniAead aad
                         (eText, aeadE) = aeadEncrypt aead plaintext
@@ -413,10 +413,10 @@ testBlockCipherAEAD cipher =
                 CryptoFailed err
                     | err == CryptoError_AEADModeNotSupported -> True
                     | otherwise                               -> error ("testProperty_AEAD: " ++ show err)
-            where mode' = updateCcmInputSize mode (B.length plaintext)
-                  updateCcmInputSize aeadmode k = case aeadmode of
-                    AEAD_CCM _ m l -> AEAD_CCM k m l
-                    aeadOther      -> aeadOther
+            where (mode', iv') = updateCcmInputSize mode (B.length plaintext) testIV
+                  updateCcmInputSize aeadmode k iv = case aeadmode of
+                    AEAD_CCM _ m l -> (AEAD_CCM k m l, B.take 13 (iv <> (B.replicate 15 0)))
+                    aeadOther      -> (aeadOther, iv)
 
 withCtx :: Cipher c => Key c -> (c -> a) -> a
 withCtx (Key key) f =