From b28b25abe7269a4fb575b6b5934417a731fedd76 Mon Sep 17 00:00:00 2001
From: Vincent Hanquez <vincent@snarc.org>
Date: Sat, 18 Apr 2015 07:58:52 +0100
Subject: [PATCH] [AES] add gcm / ocb support to base AES type

---
 Crypto/Cipher/AES.hs           | 16 ----------------
 Crypto/Cipher/AES/Primitive.hs | 27 ++++++++++++++++++++++-----
 2 files changed, 22 insertions(+), 21 deletions(-)

diff --git a/Crypto/Cipher/AES.hs b/Crypto/Cipher/AES.hs
index a6a7f08..00313df 100644
--- a/Crypto/Cipher/AES.hs
+++ b/Crypto/Cipher/AES.hs
@@ -40,22 +40,6 @@ instance Cipher AES256 where
     cipherKeySize _ = KeySizeFixed 32
     cipherInit k    = AES256 `fmap` initAES k
 
-gcmMode :: AES -> AEADModeImpl AESGCM
-gcmMode aes = AEADModeImpl
-    { aeadImplAppendHeader = gcmAppendAAD
-    , aeadImplEncrypt      = gcmAppendEncrypt aes
-    , aeadImplDecrypt      = gcmAppendDecrypt aes
-    , aeadImplFinalize     = gcmFinish aes
-    }
-
-ocbMode :: AES -> AEADModeImpl AESOCB
-ocbMode aes = AEADModeImpl
-    { aeadImplAppendHeader = ocbAppendAAD aes
-    , aeadImplEncrypt      = ocbAppendEncrypt aes
-    , aeadImplDecrypt      = ocbAppendDecrypt aes
-    , aeadImplFinalize     = ocbFinish aes
-    }
-
 #define INSTANCE_BLOCKCIPHER(CSTR) \
 instance BlockCipher CSTR where \
     { blockSize _ = 16 \
diff --git a/Crypto/Cipher/AES/Primitive.hs b/Crypto/Cipher/AES/Primitive.hs
index 85dad93..cbf95d4 100644
--- a/Crypto/Cipher/AES/Primitive.hs
+++ b/Crypto/Cipher/AES/Primitive.hs
@@ -44,6 +44,7 @@ module Crypto.Cipher.AES.Primitive
     , decryptOCB
 
     -- * incremental GCM
+    , gcmMode
     , gcmInit
     , gcmAppendAAD
     , gcmAppendEncrypt
@@ -51,6 +52,7 @@ module Crypto.Cipher.AES.Primitive
     , gcmFinish
 
     -- * incremental OCB
+    , ocbMode
     , ocbInit
     , ocbAppendAAD
     , ocbAppendEncrypt
@@ -86,15 +88,30 @@ instance BlockCipher AES where
     cbcEncrypt = encryptCBC
     cbcDecrypt = decryptCBC
     ctrCombine = encryptCTR
-    {-
-    aeadInit AEAD_GCM aes iv = Just $ AEAD aes $ AEADState $ gcmInit aes iv
-    aeadInit AEAD_OCB aes iv = Just $ AEAD aes $ AEADState $ ocbInit aes iv
-    aeadInit _        _    _ = Nothing
-    -}
+    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 _        _   _  = CryptoFailed CryptoError_AEADModeNotSupported
 instance BlockCipher128 AES where 
     xtsEncrypt = encryptXTS
     xtsDecrypt = decryptXTS
 
+gcmMode :: AES -> AEADModeImpl AESGCM
+gcmMode aes = AEADModeImpl
+    { aeadImplAppendHeader = gcmAppendAAD
+    , aeadImplEncrypt      = gcmAppendEncrypt aes
+    , aeadImplDecrypt      = gcmAppendDecrypt aes
+    , aeadImplFinalize     = gcmFinish aes
+    }
+
+ocbMode :: AES -> AEADModeImpl AESOCB
+ocbMode aes = AEADModeImpl
+    { aeadImplAppendHeader = ocbAppendAAD aes
+    , aeadImplEncrypt      = ocbAppendEncrypt aes
+    , aeadImplDecrypt      = ocbAppendDecrypt aes
+    , aeadImplFinalize     = ocbFinish aes
+    }
+
+
 -- | AES Context (pre-processed key)
 newtype AES = AES SecureBytes