    {-# LANGUAGE BangPatterns #-}
    module Crypto.Random.DRBG.HMAC
            ( State, counter
            , reseedInterval
            , instantiate
            , reseed
            , generate) where

    import qualified Data.ByteString as B
    import qualified Data.ByteString.Lazy as L
    import Data.Serialize (encode, Serialize(..))
    import Data.Serialize.Put
    import Data.Serialize.Builder (toByteString)
    import Data.Tagged (proxy)
    import Data.Word (Word64)
    import Crypto.Classes
    import Crypto.HMAC
    import Crypto.Types
    import Crypto.Random.DRBG.Types

    type Key = B.ByteString
    type Value = B.ByteString

    data State d = St
            { counter               :: {-# UNPACK #-} !Word64
            -- Start admin info
            , value                 :: !Value
            , key                   :: !Key
            }

    -- This is available with the right type in the tagged package starting from
    -- version 0.7, but ending with GHC version 7.8. Sigh.
    asProxyTypeOf :: d -> state d -> d
    asProxyTypeOf = const

    reseedInterval :: Word64
    reseedInterval = 2^48

    fc = L.fromChunks . \s -> [s]

    update :: (Hash c d) => State d -> L.ByteString -> State d
    update st input = st { value = newV , key = newK }
      where
      hm k = hmac (MacKey k)
      k    = key st
      v    = value st
      k'   = encode $ (hm k (L.concat [fc v, L.singleton 0, input]) `asProxyTypeOf` st)
      v'   = encode $ (hm k' (fc v) `asProxyTypeOf` st)
      (newK, newV) =
        if L.length input == 0
          then (k',v')
          else let k'' = encode $ hm k' (L.concat [fc v', L.singleton 1, input]) `asProxyTypeOf` st
               in (k'', encode $ hm k'' (fc v') `asProxyTypeOf` st)

    instantiate :: (Hash c d) => Entropy -> Nonce -> PersonalizationString -> State d
    instantiate ent nonce perStr = st
      where
      seedMaterial = L.fromChunks [ent, nonce, perStr]
      k = B.replicate olen 0
      v = B.replicate olen 1
      st = update (St 1 v k) seedMaterial
      olen = (outputLength `proxy` st) `div` 8

    reseed :: (Hash c d) => State d -> Entropy -> AdditionalInput -> State d
    reseed st ent ai = (update st (L.fromChunks [ent, ai])) { counter = 1 }

    generate :: (Hash c d) => State d -> BitLength -> AdditionalInput -> Maybe (RandomBits, State d)
    generate st req additionalInput =
            if(counter st > reseedInterval)
                    then Nothing
                    else Just (randBitsFinal, stFinal { counter = 1 + counter st})
      where
      st' = if B.length additionalInput == 0
                    then st
                    else update st (fc additionalInput)
      reqBytes = (req+7) `div` 8
      iterations = (reqBytes + (outlen - 1)) `div` outlen

      -- getV is the main cost.  HMACing and storing 'iterations' bytestrings at
      -- ~64 bytes each is a real waste.  Some pre-allocation and unsafe functions
      -- exported from Crypto.HMAC could cut this down, but it really isn't worth
      -- giving CPR to such a bad idea as using ByteString for crypto computations
      getV :: Value -> Int -> (Value, [B.ByteString])
      getV !u 0 = (u, [])
      getV !u i =
            let !vNew = hmac' (MacKey kFinal) u `asProxyTypeOf` st
                !encV = encode vNew
                (uFinal, rest) = getV encV (i - 1)
            in (uFinal, encV:rest)
      (vFinal, randBitsList) = getV (value st') iterations
      randBitsFinal = B.take reqBytes $ B.concat randBitsList
      kFinal = key st'
      stFinal = update (st' { key = kFinal, value = vFinal} `asTypeOf` st) (fc additionalInput)
      outlen = (outputLength `proxy` st) `div` 8