From 1e57f41e1d4e47c654412f07f69ee946bcabfcf9 Mon Sep 17 00:00:00 2001
From: Baojun Wang <wangbj@gmail.com>
Date: Thu, 1 Feb 2018 11:16:13 -0800
Subject: [PATCH] check AESCCM IV length in ccmInit instead of aeadInit

---
 Crypto/Cipher/AES.hs           | 10 +---------
 Crypto/Cipher/AES/Primitive.hs | 15 ++++++++-------
 2 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/Crypto/Cipher/AES.hs b/Crypto/Cipher/AES.hs
index 5e2c76c..61b81cf 100644
--- a/Crypto/Cipher/AES.hs
+++ b/Crypto/Cipher/AES.hs
@@ -48,14 +48,6 @@ 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 \
     { blockSize _ = 16 \
@@ -66,7 +58,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 = aeadInitCcm n m l aes iv \
+    ; aeadInit (AEAD_CCM n m l) (CSTR aes) iv = AEAD (ccmMode aes) <$> ccmInit aes iv n m l \
     ; aeadInit _        _          _  = CryptoFailed CryptoError_AEADModeNotSupported \
     }; \
 instance BlockCipher128 CSTR where \
diff --git a/Crypto/Cipher/AES/Primitive.hs b/Crypto/Cipher/AES/Primitive.hs
index 2e4f2ac..5169366 100644
--- a/Crypto/Cipher/AES/Primitive.hs
+++ b/Crypto/Cipher/AES/Primitive.hs
@@ -77,7 +77,7 @@ instance BlockCipher AES where
     ctrCombine = encryptCTR
     aeadInit AEAD_GCM aes iv = CryptoPassed $ AEAD (gcmMode aes) (gcmInit aes iv)
     aeadInit AEAD_OCB aes iv = CryptoPassed $ AEAD (ocbMode aes) (ocbInit aes iv)
-    aeadInit (AEAD_CCM n m l) aes iv = CryptoPassed $ AEAD (ccmMode aes) (ccmInit aes iv n m l)
+    aeadInit (AEAD_CCM n m l) aes iv = AEAD (ccmMode aes) <$> ccmInit aes iv n m l
     aeadInit _        _   _  = CryptoFailed CryptoError_AEADModeNotSupported
 instance BlockCipher128 AES where
     xtsEncrypt = encryptXTS
@@ -492,12 +492,13 @@ ccmGetL l = case l of
 
 -- | initialize a ccm context
 {-# NOINLINE ccmInit #-}
-ccmInit :: ByteArrayAccess iv => AES -> iv -> Int -> CCM_M -> CCM_L -> AESCCM
-ccmInit ctx iv n m l = unsafeDoIO $ do
-    sm <- B.alloc sizeCCM $ \ccmStPtr ->
-            withKeyAndIV ctx iv $ \k v ->
-            c_aes_ccm_init (castPtr ccmStPtr) k v (fromIntegral $ B.length iv) (fromIntegral n) (fromIntegral (ccmGetM m)) (fromIntegral (ccmGetL l))
-    return $ AESCCM sm
+ccmInit :: ByteArrayAccess iv => AES -> iv -> Int -> CCM_M -> CCM_L -> CryptoFailable AESCCM
+ccmInit ctx iv n m l = if 15 - ccmGetL l /= B.length iv then CryptoFailed CryptoError_IvSizeInvalid
+  else unsafeDoIO $ do
+      sm <- B.alloc sizeCCM $ \ccmStPtr ->
+        withKeyAndIV ctx iv $ \k v ->
+        c_aes_ccm_init (castPtr ccmStPtr) k v (fromIntegral $ B.length iv) (fromIntegral n) (fromIntegral (ccmGetM m)) (fromIntegral (ccmGetL l))
+      return $ CryptoPassed (AESCCM sm)
 
 -- | append data which is only going to be authenticated to the CCM context.
 --