{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}

module ResumableHash (tests) where

import Crypto.Hash ( SHAKE128(..), SHAKE256(..), SHA3_224(..), SHA3_256(..), SHA3_384(..), SHA3_512(..), Keccak_224(..), Keccak_256(..), Keccak_384(..), Keccak_512(..)
                   , HashAlgorithm, HashAlgorithmResumable, Context, hashPutContext, hashGetContext)
import qualified Crypto.Hash as Hash

import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as LB
import qualified Data.ByteString.Builder as Builder

import Data.Maybe (fromMaybe)

import Imports

data HashResume a = HashResume [ByteString] [ByteString] (Hash.Digest a)
  deriving (Show, Eq)

instance HashAlgorithm a => Arbitrary (HashResume a) where
  arbitrary = do
    (beforeChunks, afterChunks) <- oneof
      [ ([], ) <$> (choose (1,20) >>= \n -> replicateM n (arbitraryBSof 1 99))
      , (,) <$> (choose (1,20) >>= \n -> replicateM n (arbitraryBSof 1 99))
            <*> (choose (1,20) >>= \n -> replicateM n (arbitraryBSof 1 99))
      , (, []) <$> (choose (1,20) >>= \n -> replicateM n (arbitraryBSof 1 99))
      , pure ([], [])
      ]
    return $ HashResume beforeChunks afterChunks (Hash.hashlazy . LB.fromChunks $ beforeChunks ++ afterChunks)

resumeTests =
    [ testResumeProperties "SHAKE128_256" (SHAKE128 :: SHAKE128 256)
    , testResumeProperties "SHAKE256_256" (SHAKE256 :: SHAKE256 512)
    , testResumeProperties "SHA3_224" SHA3_224
    , testResumeProperties "SHA3_256" SHA3_256
    , testResumeProperties "SHA3_384" SHA3_384
    , testResumeProperties "SHA3_512" SHA3_512
    , testResumeProperties "Keccak_224" Keccak_224
    , testResumeProperties "Keccak_256" Keccak_256
    , testResumeProperties "Keccak_384" Keccak_384
    , testResumeProperties "Keccak_512" Keccak_512
    , testCase "serializes big endian" $ test_is_be 168 (SHAKE128 :: SHAKE128 256)
    ]
  where
        testResumeProperties :: HashAlgorithmResumable a => TestName -> a -> TestTree
        testResumeProperties name a = testGroup name
          [ testProperty "resume" (prop_resume_start a)
          ]

        prop_resume_start :: forall a. HashAlgorithmResumable a => a -> HashResume a -> Bool
        prop_resume_start _ (HashResume beforeChunks afterChunks result) = fromMaybe False $ do
          let beforeCtx = Hash.hashUpdates (Hash.hashInit :: Context a) beforeChunks
          ctx <- hashGetContext (hashPutContext beforeCtx :: ByteString)
          let afterCtx = Hash.hashUpdates ctx afterChunks
          return $ result `assertEq` Hash.hashFinalize afterCtx

        test_is_be :: forall a. HashAlgorithmResumable a => Word32 -> a -> Assertion
        test_is_be size _ = slice @=? size_be
          where size_be = LB.toStrict $ Builder.toLazyByteString $ Builder.word32BE size
                serialized = hashPutContext (Hash.hashInit :: Context a) :: ByteString
                slice = B.take 4 $ B.drop 4 serialized

tests = testGroup "ResumableHash" resumeTests