AESSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module BlockCipher.AESSpec (spec) where

import BlockCipher
import qualified Crypto.Cipher.AES as AES
import Crypto.Cipher.Types
import Crypto.Error
import Crypto.Hash (Digest, SHA256, hash)
import qualified Data.ByteArray as BA
import qualified Data.ByteString as B
import Data.Maybe
import Imports

import qualified BlockCipher.AES.CBC as KATCBC
import qualified BlockCipher.AES.CCM as KATCCM
import qualified BlockCipher.AES.CTR as KATCTR
import qualified BlockCipher.AES.ECB as KATECB
import qualified BlockCipher.AES.GCM as KATGCM
import qualified BlockCipher.AES.GCMLong as KATGCMLong
import qualified BlockCipher.AES.OCB3 as KATOCB3
import qualified BlockCipher.AES.XTS as KATXTS

{-
instance Show AES.AES where
    show _ = "AES"
instance Arbitrary AES.AESIV where
    arbitrary = AES.aesIV_ . B.pack <$> replicateM 16 arbitrary
instance Arbitrary AES.AES where
    arbitrary = AES.initAES . B.pack <$> replicateM 16 arbitrary
-}

toKatECB (k, p, c) = KAT_ECB{ecbKey = k, ecbPlaintext = p, ecbCiphertext = c}
toKatCBC (k, iv, p, c) = KAT_CBC{cbcKey = k, cbcIV = iv, cbcPlaintext = p, cbcCiphertext = c}
toKatCTR (k, iv, p, c) = KAT_CTR{ctrKey = k, ctrIV = iv, ctrPlaintext = p, ctrCiphertext = c}
toKatXTS (k1, k2, iv, p, _, c) =
    KAT_XTS
        { xtsKey1 = k1
        , xtsKey2 = k2
        , xtsIV = iv
        , xtsPlaintext = p
        , xtsCiphertext = c
        }
toKatAEAD mode (k, iv, h, p, c, taglen, tag) =
    KAT_AEAD
        { aeadMode = mode
        , aeadKey = k
        , aeadIV = iv
        , aeadHeader = h
        , aeadPlaintext = p
        , aeadCiphertext = c
        , aeadTaglen = taglen
        , aeadTag = tag
        }
toKatGCM = toKatAEAD AEAD_GCM
toKatOCB = toKatAEAD AEAD_OCB

toKatCCM (k, iv, h, i, o, m) =
    KAT_AEAD
        { aeadMode = AEAD_CCM (B.length i) (ccmMVal m) CCM_L2
        , aeadKey = k
        , aeadIV = iv
        , aeadHeader = h
        , aeadPlaintext = i
        , aeadCiphertext = ct
        , aeadTaglen = m
        , aeadTag = at
        }
  where
    ccmMVal x =
        fromMaybe (error $ "unsupported CCM tag length: " ++ show x) $
            lookup
                x
                [ (4, CCM_M4)
                , (6, CCM_M6)
                , (8, CCM_M8)
                , (10, CCM_M10)
                , (12, CCM_M12)
                , (14, CCM_M14)
                , (16, CCM_M16)
                ]
    ctWithTag = B.drop (B.length h) o
    (ct, at) = B.splitAt (B.length ctWithTag - m) ctWithTag

kats128 =
    defaultKATs
        { kat_ECB = map toKatECB KATECB.vectors_aes128_enc
        , kat_CBC = map toKatCBC KATCBC.vectors_aes128_enc
        , kat_CTR = map toKatCTR KATCTR.vectors_aes128_enc
        , kat_CFB =
            [ KAT_CFB
                { cfbKey =
                    "\x2b\x7e\x15\x16\x28\xae\xd2\xa6\xab\xf7\x15\x88\x09\xcf\x4f\x3c"
                , cfbIV =
                    "\xC8\xA6\x45\x37\xA0\xB3\xA9\x3F\xCD\xE3\xCD\xAD\x9F\x1C\xE5\x8B"
                , cfbPlaintext =
                    "\x30\xc8\x1c\x46\xa3\x5c\xe4\x11\xe5\xfb\xc1\x19\x1a\x0a\x52\xef"
                , cfbCiphertext =
                    "\x26\x75\x1f\x67\xa3\xcb\xb1\x40\xb1\x80\x8c\xf1\x87\xa4\xf4\xdf"
                }
            ]
        , kat_XTS = map toKatXTS KATXTS.vectors_aes128_enc
        , kat_AEAD =
            map toKatGCM KATGCM.vectors_aes128_enc
                ++ map toKatOCB KATOCB3.vectors_aes128_enc
                ++ map toKatCCM KATCCM.vectors_aes128_enc
        }

kats192 =
    defaultKATs
        { kat_ECB = map toKatECB KATECB.vectors_aes192_enc
        , kat_CBC = map toKatCBC KATCBC.vectors_aes192_enc
        , kat_CTR = map toKatCTR KATCTR.vectors_aes192_enc
        , kat_AEAD =
            map toKatGCM KATGCM.vectors_aes192_enc
                ++ map toKatOCB KATOCB3.vectors_aes192_enc
                ++ map toKatCCM KATCCM.vectors_aes192_enc
        }

kats256 =
    defaultKATs
        { kat_ECB = map toKatECB KATECB.vectors_aes256_enc
        , kat_CBC = map toKatCBC KATCBC.vectors_aes256_enc
        , kat_CTR = map toKatCTR KATCTR.vectors_aes256_enc
        , kat_XTS = map toKatXTS KATXTS.vectors_aes256_enc
        , kat_AEAD =
            map toKatGCM KATGCM.vectors_aes256_enc
                ++ map toKatOCB KATOCB3.vectors_aes256_enc
                ++ map toKatCCM KATCCM.vectors_aes256_enc
        }

-- SP 800-38D 5.2.1.1: 1 <= len(IV) <= 2^64 - 1.  A zero-length IV makes
-- J0 the GHASH of the empty string, which leaks the authentication key.
aeadIVLengthTests :: Spec
aeadIVLengthTests =
    describe "AEAD IV length" $ do
        it "96-bit IV accepted" $
            isRight (initWith (B.replicate 12 0)) `shouldBe` True
        it "8-bit IV accepted" $
            isRight (initWith (B.replicate 1 0)) `shouldBe` True
        it "empty IV rejected" $
            initWith B.empty `shouldBe` Left CryptoError_IvSizeInvalid
  where
    ctx = throwCryptoError (cipherInit (B.replicate 16 0)) :: AES.AES128
    initWith iv =
        eitherCryptoError (() <$ aeadInit AEAD_GCM ctx (iv :: ByteString))
    isRight = either (const False) (const True)

aeadTagLengthTests :: Spec
aeadTagLengthTests =
    describe "AEAD tag length" $ do
        it "full tag verifies" $ openWith fullTag `shouldBe` Just message
        it "empty tag rejected" $ openWith B.empty `shouldBe` Nothing
        it "1-byte tag rejected" $ openWith (B.take 1 fullTag) `shouldBe` Nothing
        it "3-byte tag rejected" $ openWith (B.take 3 fullTag) `shouldBe` Nothing
        it "wrong tag rejected" $
            openWith (B.map (+ 1) fullTag) `shouldBe` Nothing
        -- a truncated tag is still at or above the minimum, so the length
        -- taken from the tag is the peer's choice of how much to verify
        it "4-byte tag accepted, since the tag sets the length" $
            openWith (B.take 4 fullTag) `shouldBe` Just message
        it "tryAeadSimpleDecrypt verifies the full tag" $
            openWith' 16 fullTag `shouldBe` Just message
        it "tryAeadSimpleDecrypt refuses a truncated tag" $
            openWith' 16 (B.take 4 fullTag) `shouldBe` Nothing
        it "tryAeadSimpleDecrypt refuses an overlong tag" $
            openWith' 16 (fullTag `B.append` B.singleton 0) `shouldBe` Nothing
        it "tryAeadSimpleDecrypt refuses a length below the minimum" $
            openWith' 3 (B.take 3 fullTag) `shouldBe` Nothing
        it "tryAeadSimpleDecrypt verifies a short tag the caller asked for" $
            openWith' 8 (B.take 8 fullTag) `shouldBe` Just message
        it "tryAeadSimpleDecrypt refuses a wrong tag" $
            openWith' 16 (B.map (+ 1) fullTag) `shouldBe` Nothing
  where
    key = B.replicate 16 0
    iv = B.replicate 12 0
    aad = "additional data" :: ByteString
    message = "authenticated message" :: ByteString
    ctx = throwCryptoError (cipherInit key) :: AES.AES128
    aead = throwCryptoError (aeadInit AEAD_GCM ctx iv)
    (AuthTag tag, ciphertext) = aeadSimpleEncrypt aead aad message 16
    fullTag = BA.convert tag :: ByteString
    openWith t = aeadSimpleDecrypt aead aad ciphertext (AuthTag (BA.convert t))
    openWith' n t = tryAeadSimpleDecrypt aead aad ciphertext n (AuthTag (BA.convert t))

-- The bulk loops -- eight blocks at a time under AES-NI, six at a time in
-- the assembly -- only start once the message is long enough to fill them,
-- and what they leave over goes down a different path.  These lengths sit
-- either side of each of those boundaries, so a group that hashes the wrong
-- blocks or a tail that is picked up at the wrong offset shows up here.
gcmLongTests :: Spec
gcmLongTests =
    describe "GCM long messages" $ mapM_ test KATGCMLong.vectors
  where
    test v@(klen, aadlen, ptlen, _, _) =
        it
            ( show klen
                ++ "-byte key, "
                ++ show aadlen
                ++ "-byte AAD, "
                ++ show ptlen
                ++ "-byte message"
            ) $
            case klen of
                16 -> run (undefined :: AES.AES128) v
                24 -> run (undefined :: AES.AES192) v
                _ -> run (undefined :: AES.AES256) v
    run
        :: BlockCipher cipher
        => cipher
        -> KATGCMLong.KATGCMLong
        -> Expectation
    run cipherWitness (klen, aadlen, ptlen, tag, ctHash) = do
        BA.convert authTag `shouldBe` tag
        digest ciphertext `shouldBe` ctHash
        aeadSimpleDecrypt aead aad ciphertext authTag `shouldBe` Just plaintext
      where
        cipher = throwCryptoError (cipherInit (KATGCMLong.gcmKey klen)) `asTypeOf` cipherWitness
        aead = throwCryptoError (aeadInit AEAD_GCM cipher KATGCMLong.gcmIV)
        aad = KATGCMLong.gcmAAD aadlen
        plaintext = KATGCMLong.gcmPlaintext ptlen
        (authTag, ciphertext) = aeadSimpleEncrypt aead aad plaintext 16
    digest bs = BA.convert (hash bs :: Digest SHA256) :: ByteString

spec :: Spec
spec = do
    testBlockCipher128 kats128 (undefined :: AES.AES128)
    testBlockCipher128 kats192 (undefined :: AES.AES192)
    testBlockCipher128 kats256 (undefined :: AES.AES256)
    aeadIVLengthTests
    aeadTagLengthTests
    gcmLongTests