{-# OPTIONS_HADDOCK prune #-}
{-# OPTIONS_GHC -funbox-small-strict-fields #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE LambdaCase #-}

-- |
-- Module: Crypto.DRBG.HMAC
-- 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 (
  -- * DRBG and HMAC function types
    DRBG
  , _read_v
  , _read_k
  , HMAC

  -- * DRBG interaction
  , new
  , gen
  , reseed
  ) where

import Control.Monad.Primitive (PrimMonad, PrimState)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as BSB
import qualified Data.Primitive.MutVar as P
import Data.Word (Word64)

-- keystroke savers and 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 #-}

toStrict :: BSB.Builder -> BS.ByteString
toStrict :: Builder -> ByteString
toStrict = ByteString -> ByteString
BS.toStrict (ByteString -> ByteString)
-> (Builder -> ByteString) -> Builder -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> ByteString
BSB.toLazyByteString
{-# INLINE toStrict #-}

-- dumb strict pair
data Pair a b = Pair !a !b
  deriving Int -> Pair a b -> ShowS
[Pair a b] -> ShowS
Pair a b -> String
(Int -> Pair a b -> ShowS)
-> (Pair a b -> String) -> ([Pair a b] -> ShowS) -> Show (Pair a b)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall a b. (Show a, Show b) => Int -> Pair a b -> ShowS
forall a b. (Show a, Show b) => [Pair a b] -> ShowS
forall a b. (Show a, Show b) => Pair a b -> String
$cshowsPrec :: forall a b. (Show a, Show b) => Int -> Pair a b -> ShowS
showsPrec :: Int -> Pair a b -> ShowS
$cshow :: forall a b. (Show a, Show b) => Pair a b -> String
show :: Pair a b -> String
$cshowList :: forall a b. (Show a, Show b) => [Pair a b] -> ShowS
showList :: [Pair a b] -> ShowS
Show

-- types ----------------------------------------------------------------------

-- see SP 800-90A table 2
_RESEED_COUNTER :: Word64
_RESEED_COUNTER :: Word64
_RESEED_COUNTER = (Word64
2 :: Word64) Word64 -> Word64 -> Word64
forall a b. (Num a, Integral b) => a -> b -> a
^ (Word64
48 :: Word64)

-- | A deterministic random bit generator (DRBG).
--
--   Create a DRBG with 'new', and then use and reuse it to generate
--   bytes as needed.
newtype DRBG s = DRBG (P.MutVar s DRBGState)

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

-- DRBG environment data and state
data DRBGState = DRBGState
                 !HMACEnv       -- hmac function & outlen
                 !Word64        -- reseed counter
  {-# UNPACK #-} !BS.ByteString -- v
  {-# UNPACK #-} !BS.ByteString -- key

-- NB following synonym really only exists to make haddocks more
--    readable

-- | A HMAC function, taking a key as the first argument and the input
--   value as the second, producing a MAC digest.
--
--   >>> import qualified Crypto.Hash.SHA256 as SHA256
--   >>> :t SHA256.hmac
--   SHA256.hmac :: BS.ByteString -> BS.ByteString -> BS.ByteString
type HMAC = BS.ByteString -> BS.ByteString -> BS.ByteString

-- HMAC function and its associated outlength
data HMACEnv = HMACEnv
                 !HMAC
  {-# UNPACK #-} !Word64

-- the following convenience functions are useful 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 MutVar (PrimState m) DRBGState
mut) = do
  DRBGState HMACEnv
_ Word64
_ ByteString
v ByteString
_ <- MutVar (PrimState m) DRBGState -> m DRBGState
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
P.readMutVar MutVar (PrimState m) DRBGState
mut
  ByteString -> m ByteString
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
v

_read_k
  :: PrimMonad m
  => DRBG (PrimState m)
  -> m BS.ByteString
_read_k :: forall (m :: * -> *).
PrimMonad m =>
DRBG (PrimState m) -> m ByteString
_read_k (DRBG MutVar (PrimState m) DRBGState
mut) = do
  DRBGState HMACEnv
_ Word64
_ ByteString
_ ByteString
k <- MutVar (PrimState m) DRBGState -> m DRBGState
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
P.readMutVar MutVar (PrimState m) DRBGState
mut
  ByteString -> m ByteString
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
k

-- drbg interaction ------------------------------------------------------

-- | Create a DRBG from the supplied HMAC function, entropy, nonce, and
--   personalization string.
--
--   You can instantiate the DRBG using any appropriate HMAC function;
--   it should merely take a key and value as input, as is standard, and
--   return a MAC digest, each being a strict 'ByteString'.
--
--   The DRBG is returned in any 'PrimMonad', e.g. 'ST' or 'IO'.
--
--   >>> import qualified Crypto.Hash.SHA256 as SHA256
--   >>> new SHA256.hmac entropy nonce personalization_string
--   "<drbg>"
new
  :: PrimMonad m
  => HMAC           -- ^ HMAC function
  -> BS.ByteString  -- ^ entropy
  -> BS.ByteString  -- ^ nonce
  -> BS.ByteString  -- ^ personalization string
  -> m (DRBG (PrimState m))
new :: forall (m :: * -> *).
PrimMonad m =>
HMAC
-> ByteString -> ByteString -> ByteString -> m (DRBG (PrimState m))
new HMAC
hmac ByteString
entropy ByteString
nonce ByteString
ps = do
  let !drbg :: DRBGState
drbg = HMAC -> ByteString -> ByteString -> ByteString -> DRBGState
new_pure HMAC
hmac ByteString
entropy ByteString
nonce ByteString
ps
  MutVar (PrimState m) DRBGState
mut <- DRBGState -> m (MutVar (PrimState m) DRBGState)
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
P.newMutVar DRBGState
drbg
  DRBG (PrimState m) -> m (DRBG (PrimState m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MutVar (PrimState m) DRBGState -> DRBG (PrimState m)
forall s. MutVar s DRBGState -> DRBG s
DRBG MutVar (PrimState m) DRBGState
mut)

-- | 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
  => BS.ByteString        -- ^ entropy to inject
  -> BS.ByteString        -- ^ additional bytes to inject
  -> DRBG (PrimState m)
  -> m ()
reseed :: forall (m :: * -> *).
PrimMonad m =>
ByteString -> ByteString -> DRBG (PrimState m) -> m ()
reseed ByteString
ent ByteString
add (DRBG MutVar (PrimState m) DRBGState
drbg) = MutVar (PrimState m) DRBGState -> (DRBGState -> DRBGState) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> (a -> a) -> m ()
P.modifyMutVar' MutVar (PrimState m) DRBGState
drbg (ByteString -> ByteString -> DRBGState -> DRBGState
reseed_pure ByteString
ent ByteString
add)

-- | Generate bytes from a DRBG, optionally injecting additional bytes
--   per SP 800-90A.
--
--   >>> import qualified Data.ByteString.Base16 as B16
--   >>> drbg <- new SHA256.hmac entropy nonce personalization_string
--   >>> bytes0 <- gen addl_bytes 16 drbg
--   >>> bytes1 <- gen addl_bytes 16 drbg
--   >>> B16.encode bytes0
--   "938d6ca6d0b797f7b3c653349d6e3135"
--   >>> B16.encode bytes1
--   "5f379d16de6f2c6f8a35c56f13f9e5a5"
gen
  :: PrimMonad m
  => BS.ByteString       -- ^ additional bytes to inject
  -> Word64              -- ^ number of bytes to generate
  -> DRBG (PrimState m)
  -> m BS.ByteString
gen :: forall (m :: * -> *).
PrimMonad m =>
ByteString -> Word64 -> DRBG (PrimState m) -> m ByteString
gen ByteString
addl Word64
bytes (DRBG MutVar (PrimState m) DRBGState
mut) = do
  DRBGState
drbg0 <- MutVar (PrimState m) DRBGState -> m DRBGState
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
P.readMutVar MutVar (PrimState m) DRBGState
mut
  let !(Pair ByteString
bs DRBGState
drbg1) = ByteString -> Word64 -> DRBGState -> Pair ByteString DRBGState
gen_pure ByteString
addl Word64
bytes DRBGState
drbg0
  MutVar (PrimState m) DRBGState -> DRBGState -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> a -> m ()
P.writeMutVar MutVar (PrimState m) DRBGState
mut DRBGState
drbg1
  ByteString -> m ByteString
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
bs

-- pure drbg interaction ------------------------------------------------------

-- SP 800-90A 10.1.2.2
update_pure
  :: BS.ByteString
  -> DRBGState
  -> DRBGState
update_pure :: ByteString -> DRBGState -> DRBGState
update_pure ByteString
provided_data (DRBGState h :: HMACEnv
h@(HMACEnv HMAC
hmac Word64
_) Word64
r ByteString
v0 ByteString
k0) =
    let !k1 :: ByteString
k1 = HMAC
hmac ByteString
k0 (ByteString -> Word8 -> ByteString -> ByteString
cat ByteString
v0 Word8
0x00 ByteString
provided_data)
        !v1 :: ByteString
v1 = HMAC
hmac ByteString
k1 ByteString
v0
    in  if   ByteString -> Bool
BS.null ByteString
provided_data
        then HMACEnv -> Word64 -> ByteString -> ByteString -> DRBGState
DRBGState HMACEnv
h Word64
r ByteString
v1 ByteString
k1
        else let !k2 :: ByteString
k2 = HMAC
hmac ByteString
k1 (ByteString -> Word8 -> ByteString -> ByteString
cat ByteString
v1 Word8
0x01 ByteString
provided_data)
                 !v2 :: ByteString
v2 = HMAC
hmac ByteString
k2 ByteString
v1
             in  HMACEnv -> Word64 -> ByteString -> ByteString -> DRBGState
DRBGState HMACEnv
h Word64
r ByteString
v2 ByteString
k2
  where
    cat :: ByteString -> Word8 -> ByteString -> ByteString
cat ByteString
bs Word8
byte ByteString
suf = Builder -> ByteString
toStrict (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$
      ByteString -> Builder
BSB.byteString ByteString
bs Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BSB.word8 Word8
byte Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Builder
BSB.byteString ByteString
suf

-- SP 800-90A 10.1.2.3
new_pure
  :: HMAC           -- HMAC function
  -> BS.ByteString  -- entropy
  -> BS.ByteString  -- nonce
  -> BS.ByteString  -- personalization string
  -> DRBGState
new_pure :: HMAC -> ByteString -> ByteString -> ByteString -> DRBGState
new_pure HMAC
hmac ByteString
entropy ByteString
nonce ByteString
ps =
    let !drbg :: DRBGState
drbg = HMACEnv -> Word64 -> ByteString -> ByteString -> DRBGState
DRBGState (HMAC -> Word64 -> HMACEnv
HMACEnv HMAC
hmac Word64
outlen) Word64
1 ByteString
v0 ByteString
k0
    in  ByteString -> DRBGState -> DRBGState
update_pure ByteString
seed_material DRBGState
drbg
  where
    seed_material :: ByteString
seed_material = ByteString
entropy HMAC
forall a. Semigroup a => a -> a -> a
<> ByteString
nonce HMAC
forall a. Semigroup a => a -> a -> a
<> ByteString
ps
    outlen :: Word64
outlen = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fi (ByteString -> Int
BS.length (HMAC
hmac ByteString
forall a. Monoid a => a
mempty ByteString
forall a. Monoid a => a
mempty))
    k0 :: ByteString
k0 = Int -> Word8 -> ByteString
BS.replicate (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fi Word64
outlen) Word8
0x00
    v0 :: ByteString
v0 = Int -> Word8 -> ByteString
BS.replicate (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fi Word64
outlen) Word8
0x01

-- SP 800-90A 10.1.2.4
reseed_pure :: BS.ByteString -> BS.ByteString -> DRBGState -> DRBGState
reseed_pure :: ByteString -> ByteString -> DRBGState -> DRBGState
reseed_pure ByteString
entropy ByteString
addl DRBGState
drbg =
  let !(DRBGState HMACEnv
h Word64
_ ByteString
v ByteString
k) = ByteString -> DRBGState -> DRBGState
update_pure (ByteString
entropy HMAC
forall a. Semigroup a => a -> a -> a
<> ByteString
addl) DRBGState
drbg
  in  HMACEnv -> Word64 -> ByteString -> ByteString -> DRBGState
DRBGState HMACEnv
h Word64
1 ByteString
v ByteString
k

-- SP 800-90A 10.1.2.5
gen_pure
  :: BS.ByteString
  -> Word64
  -> DRBGState
  -> Pair BS.ByteString DRBGState
gen_pure :: ByteString -> Word64 -> DRBGState -> Pair ByteString DRBGState
gen_pure ByteString
addl Word64
bytes drbg0 :: DRBGState
drbg0@(DRBGState h :: HMACEnv
h@(HMACEnv HMAC
hmac Word64
outlen) Word64
_ ByteString
_ ByteString
_)
    | Word64
r Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
_RESEED_COUNTER = String -> Pair ByteString DRBGState
forall a. HasCallStack => String -> a
error String
"ppad-hmac-drbg: reseed required"
    | Bool
otherwise =
        let !(Pair ByteString
temp DRBGState
drbg1) = Builder -> Word64 -> ByteString -> Pair ByteString DRBGState
loop Builder
forall a. Monoid a => a
mempty Word64
0 ByteString
v1
            returned_bits :: ByteString
returned_bits = Int -> ByteString -> ByteString
BS.take (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fi Word64
bytes) ByteString
temp
            drbg :: DRBGState
drbg = ByteString -> DRBGState -> DRBGState
update_pure ByteString
addl DRBGState
drbg1
        in  ByteString -> DRBGState -> Pair ByteString DRBGState
forall a b. a -> b -> Pair a b
Pair ByteString
returned_bits DRBGState
drbg
  where
    !(DRBGState HMACEnv
_ Word64
r ByteString
v1 ByteString
k1)
      | ByteString -> Bool
BS.null ByteString
addl = DRBGState
drbg0
      | Bool
otherwise = ByteString -> DRBGState -> DRBGState
update_pure ByteString
addl DRBGState
drbg0

    loop :: Builder -> Word64 -> ByteString -> Pair ByteString DRBGState
loop !Builder
acc !Word64
len !ByteString
vl
      | Word64
len Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
bytes =
          let nv :: ByteString
nv   = HMAC
hmac ByteString
k1 ByteString
vl
              nacc :: Builder
nacc = Builder
acc Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Builder
BSB.byteString ByteString
nv
              nlen :: Word64
nlen = Word64
len Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
outlen
          in  Builder -> Word64 -> ByteString -> Pair ByteString DRBGState
loop Builder
nacc Word64
nlen ByteString
nv

      | Bool
otherwise =
          let facc :: ByteString
facc = Builder -> ByteString
toStrict Builder
acc
          in  ByteString -> DRBGState -> Pair ByteString DRBGState
forall a b. a -> b -> Pair a b
Pair ByteString
facc (HMACEnv -> Word64 -> ByteString -> ByteString -> DRBGState
DRBGState HMACEnv
h (Word64 -> Word64
forall a. Enum a => a -> a
succ Word64
r) ByteString
vl ByteString
k1)