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

-- |
-- 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
  , HMAC
  , Error(..)
  , _read_v
  , _read_k

  -- * 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.ByteString.Builder.Extra as BE
import qualified Data.ByteString.Internal as BI
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 #-}

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

to_strict_small :: BSB.Builder -> BS.ByteString
to_strict_small :: Builder -> ByteString
to_strict_small = LazyByteString -> ByteString
BS.toStrict (LazyByteString -> ByteString)
-> (Builder -> LazyByteString) -> Builder -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AllocationStrategy -> LazyByteString -> Builder -> LazyByteString
BE.toLazyByteStringWith
  (Int -> Int -> AllocationStrategy
BE.safeStrategy Int
128 Int
BE.smallChunkSize) LazyByteString
forall a. Monoid a => a
mempty
{-# INLINE to_strict_small #-}

-- 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 ----------------------------------------------------------------------

-- | A DRBG error.
data Error =
    MaxBytesExceeded -- ^ More than 65536 bytes have been requested.
  | ReseedRequired   -- ^ The DRBG must be reseeded (via 'reseed').
  deriving (Error -> Error -> Bool
(Error -> Error -> Bool) -> (Error -> Error -> Bool) -> Eq Error
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Error -> Error -> Bool
== :: Error -> Error -> Bool
$c/= :: Error -> Error -> Bool
/= :: Error -> Error -> Bool
Eq, Int -> Error -> ShowS
[Error] -> ShowS
Error -> String
(Int -> Error -> ShowS)
-> (Error -> String) -> ([Error] -> ShowS) -> Show Error
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Error -> ShowS
showsPrec :: Int -> Error -> ShowS
$cshow :: Error -> String
show :: Error -> String
$cshowList :: [Error] -> ShowS
showList :: [Error] -> ShowS
Show)

-- 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.
--
--   >>> drbg <- new hmac entropy nonce personalization_string
--   >>> bytes0 <- gen addl_bytes 16 drbg
--   >>> bytes1 <- gen addl_bytes 16 drbg
--   >>> drbg
--   "<drbg>"
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
  {-# UNPACK #-} !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
--   >>> let hmac k b = let SHA256.MAC m = SHA256.hmac k b in m
--   >>> :t hmac
--   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 _ _ v _ <- MutVar (PrimState m) DRBGState -> m DRBGState
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
P.readMutVar MutVar (PrimState m) DRBGState
mut
  pure 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 _ _ _ 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
  pure 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
--   >>> let hmac k b = let SHA256.MAC m = SHA256.hmac k b in m
--   >>> new 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
  mut <- DRBGState -> m (MutVar (PrimState m) DRBGState)
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
P.newMutVar DRBGState
drbg
  pure (DRBG mut)

-- | 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 hmac entropy nonce personalization_string
--   >>> Right bytes0 <- gen addl_bytes 16 drbg
--   >>> Right 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 (Either Error BS.ByteString)
gen :: forall (m :: * -> *).
PrimMonad m =>
ByteString
-> Word64 -> DRBG (PrimState m) -> m (Either Error ByteString)
gen ByteString
addl Word64
bytes (DRBG MutVar (PrimState m) DRBGState
mut) = do
  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
  case gen_pure addl bytes drbg0 of
    Left Error
e -> Either Error ByteString -> m (Either Error ByteString)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Error -> Either Error ByteString
forall a b. a -> Either a b
Left Error
e)
    Right !(Pair ByteString
bs DRBGState
drbg1) -> do
      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
      Either Error ByteString -> m (Either Error ByteString)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Either Error ByteString
forall a b. b -> Either a b
Right ByteString
bs)

-- | 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)

-- 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 suf :: ByteString
suf@(BI.PS ForeignPtr Word8
_ Int
_ Int
l) =
      let bil :: Builder
bil = 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
      in  if   Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
64
          then Builder -> ByteString
to_strict_small Builder
bil
          else Builder -> ByteString
to_strict Builder
bil
    {-# INLINE cat #-}

-- 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
  -> Either Error (Pair BS.ByteString DRBGState)
gen_pure :: ByteString
-> Word64 -> DRBGState -> Either Error (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
bytes Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
0x10000     = Error -> Either Error (Pair ByteString DRBGState)
forall a b. a -> Either a b
Left Error
MaxBytesExceeded
    | Word64
r Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
_RESEED_COUNTER = Error -> Either Error (Pair ByteString DRBGState)
forall a b. a -> Either a b
Left Error
ReseedRequired
    | 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  Pair ByteString DRBGState
-> Either Error (Pair ByteString DRBGState)
forall a b. b -> Either a b
Right (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 | Word64
bytes Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
128 = Builder -> ByteString
to_strict_small Builder
acc
                   | Bool
otherwise   = Builder -> ByteString
to_strict 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)
{-# INLINE gen_pure #-}