From 1e04890d733667c433f0ac818ffd18754e0b7752 Mon Sep 17 00:00:00 2001
From: Baojun Wang <wangbj@gmail.com>
Date: Thu, 1 Jun 2017 10:25:09 -0700
Subject: [PATCH] Add ccm decrypt

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

diff --git a/Crypto/Cipher/AES/Primitive.hs b/Crypto/Cipher/AES/Primitive.hs
index f139c17..6c684b1 100644
--- a/Crypto/Cipher/AES/Primitive.hs
+++ b/Crypto/Cipher/AES/Primitive.hs
@@ -535,12 +535,7 @@ ccmEncrypt ctx ccm input = unsafeDoIO $ (withCCMKeyAndCopySt ctx ccm cbcmacAndIv
 -- needs to happen after AAD appending, or after initialization if no AAD data.
 {-# NOINLINE ccmDecrypt #-}
 ccmDecrypt :: ByteArray ba => AES -> AESCCM -> ba -> (ba, AESCCM)
-ccmDecrypt ctx ccm input = unsafeDoIO $ withCCMKeyAndCopySt ctx ccm doDec
-  where len = B.length input
-        doDec ccmStPtr aesPtr =
-            B.alloc len $ \o ->
-            withByteArray input $ \i ->
-            c_aes_ccm_decrypt (castPtr o) ccmStPtr aesPtr i (fromIntegral len)
+ccmDecrypt = ccmEncrypt
 
 -- | Generate the Tag from CCM context
 {-# NOINLINE ccmFinish #-}
diff --git a/tests/BlockCipher.hs b/tests/BlockCipher.hs
index 3243f70..988c053 100644
--- a/tests/BlockCipher.hs
+++ b/tests/BlockCipher.hs
@@ -16,7 +16,7 @@ import           Imports
 import           Data.Maybe
 import           Crypto.Error
 import           Crypto.Cipher.Types
-import           Data.ByteArray as B hiding (pack, null)
+import           Data.ByteArray as B hiding (pack, null, length)
 import qualified Data.ByteString as B hiding (all)
 
 ------------------------------------------------------------------------
@@ -389,7 +389,7 @@ testBlockCipherModes cipher =
 testBlockCipherAEAD :: BlockCipher a => a -> [TestTree]
 testBlockCipherAEAD cipher =
     [ testProperty "OCB" (aeadProp AEAD_OCB)
-    , testProperty "CCM" (aeadProp AEAD_CCM)
+    , testProperty "CCM" (aeadProp (AEAD_CCM 0 CCM_M16 CCM_L2))
     , testProperty "EAX" (aeadProp AEAD_EAX)
     , testProperty "CWC" (aeadProp AEAD_CWC)
     , testProperty "GCM" (aeadProp AEAD_GCM)
@@ -398,7 +398,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 testIV of
                 CryptoPassed iniAead ->
                     let aead           = aeadAppendHeader iniAead aad
                         (eText, aeadE) = aeadEncrypt aead plaintext
@@ -409,6 +409,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
 
 withCtx :: Cipher c => Key c -> (c -> a) -> a
 withCtx (Key key) f =