{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE UnboxedTuples #-}

-- |
-- Module: Crypto.DRBG.HMAC.SHA256
-- Copyright: (c) 2024 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- A pure HMAC-DRBG implementation, as specified by
-- [NIST SP-800-90A](https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-90Ar1.pdf).

module Crypto.DRBG.HMAC.SHA256 (
  -- * DRBG and HMAC function types
    DRBG
  , Error(..)

  -- * DRBG interaction
  , new
  , gen
  , reseed
  , wipe

  -- for testing
  , _read_v
  , _read_k
  ) where

import Crypto.DRBG.HMAC.Internal (Error(..), _RESEED_COUNTER, _MAX_BYTES)
import qualified Crypto.Hash.SHA256 as SHA256
import Crypto.Hash.SHA256.Internal (Registers(..))
import qualified Crypto.Hash.SHA256.Internal as SHA256 (cat)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Control.Monad.ST (ST)
import GHC.Exts (RealWorld)
import qualified Control.Monad.Primitive as Prim (unsafeIOToPrim)
import Data.Bits ((.<<.), (.>>.), (.|.))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as BSB
import qualified Data.ByteString.Internal as BI
import qualified Data.Primitive.PrimArray as PA
import Data.Word (Word32, Word64)
import qualified GHC.Word
import qualified Foreign.Ptr as FP

-- api ------------------------------------------------------------------------

-- | A deterministic random bit generator (DRBG).
--
--   Create a DRBG with 'new', and then use and reuse it to generate
--   bytes as needed.
--
--   >>> drbg <- new entropy nonce personalization_string
--   >>> bytes0 <- gen drbg mempty 10
--   >>> bytes1 <- gen drbg mempty 10
--   >>> drbg
--   "<drbg>"

-- first two elements are hi/lo bits of word64 counter
-- next eight elements are k
-- next eight elements are v
-- next sixteen elements are scratch space
newtype DRBG s = DRBG (PA.MutablePrimArray s Word32)

instance Show (DRBG s) where
  show :: DRBG s -> String
show DRBG s
_ = String
"<drbg>"

-- | Create a HMAC-SHA256 DRBG from the supplied entropy, nonce, and
--   personalization string.
--
--   The DRBG is returned in any 'PrimMonad', e.g. 'ST s' or 'IO'.
--
--   >>> new entropy nonce personalization_string
--   "<drbg>"
new
  :: PrimMonad m
  => BS.ByteString    -- ^ entropy
  -> BS.ByteString    -- ^ nonce
  -> BS.ByteString    -- ^ personalization string
  -> m (DRBG (PrimState m))
new :: forall (m :: * -> *).
PrimMonad m =>
ByteString -> ByteString -> ByteString -> m (DRBG (PrimState m))
new ByteString
entropy ByteString
nonce ByteString
ps = do
  drbg <- Int -> m (MutablePrimArray (PrimState m) Word32)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
PA.newPinnedPrimArray Int
34 -- 2 (ctr) + 16 (k, v) + 16 (scratch)
  init_counter drbg
  PA.setPrimArray drbg 02 08 (0x00000000 :: Word32) -- init k
  PA.setPrimArray drbg 10 08 (0x01010101 :: Word32) -- init v
  PA.setPrimArray drbg 18 16 (0x00000000 :: Word32) -- scratch
  update drbg (entropy <> nonce <> ps)
  pure $! DRBG drbg
{-# INLINABLE new #-}
{-# SPECIALIZE new
  :: BS.ByteString -> BS.ByteString -> BS.ByteString -> IO (DRBG RealWorld) #-}
{-# SPECIALIZE new
  :: BS.ByteString -> BS.ByteString -> BS.ByteString -> ST s (DRBG s) #-}

-- | Reseed a DRBG.
--
--   Each DRBG has an internal /reseed counter/ that tracks the number
--   of requests made to the generator (note /requests made/, not bytes
--   generated). SP 800-90A specifies that a HMAC-DRBG should support
--   2 ^ 48 requests before requiring a reseed, so in practice you're
--   unlikely to ever need to use this to actually reset the counter.
--
--   Note however that 'reseed' can be used to implement "explicit"
--   prediction resistance, per SP 800-90A, by injecting entropy generated
--   elsewhere into the DRBG.
--
--   >>> import qualified System.Entropy as E
--   >>> entropy <- E.getEntropy 32
--   >>> reseed entropy addl_bytes drbg
--   "<reseeded drbg>"
reseed
  :: PrimMonad m
  => DRBG (PrimState m)
  -> BS.ByteString
  -> BS.ByteString
  -> m ()
reseed :: forall (m :: * -> *).
PrimMonad m =>
DRBG (PrimState m) -> ByteString -> ByteString -> m ()
reseed (DRBG MutablePrimArray (PrimState m) Word32
drbg) ByteString
entr ByteString
addl = do
  MutablePrimArray (PrimState m) Word32 -> ByteString -> m ()
forall (m :: * -> *).
PrimMonad m =>
MutablePrimArray (PrimState m) Word32 -> ByteString -> m ()
update MutablePrimArray (PrimState m) Word32
drbg (ByteString
entr ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
addl)
  MutablePrimArray (PrimState m) Word32 -> m ()
forall (m :: * -> *).
PrimMonad m =>
MutablePrimArray (PrimState m) Word32 -> m ()
init_counter MutablePrimArray (PrimState m) Word32
drbg
{-# INLINE reseed #-}

-- | Generate bytes from a DRBG, optionally injecting additional bytes
--   per SP 800-90A.
--
--   Per SP 800-90A, the maximum number of bytes that can be requested
--   on any invocation is 65536. Larger requests will return
--   'MaxBytesExceeded'.
--
--   >>> import qualified Data.ByteString.Base16 as B16
--   >>> drbg <- new entropy nonce personalization_string
--   >>> Right bytes0 <- gen drbg addl_bytes 16
--   >>> Right bytes1 <- gen drbg addl_bytes 16
--   >>> B16.encode bytes0
--   "938d6ca6d0b797f7b3c653349d6e3135"
--   >>> B16.encode bytes1
--   "5f379d16de6f2c6f8a35c56f13f9e5a5"
gen
  :: PrimMonad m
  => DRBG (PrimState m)
  -> BS.ByteString
  -> Word64
  -> m (Either Error BS.ByteString)
gen :: forall (m :: * -> *).
PrimMonad m =>
DRBG (PrimState m)
-> ByteString -> Word64 -> m (Either Error ByteString)
gen (DRBG MutablePrimArray (PrimState m) Word32
drbg) addl :: ByteString
addl@(BI.PS ForeignPtr Word8
_ Int
_ Int
l) Word64
bytes
  | Word64
bytes Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
_MAX_BYTES = Either Error ByteString -> m (Either Error ByteString)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Error ByteString -> m (Either Error ByteString))
-> Either Error ByteString -> m (Either Error ByteString)
forall a b. (a -> b) -> a -> b
$! Error -> Either Error ByteString
forall a b. a -> Either a b
Left Error
MaxBytesExceeded
  | Bool
otherwise = do
      ctr <- MutablePrimArray (PrimState m) Word32 -> m Word64
forall (m :: * -> *).
PrimMonad m =>
MutablePrimArray (PrimState m) Word32 -> m Word64
read_counter MutablePrimArray (PrimState m) Word32
drbg
      if   ctr > _RESEED_COUNTER
      then pure $! Left ReseedRequired
      else do
        if l == 0 then pure () else update drbg addl
        !(GHC.Word.W32# k00) <- PA.readPrimArray drbg 02
        !(GHC.Word.W32# k01) <- PA.readPrimArray drbg 03
        !(GHC.Word.W32# k02) <- PA.readPrimArray drbg 04
        !(GHC.Word.W32# k03) <- PA.readPrimArray drbg 05
        !(GHC.Word.W32# k04) <- PA.readPrimArray drbg 06
        !(GHC.Word.W32# k05) <- PA.readPrimArray drbg 07
        !(GHC.Word.W32# k06) <- PA.readPrimArray drbg 08
        !(GHC.Word.W32# k07) <- PA.readPrimArray drbg 09
        !(GHC.Word.W32# v00) <- PA.readPrimArray drbg 10
        !(GHC.Word.W32# v01) <- PA.readPrimArray drbg 11
        !(GHC.Word.W32# v02) <- PA.readPrimArray drbg 12
        !(GHC.Word.W32# v03) <- PA.readPrimArray drbg 13
        !(GHC.Word.W32# v04) <- PA.readPrimArray drbg 14
        !(GHC.Word.W32# v05) <- PA.readPrimArray drbg 15
        !(GHC.Word.W32# v06) <- PA.readPrimArray drbg 16
        !(GHC.Word.W32# v07) <- PA.readPrimArray drbg 17
        let !k0 = (# Word32#, Word32#, Word32#, Word32#, Word32#, Word32#, Word32#,
   Word32# #)
-> Registers
Registers (# Word32#
k00, Word32#
k01, Word32#
k02, Word32#
k03, Word32#
k04, Word32#
k05, Word32#
k06, Word32#
k07 #)
            !v0 = (# Word32#, Word32#, Word32#, Word32#, Word32#, Word32#, Word32#,
   Word32# #)
-> Registers
Registers (# Word32#
v00, Word32#
v01, Word32#
v02, Word32#
v03, Word32#
v04, Word32#
v05, Word32#
v06, Word32#
v07 #)
        !res <- gen_loop drbg k0 v0 bytes
        update drbg addl
        write_counter drbg (ctr + 1)
        pure $! Right res
{-# INLINABLE gen #-}
{-# SPECIALIZE gen
  :: DRBG RealWorld -> BS.ByteString -> Word64
  -> IO (Either Error BS.ByteString) #-}
{-# SPECIALIZE gen
  :: DRBG s -> BS.ByteString -> Word64
  -> ST s (Either Error BS.ByteString) #-}

-- | Wipe the state of a DRBG.
--
--   You should call this when you're finished with a DRBG to ensure that its
--   state is wiped from memory.
--
--   >>> drbg <- new mempty mempty mempty
--   >>> Right bytes <- gen drbg addl_bytes 16
--   >>> wipe drbg
--   >>> -- do something with bytes
wipe
  :: PrimMonad m
  => DRBG (PrimState m)
  -> m ()
wipe :: forall (m :: * -> *). PrimMonad m => DRBG (PrimState m) -> m ()
wipe (DRBG MutablePrimArray (PrimState m) Word32
drbg) = do
  MutablePrimArray (PrimState m) Word32 -> m ()
forall (m :: * -> *).
PrimMonad m =>
MutablePrimArray (PrimState m) Word32 -> m ()
init_counter MutablePrimArray (PrimState m) Word32
drbg
  MutablePrimArray (PrimState m) Word32
-> Int -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> a -> m ()
PA.setPrimArray MutablePrimArray (PrimState m) Word32
drbg Int
02 Int
08 (Word32
0x00000000 :: Word32) -- init k
  MutablePrimArray (PrimState m) Word32
-> Int -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> a -> m ()
PA.setPrimArray MutablePrimArray (PrimState m) Word32
drbg Int
10 Int
08 (Word32
0x01010101 :: Word32) -- init v
  MutablePrimArray (PrimState m) Word32
-> Int -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> a -> m ()
PA.setPrimArray MutablePrimArray (PrimState m) Word32
drbg Int
18 Int
16 (Word32
0x00000000 :: Word32) -- init scratch
{-# INLINE wipe #-}
-- utilities ------------------------------------------------------------------

fi :: (Integral a, Num b) => a -> b
fi :: forall a b. (Integral a, Num b) => a -> b
fi = a -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE fi #-}

-- drbg utilities -------------------------------------------------------------

gen_loop
  :: PrimMonad m
  => PA.MutablePrimArray (PrimState m) Word32
  -> Registers
  -> Registers
  -> Word64
  -> m BS.ByteString
gen_loop :: forall (m :: * -> *).
PrimMonad m =>
MutablePrimArray (PrimState m) Word32
-> Registers -> Registers -> Word64 -> m ByteString
gen_loop MutablePrimArray (PrimState m) Word32
drbg Registers
k0 Registers
v0 Word64
bytes = ByteString -> Registers -> Word64 -> m ByteString
forall {m :: * -> *}.
(PrimState m ~ PrimState m, PrimMonad m) =>
ByteString -> Registers -> Word64 -> m ByteString
loop ByteString
forall a. Monoid a => a
mempty Registers
v0 Word64
0 where
  !vp :: Ptr b
vp = MutablePrimArray (PrimState m) Word32 -> Ptr Word32
forall s a. MutablePrimArray s a -> Ptr a
PA.mutablePrimArrayContents MutablePrimArray (PrimState m) Word32
drbg Ptr Word32 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`FP.plusPtr` Int
40 -- 10 * 4
  !sp :: Ptr b
sp = MutablePrimArray (PrimState m) Word32 -> Ptr Word32
forall s a. MutablePrimArray s a -> Ptr a
PA.mutablePrimArrayContents MutablePrimArray (PrimState m) Word32
drbg Ptr Word32 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`FP.plusPtr` Int
72 -- 18 * 4
  loop :: ByteString -> Registers -> Word64 -> m ByteString
loop !ByteString
acc Registers
v Word64
l
    | Word64
l Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word64
bytes = do
        MutablePrimArray (PrimState m) Word32 -> Registers -> m ()
forall (m :: * -> *).
PrimMonad m =>
MutablePrimArray (PrimState m) Word32 -> Registers -> m ()
write_v MutablePrimArray (PrimState m) Word32
MutablePrimArray (PrimState m) Word32
drbg Registers
v
        ByteString -> m ByteString
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
acc
    | Bool
otherwise = do
        IO () -> m ()
forall (m :: * -> *) a. PrimMonad m => IO a -> m a
Prim.unsafeIOToPrim (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Ptr Word32 -> Ptr Word32 -> Registers -> Registers -> IO ()
SHA256._hmac_rr Ptr Word32
forall {b}. Ptr b
vp Ptr Word32
forall {b}. Ptr b
sp Registers
k0 Registers
v
        !(GHC.Word.W32# nv0) <- MutablePrimArray (PrimState m) Word32 -> Int -> m Word32
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
PA.readPrimArray MutablePrimArray (PrimState m) Word32
MutablePrimArray (PrimState m) Word32
drbg Int
10
        !(GHC.Word.W32# nv1) <- PA.readPrimArray drbg 11
        !(GHC.Word.W32# nv2) <- PA.readPrimArray drbg 12
        !(GHC.Word.W32# nv3) <- PA.readPrimArray drbg 13
        !(GHC.Word.W32# nv4) <- PA.readPrimArray drbg 14
        !(GHC.Word.W32# nv5) <- PA.readPrimArray drbg 15
        !(GHC.Word.W32# nv6) <- PA.readPrimArray drbg 16
        !(GHC.Word.W32# nv7) <- PA.readPrimArray drbg 17
        let !nv = (# Word32#, Word32#, Word32#, Word32#, Word32#, Word32#, Word32#,
   Word32# #)
-> Registers
Registers (# Word32#
nv0, Word32#
nv1, Word32#
nv2, Word32#
nv3, Word32#
nv4, Word32#
nv5, Word32#
nv6, Word32#
nv7 #)
            !na = ByteString
acc ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Registers -> ByteString
SHA256.cat Registers
nv
            !nl = Word64
l Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
32
        loop na nv nl
{-# INLINE gen_loop #-}

update
  :: PrimMonad m
  => PA.MutablePrimArray (PrimState m) Word32
  -> BS.ByteString
  -> m ()
update :: forall (m :: * -> *).
PrimMonad m =>
MutablePrimArray (PrimState m) Word32 -> ByteString -> m ()
update MutablePrimArray (PrimState m) Word32
drbg provided_data :: ByteString
provided_data@(BI.PS ForeignPtr Word8
_ Int
_ Int
l) = do
  !(GHC.Word.W32# k00) <- MutablePrimArray (PrimState m) Word32 -> Int -> m Word32
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
PA.readPrimArray MutablePrimArray (PrimState m) Word32
drbg Int
02
  !(GHC.Word.W32# k01) <- PA.readPrimArray drbg 03
  !(GHC.Word.W32# k02) <- PA.readPrimArray drbg 04
  !(GHC.Word.W32# k03) <- PA.readPrimArray drbg 05
  !(GHC.Word.W32# k04) <- PA.readPrimArray drbg 06
  !(GHC.Word.W32# k05) <- PA.readPrimArray drbg 07
  !(GHC.Word.W32# k06) <- PA.readPrimArray drbg 08
  !(GHC.Word.W32# k07) <- PA.readPrimArray drbg 09
  !(GHC.Word.W32# v00) <- PA.readPrimArray drbg 10
  !(GHC.Word.W32# v01) <- PA.readPrimArray drbg 11
  !(GHC.Word.W32# v02) <- PA.readPrimArray drbg 12
  !(GHC.Word.W32# v03) <- PA.readPrimArray drbg 13
  !(GHC.Word.W32# v04) <- PA.readPrimArray drbg 14
  !(GHC.Word.W32# v05) <- PA.readPrimArray drbg 15
  !(GHC.Word.W32# v06) <- PA.readPrimArray drbg 16
  !(GHC.Word.W32# v07) <- PA.readPrimArray drbg 17
  let !k0 = (# Word32#, Word32#, Word32#, Word32#, Word32#, Word32#, Word32#,
   Word32# #)
-> Registers
Registers (# Word32#
k00, Word32#
k01, Word32#
k02, Word32#
k03, Word32#
k04, Word32#
k05, Word32#
k06, Word32#
k07 #)
      !v0 = (# Word32#, Word32#, Word32#, Word32#, Word32#, Word32#, Word32#,
   Word32# #)
-> Registers
Registers (# Word32#
v00, Word32#
v01, Word32#
v02, Word32#
v03, Word32#
v04, Word32#
v05, Word32#
v06, Word32#
v07 #)
      !kp = MutablePrimArray (PrimState m) Word32 -> Ptr Word32
forall s a. MutablePrimArray s a -> Ptr a
PA.mutablePrimArrayContents MutablePrimArray (PrimState m) Word32
drbg Ptr Word32 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`FP.plusPtr` Int
08 --  2 * 4
      !vp = MutablePrimArray (PrimState m) Word32 -> Ptr Word32
forall s a. MutablePrimArray s a -> Ptr a
PA.mutablePrimArrayContents MutablePrimArray (PrimState m) Word32
drbg Ptr Word32 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`FP.plusPtr` Int
40 -- 10 * 4
      !sp = MutablePrimArray (PrimState m) Word32 -> Ptr Word32
forall s a. MutablePrimArray s a -> Ptr a
PA.mutablePrimArrayContents MutablePrimArray (PrimState m) Word32
drbg Ptr Word32 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`FP.plusPtr` Int
72 -- 18 * 4
  Prim.unsafeIOToPrim $ SHA256._hmac_rsb kp sp k0 v0 0x00 provided_data
  !(GHC.Word.W32# k10) <- PA.readPrimArray drbg 02
  !(GHC.Word.W32# k11) <- PA.readPrimArray drbg 03
  !(GHC.Word.W32# k12) <- PA.readPrimArray drbg 04
  !(GHC.Word.W32# k13) <- PA.readPrimArray drbg 05
  !(GHC.Word.W32# k14) <- PA.readPrimArray drbg 06
  !(GHC.Word.W32# k15) <- PA.readPrimArray drbg 07
  !(GHC.Word.W32# k16) <- PA.readPrimArray drbg 08
  !(GHC.Word.W32# k17) <- PA.readPrimArray drbg 09
  let !k1 = (# Word32#, Word32#, Word32#, Word32#, Word32#, Word32#, Word32#,
   Word32# #)
-> Registers
Registers (# Word32#
k10, Word32#
k11, Word32#
k12, Word32#
k13, Word32#
k14, Word32#
k15, Word32#
k16, Word32#
k17 #)
  Prim.unsafeIOToPrim $ SHA256._hmac_rr vp sp k1 v0
  if   l == 0
  then pure ()
  else do
    !(GHC.Word.W32# v10) <- PA.readPrimArray drbg 10
    !(GHC.Word.W32# v11) <- PA.readPrimArray drbg 11
    !(GHC.Word.W32# v12) <- PA.readPrimArray drbg 12
    !(GHC.Word.W32# v13) <- PA.readPrimArray drbg 13
    !(GHC.Word.W32# v14) <- PA.readPrimArray drbg 14
    !(GHC.Word.W32# v15) <- PA.readPrimArray drbg 15
    !(GHC.Word.W32# v16) <- PA.readPrimArray drbg 16
    !(GHC.Word.W32# v17) <- PA.readPrimArray drbg 17
    let !v1 = (# Word32#, Word32#, Word32#, Word32#, Word32#, Word32#, Word32#,
   Word32# #)
-> Registers
Registers (# Word32#
v10, Word32#
v11, Word32#
v12, Word32#
v13, Word32#
v14, Word32#
v15, Word32#
v16, Word32#
v17 #)
    Prim.unsafeIOToPrim $ SHA256._hmac_rsb kp sp k1 v1 0x01 provided_data
    !(GHC.Word.W32# k20) <- PA.readPrimArray drbg 02
    !(GHC.Word.W32# k21) <- PA.readPrimArray drbg 03
    !(GHC.Word.W32# k22) <- PA.readPrimArray drbg 04
    !(GHC.Word.W32# k23) <- PA.readPrimArray drbg 05
    !(GHC.Word.W32# k24) <- PA.readPrimArray drbg 06
    !(GHC.Word.W32# k25) <- PA.readPrimArray drbg 07
    !(GHC.Word.W32# k26) <- PA.readPrimArray drbg 08
    !(GHC.Word.W32# k27) <- PA.readPrimArray drbg 09
    let !k2 = (# Word32#, Word32#, Word32#, Word32#, Word32#, Word32#, Word32#,
   Word32# #)
-> Registers
Registers (# Word32#
k20, Word32#
k21, Word32#
k22, Word32#
k23, Word32#
k24, Word32#
k25, Word32#
k26, Word32#
k27 #)
    Prim.unsafeIOToPrim $ SHA256._hmac_rr vp sp k2 v1
{-# INLINABLE update #-}
{-# SPECIALIZE update
  :: PA.MutablePrimArray RealWorld Word32 -> BS.ByteString -> IO () #-}
{-# SPECIALIZE update
  :: PA.MutablePrimArray s Word32 -> BS.ByteString -> ST s () #-}

init_counter
  :: PrimMonad m
  => PA.MutablePrimArray (PrimState m) Word32
  -> m ()
init_counter :: forall (m :: * -> *).
PrimMonad m =>
MutablePrimArray (PrimState m) Word32 -> m ()
init_counter MutablePrimArray (PrimState m) Word32
drbg = do
  MutablePrimArray (PrimState m) Word32 -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray (PrimState m) Word32
drbg Int
0 (Word32
0x00 :: Word32) -- init high word, counter
  MutablePrimArray (PrimState m) Word32 -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray (PrimState m) Word32
drbg Int
1 (Word32
0x01 :: Word32) -- init low word, counter
{-# INLINE init_counter #-}

read_counter
  :: PrimMonad m
  => PA.MutablePrimArray (PrimState m) Word32
  -> m Word64
read_counter :: forall (m :: * -> *).
PrimMonad m =>
MutablePrimArray (PrimState m) Word32 -> m Word64
read_counter MutablePrimArray (PrimState m) Word32
drbg = do
  !hi <- MutablePrimArray (PrimState m) Word32 -> Int -> m Word32
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
PA.readPrimArray MutablePrimArray (PrimState m) Word32
drbg Int
0
  !lo <- PA.readPrimArray drbg 1
  let !ctr = Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fi Word32
hi Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
.<<. Int
32 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fi Word32
lo
  pure $! ctr
{-# INLINE read_counter #-}

write_counter
  :: PrimMonad m
  => PA.MutablePrimArray (PrimState m) Word32
  -> Word64
  -> m ()
write_counter :: forall (m :: * -> *).
PrimMonad m =>
MutablePrimArray (PrimState m) Word32 -> Word64 -> m ()
write_counter MutablePrimArray (PrimState m) Word32
drbg Word64
ctr = do
  let !hi :: Word32
hi = Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fi (Word64
ctr Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
.>>. Int
32)
      !lo :: Word32
lo = Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fi Word64
ctr
  MutablePrimArray (PrimState m) Word32 -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray (PrimState m) Word32
drbg Int
0 Word32
hi
  MutablePrimArray (PrimState m) Word32 -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray (PrimState m) Word32
drbg Int
1 Word32
lo
{-# INLINE write_counter #-}

write_v
  :: PrimMonad m
  => PA.MutablePrimArray (PrimState m) Word32
  -> Registers
  -> m ()
write_v :: forall (m :: * -> *).
PrimMonad m =>
MutablePrimArray (PrimState m) Word32 -> Registers -> m ()
write_v MutablePrimArray (PrimState m) Word32
drbg (R Word32#
v0 Word32#
v1 Word32#
v2 Word32#
v3 Word32#
v4 Word32#
v5 Word32#
v6 Word32#
v7) = do
  MutablePrimArray (PrimState m) Word32 -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray (PrimState m) Word32
drbg Int
10 (Word32# -> Word32
GHC.Word.W32# Word32#
v0)
  MutablePrimArray (PrimState m) Word32 -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray (PrimState m) Word32
drbg Int
11 (Word32# -> Word32
GHC.Word.W32# Word32#
v1)
  MutablePrimArray (PrimState m) Word32 -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray (PrimState m) Word32
drbg Int
12 (Word32# -> Word32
GHC.Word.W32# Word32#
v2)
  MutablePrimArray (PrimState m) Word32 -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray (PrimState m) Word32
drbg Int
13 (Word32# -> Word32
GHC.Word.W32# Word32#
v3)
  MutablePrimArray (PrimState m) Word32 -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray (PrimState m) Word32
drbg Int
14 (Word32# -> Word32
GHC.Word.W32# Word32#
v4)
  MutablePrimArray (PrimState m) Word32 -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray (PrimState m) Word32
drbg Int
15 (Word32# -> Word32
GHC.Word.W32# Word32#
v5)
  MutablePrimArray (PrimState m) Word32 -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray (PrimState m) Word32
drbg Int
16 (Word32# -> Word32
GHC.Word.W32# Word32#
v6)
  MutablePrimArray (PrimState m) Word32 -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray (PrimState m) Word32
drbg Int
17 (Word32# -> Word32
GHC.Word.W32# Word32#
v7)
{-# INLINE write_v #-}

-- read secret drbg state (for testing)
_read_v
  :: PrimMonad m
  => DRBG (PrimState m)
  -> m BS.ByteString
_read_v :: forall (m :: * -> *).
PrimMonad m =>
DRBG (PrimState m) -> m ByteString
_read_v (DRBG MutablePrimArray (PrimState m) Word32
drbg) = do
  !v00 <- MutablePrimArray (PrimState m) Word32 -> Int -> m Word32
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
PA.readPrimArray MutablePrimArray (PrimState m) Word32
drbg Int
10
  !v01 <- PA.readPrimArray drbg 11
  !v02 <- PA.readPrimArray drbg 12
  !v03 <- PA.readPrimArray drbg 13
  !v04 <- PA.readPrimArray drbg 14
  !v05 <- PA.readPrimArray drbg 15
  !v06 <- PA.readPrimArray drbg 16
  !v07 <- PA.readPrimArray drbg 17
  pure . BS.toStrict . BSB.toLazyByteString $ mconcat [
      BSB.word32BE v00
    , BSB.word32BE v01
    , BSB.word32BE v02
    , BSB.word32BE v03
    , BSB.word32BE v04
    , BSB.word32BE v05
    , BSB.word32BE v06
    , BSB.word32BE v07
    ]

-- read secret drbg state (for testing)
_read_k
  :: PrimMonad m
  => DRBG (PrimState m)
  -> m BS.ByteString
_read_k :: forall (m :: * -> *).
PrimMonad m =>
DRBG (PrimState m) -> m ByteString
_read_k (DRBG MutablePrimArray (PrimState m) Word32
drbg) = do
  !k00 <- MutablePrimArray (PrimState m) Word32 -> Int -> m Word32
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
PA.readPrimArray MutablePrimArray (PrimState m) Word32
drbg Int
02
  !k01 <- PA.readPrimArray drbg 03
  !k02 <- PA.readPrimArray drbg 04
  !k03 <- PA.readPrimArray drbg 05
  !k04 <- PA.readPrimArray drbg 06
  !k05 <- PA.readPrimArray drbg 07
  !k06 <- PA.readPrimArray drbg 08
  !k07 <- PA.readPrimArray drbg 09
  pure . BS.toStrict . BSB.toLazyByteString $ mconcat [
      BSB.word32BE k00
    , BSB.word32BE k01
    , BSB.word32BE k02
    , BSB.word32BE k03
    , BSB.word32BE k04
    , BSB.word32BE k05
    , BSB.word32BE k06
    , BSB.word32BE k07
    ]