{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE PackageImports #-}
---------------------------------------------------------
--
-- |
--
-- Module        : Web.ClientSession
-- Copyright     : Michael Snoyman
-- License       : BSD3
--
-- Maintainer    : Michael Snoyman <michael@snoyman.com>
-- Stability     : Stable
-- Portability   : portable
--
-- Stores session data in a client cookie.  In order to do so,
-- we:
--
-- * Encrypt the cookie data using AES in CTR mode.  This allows
-- you to store sensitive information on the client side without
-- worrying about eavesdropping.
--
-- * Authenticate the encrypted cookie data using
-- Skein-MAC-512-256.  Besides detecting potential errors in
-- storage or transmission of the cookies (integrity), the MAC
-- also avoids malicious modifications of the cookie data by
-- assuring you that the cookie data really was generated by this
-- server (authenticity).
--
-- * Encode everything using Base64.  Thus we avoid problems with
-- non-printable characters by giving the browser a simple
-- string.
--
-- Simple usage of the library involves just calling
-- 'getDefaultKey' on the startup of your server, 'encryptIO'
-- when serializing cookies and 'decrypt' when parsing then back.
--
---------------------------------------------------------
module Web.ClientSession
    ( -- * Automatic key generation
      Key
    , IV
    , randomIV
    , mkIV
    , getKey
    , getKeyEnv
    , defaultKeyFile
    , getDefaultKey
    , initKey
    , randomKey
    , randomKeyEnv
      -- * Actual encryption/decryption
    , encrypt
    , encryptIO
    , decrypt
    ) where

-- from base
import Control.Applicative ((<$>))
import Control.Concurrent (forkIO)
import Control.Monad (guard, when)
import Data.Function (on)

#if MIN_VERSION_base(4,7,0)
import System.Environment (lookupEnv, setEnv)
#elif MIN_VERSION_base(4,6,0)
import System.Environment (lookupEnv)
import System.SetEnv (setEnv)
#else
import System.LookupEnv (lookupEnv)
import System.SetEnv (setEnv)
#endif

import System.IO.Unsafe (unsafePerformIO)
import qualified Data.IORef as I

-- from directory
import System.Directory (doesFileExist)

-- from bytestring
import qualified Data.ByteString as S
import qualified Data.ByteString.Char8 as C
import qualified Data.ByteString.Base64 as B

-- from cereal
import Data.Serialize (encode, Serialize (put, get), getBytes, putByteString)

-- from tagged
import Data.Tagged (Tagged, untag)

-- from crypto-api
import Crypto.Classes (constTimeEq)
import "crypto-api" Crypto.Random (genSeedLength, reseed)
import Crypto.Types (ByteLength)

-- from cipher-aes
import qualified Crypto.Cipher.AES as A

-- from skein
import Crypto.Skein (skeinMAC', Skein_512_256)

-- from entropy
import System.Entropy (getEntropy)

-- from cprng-aes
#if MIN_VERSION_cprng_aes(0,5,0)
import Crypto.Random.AESCtr (AESRNG, makeSystem)
import "crypto-random" Crypto.Random (cprgGenerate)
#else
import Crypto.Random.AESCtr (AESRNG, makeSystem, genRandomBytes)
#endif


-- | The keys used to store the cookies.  We have an AES key used
-- to encrypt the cookie and a Skein-MAC-512-256 key used verify
-- the authencity and integrity of the cookie.  The AES key must
-- have exactly 32 bytes (256 bits) while Skein-MAC-512-256 must
-- have 64 bytes (512 bits).
--
-- See also 'getDefaultKey' and 'initKey'.
data Key = Key { Key -> AES
aesKey ::
#if MIN_VERSION_cipher_aes(0, 2, 0)
                    !A.AES
#else
                    !A.Key
#endif
                 -- ^ AES key with 32 bytes.
               , Key -> ByteString -> Skein_512_256
macKey :: !(S.ByteString -> Skein_512_256)
                 -- ^ Skein-MAC key.  Instead of storing the key
                 -- data, we store a partially applied function
                 -- for calculating the MAC (see 'skeinMAC'').
               , Key -> ByteString
keyRaw :: !S.ByteString
               }

instance Eq Key where
    Key _ _ r1 :: ByteString
r1 == :: Key -> Key -> Bool
== Key _ _ r2 :: ByteString
r2 = ByteString
r1 ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
r2

instance Serialize Key where
    put :: Putter Key
put = Putter ByteString
putByteString Putter ByteString -> (Key -> ByteString) -> Putter Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key -> ByteString
keyRaw
    get :: Get Key
get = ([Char] -> Key) -> (Key -> Key) -> Either [Char] Key -> Key
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either [Char] -> Key
forall a. HasCallStack => [Char] -> a
error Key -> Key
forall a. a -> a
id (Either [Char] Key -> Key)
-> (ByteString -> Either [Char] Key) -> ByteString -> Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Either [Char] Key
initKey (ByteString -> Key) -> Get ByteString -> Get Key
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Get ByteString
getBytes 96

-- | Dummy 'Show' instance.
instance Show Key where
    show :: Key -> [Char]
show _ = "<Web.ClientSession.Key>"

-- | The initialization vector used by AES.  Must be exactly 16
-- bytes long.
newtype IV = IV S.ByteString

unsafeMkIV :: S.ByteString -> IV
unsafeMkIV :: ByteString -> IV
unsafeMkIV bs :: ByteString
bs = (ByteString -> IV
IV ByteString
bs)

unIV :: IV -> S.ByteString
unIV :: IV -> ByteString
unIV (IV bs :: ByteString
bs) = ByteString
bs

instance Eq IV where
  == :: IV -> IV -> Bool
(==) = ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
(==) (ByteString -> ByteString -> Bool)
-> (IV -> ByteString) -> IV -> IV -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IV -> ByteString
unIV
  /= :: IV -> IV -> Bool
(/=) = ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
(/=) (ByteString -> ByteString -> Bool)
-> (IV -> ByteString) -> IV -> IV -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IV -> ByteString
unIV

instance Ord IV where
  compare :: IV -> IV -> Ordering
compare = ByteString -> ByteString -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (ByteString -> ByteString -> Ordering)
-> (IV -> ByteString) -> IV -> IV -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IV -> ByteString
unIV
  <= :: IV -> IV -> Bool
(<=) = ByteString -> ByteString -> Bool
forall a. Ord a => a -> a -> Bool
(<=) (ByteString -> ByteString -> Bool)
-> (IV -> ByteString) -> IV -> IV -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IV -> ByteString
unIV
  < :: IV -> IV -> Bool
(<)  = ByteString -> ByteString -> Bool
forall a. Ord a => a -> a -> Bool
(<)  (ByteString -> ByteString -> Bool)
-> (IV -> ByteString) -> IV -> IV -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IV -> ByteString
unIV
  >= :: IV -> IV -> Bool
(>=) = ByteString -> ByteString -> Bool
forall a. Ord a => a -> a -> Bool
(>=) (ByteString -> ByteString -> Bool)
-> (IV -> ByteString) -> IV -> IV -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IV -> ByteString
unIV
  > :: IV -> IV -> Bool
(>)  = ByteString -> ByteString -> Bool
forall a. Ord a => a -> a -> Bool
(>)  (ByteString -> ByteString -> Bool)
-> (IV -> ByteString) -> IV -> IV -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IV -> ByteString
unIV

instance Show IV where
  show :: IV -> [Char]
show = ByteString -> [Char]
forall a. Show a => a -> [Char]
show (ByteString -> [Char]) -> (IV -> ByteString) -> IV -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IV -> ByteString
unIV

instance Serialize IV where
  put :: Putter IV
put = Putter ByteString
forall t. Serialize t => Putter t
put Putter ByteString -> (IV -> ByteString) -> Putter IV
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IV -> ByteString
unIV
  get :: Get IV
get = ByteString -> IV
unsafeMkIV (ByteString -> IV) -> Get ByteString -> Get IV
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get ByteString
forall t. Serialize t => Get t
get

-- | Construct an initialization vector from a 'S.ByteString'.
-- Fails if there isn't exactly 16 bytes.
mkIV :: S.ByteString -> Maybe IV
mkIV :: ByteString -> Maybe IV
mkIV bs :: ByteString
bs | ByteString -> Int
S.length ByteString
bs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 16 = IV -> Maybe IV
forall a. a -> Maybe a
Just (ByteString -> IV
unsafeMkIV ByteString
bs)
        | Bool
otherwise         = Maybe IV
forall a. Maybe a
Nothing

-- | Randomly construct a fresh initialization vector.  You
-- /MUST NOT/ reuse initialization vectors.
randomIV :: IO IV
randomIV :: IO IV
randomIV = IO IV
aesRNG

-- | The default key file.
defaultKeyFile :: FilePath
defaultKeyFile :: [Char]
defaultKeyFile = "client_session_key.aes"

-- | Simply calls 'getKey' 'defaultKeyFile'.
getDefaultKey :: IO Key
getDefaultKey :: IO Key
getDefaultKey = [Char] -> IO Key
getKey [Char]
defaultKeyFile

-- | Get a key from the given text file.
--
-- If the file does not exist or is corrupted a random key will
-- be generated and stored in that file.
getKey :: FilePath     -- ^ File name where key is stored.
       -> IO Key       -- ^ The actual key.
getKey :: [Char] -> IO Key
getKey keyFile :: [Char]
keyFile = do
    Bool
exists <- [Char] -> IO Bool
doesFileExist [Char]
keyFile
    if Bool
exists
        then [Char] -> IO ByteString
S.readFile [Char]
keyFile IO ByteString -> (ByteString -> IO Key) -> IO Key
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ([Char] -> IO Key)
-> (Key -> IO Key) -> Either [Char] Key -> IO Key
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (IO Key -> [Char] -> IO Key
forall a b. a -> b -> a
const IO Key
newKey) Key -> IO Key
forall (m :: * -> *) a. Monad m => a -> m a
return (Either [Char] Key -> IO Key)
-> (ByteString -> Either [Char] Key) -> ByteString -> IO Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Either [Char] Key
initKey
        else IO Key
newKey
  where
    newKey :: IO Key
newKey = do
        (bs :: ByteString
bs, key' :: Key
key') <- IO (ByteString, Key)
randomKey
        [Char] -> ByteString -> IO ()
S.writeFile [Char]
keyFile ByteString
bs
        Key -> IO Key
forall (m :: * -> *) a. Monad m => a -> m a
return Key
key'

-- | Get the key from the named environment variable
--
-- Assumes the value is a Base64-encoded string. If the variable is not set, a
-- random key will be generated, set in the environment, and the Base64-encoded
-- version printed on @/dev/stdout@.
getKeyEnv :: String     -- ^ Name of the environment variable
          -> IO Key     -- ^ The actual key.
getKeyEnv :: [Char] -> IO Key
getKeyEnv envVar :: [Char]
envVar = do
    Maybe [Char]
mvalue <- [Char] -> IO (Maybe [Char])
lookupEnv [Char]
envVar
    case Maybe [Char]
mvalue of
        Just value :: [Char]
value -> ([Char] -> IO Key)
-> (Key -> IO Key) -> Either [Char] Key -> IO Key
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (IO Key -> [Char] -> IO Key
forall a b. a -> b -> a
const IO Key
newKey) Key -> IO Key
forall (m :: * -> *) a. Monad m => a -> m a
return (Either [Char] Key -> IO Key) -> Either [Char] Key -> IO Key
forall a b. (a -> b) -> a -> b
$ ByteString -> Either [Char] Key
initKey (ByteString -> Either [Char] Key)
-> Either [Char] ByteString -> Either [Char] Key
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Either [Char] ByteString
decode [Char]
value
        Nothing -> IO Key
newKey
  where
    decode :: [Char] -> Either [Char] ByteString
decode = ByteString -> Either [Char] ByteString
B.decode (ByteString -> Either [Char] ByteString)
-> ([Char] -> ByteString) -> [Char] -> Either [Char] ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> ByteString
C.pack
    newKey :: IO Key
newKey = [Char] -> IO Key
randomKeyEnv [Char]
envVar

-- | Generate a random 'Key'.  Besides the 'Key', the
-- 'ByteString' passed to 'initKey' is returned so that it can be
-- saved for later use.
randomKey :: IO (S.ByteString, Key)
randomKey :: IO (ByteString, Key)
randomKey = do
    ByteString
bs <- Int -> IO ByteString
getEntropy 96
    case ByteString -> Either [Char] Key
initKey ByteString
bs of
        Left e :: [Char]
e -> [Char] -> IO (ByteString, Key)
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO (ByteString, Key)) -> [Char] -> IO (ByteString, Key)
forall a b. (a -> b) -> a -> b
$ "Web.ClientSession.randomKey: never here, " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
e
        Right key :: Key
key -> (ByteString, Key) -> IO (ByteString, Key)
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
bs, Key
key)

-- | Generate a random 'Key', set a Base64-encoded version of it in the given
-- environment variable, then return it. Also prints the generated string to
-- @/dev/stdout@.
randomKeyEnv :: String -> IO Key
randomKeyEnv :: [Char] -> IO Key
randomKeyEnv envVar :: [Char]
envVar = do
    (bs :: ByteString
bs, key :: Key
key) <- IO (ByteString, Key)
randomKey
    let encoded :: [Char]
encoded = ByteString -> [Char]
C.unpack (ByteString -> [Char]) -> ByteString -> [Char]
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
B.encode ByteString
bs
    [Char] -> [Char] -> IO ()
setEnv [Char]
envVar [Char]
encoded
    [Char] -> IO ()
putStrLn ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char]
envVar [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ "=" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
encoded
    Key -> IO Key
forall (m :: * -> *) a. Monad m => a -> m a
return Key
key

-- | Initializes a 'Key' from a random 'S.ByteString'.  Fails if
-- there isn't exactly 96 bytes (256 bits for AES and 512 bits
-- for Skein-MAC-512-512).
--
-- Note that the input string is assumed to be uniformly chosen
-- from the set of all 96-byte strings.  In other words, each
-- byte should be chosen from the set of all byte values (0-255)
-- with the same probability.
--
-- In particular, this function does not do any kind of key
-- stretching.  You should never feed it a password, for example.
--
-- It's /highly/ recommended to feed @initKey@ only with values
-- generated by 'randomKey', unless you really know what you're
-- doing.
initKey :: S.ByteString -> Either String Key
initKey :: ByteString -> Either [Char] Key
initKey bs :: ByteString
bs | ByteString -> Int
S.length ByteString
bs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= 96 = [Char] -> Either [Char] Key
forall a b. a -> Either a b
Left ([Char] -> Either [Char] Key) -> [Char] -> Either [Char] Key
forall a b. (a -> b) -> a -> b
$ "Web.ClientSession.initKey: length of " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++
                                         Int -> [Char]
forall a. Show a => a -> [Char]
show (ByteString -> Int
S.length ByteString
bs) [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ " /= 96."
initKey bs :: ByteString
bs = Key -> Either [Char] Key
forall a b. b -> Either a b
Right (Key -> Either [Char] Key) -> Key -> Either [Char] Key
forall a b. (a -> b) -> a -> b
$ $WKey :: AES -> (ByteString -> Skein_512_256) -> ByteString -> Key
Key { aesKey :: AES
aesKey = ByteString -> AES
forall b. Byteable b => b -> AES
A.initKey ByteString
preAesKey
                         , macKey :: ByteString -> Skein_512_256
macKey = ByteString -> ByteString -> Skein_512_256
forall skeinCtx digest.
(SkeinMAC skeinCtx, Hash skeinCtx digest) =>
ByteString -> ByteString -> digest
skeinMAC' ByteString
preMacKey
                         , keyRaw :: ByteString
keyRaw = ByteString
bs
                         }
    where
      (preMacKey :: ByteString
preMacKey, preAesKey :: ByteString
preAesKey) = Int -> ByteString -> (ByteString, ByteString)
S.splitAt 64 ByteString
bs

-- | Same as 'encrypt', however randomly generates the
-- initialization vector for you.
encryptIO :: Key -> S.ByteString -> IO S.ByteString
encryptIO :: Key -> ByteString -> IO ByteString
encryptIO key :: Key
key x :: ByteString
x = do
    IV
iv <- IO IV
randomIV
    ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Key -> IV -> ByteString -> ByteString
encrypt Key
key IV
iv ByteString
x

-- | Encrypt (AES-CTR), authenticate (Skein-MAC-512-256) and
-- encode (Base64) the given cookie data.  The returned byte
-- string is ready to be used in a response header.
encrypt :: Key          -- ^ Key of the server.
        -> IV           -- ^ New, random initialization vector (see 'randomIV').
        -> S.ByteString -- ^ Serialized cookie data.
        -> S.ByteString -- ^ Encoded cookie data to be given to
                        -- the client browser.
encrypt :: Key -> IV -> ByteString -> ByteString
encrypt key :: Key
key (IV iv :: ByteString
iv) x :: ByteString
x = ByteString -> ByteString
B.encode ByteString
final
  where
#if MIN_VERSION_cipher_aes(0, 2, 0)
    encrypted :: ByteString
encrypted  = AES -> ByteString -> ByteString -> ByteString
forall iv. Byteable iv => AES -> iv -> ByteString -> ByteString
A.encryptCTR (Key -> AES
aesKey Key
key) ByteString
iv ByteString
x
#else
    encrypted  = A.encryptCTR (aesKey key) (A.IV iv) x
#endif
    toBeAuthed :: ByteString
toBeAuthed = ByteString
iv ByteString -> ByteString -> ByteString
`S.append` ByteString
encrypted
    auth :: Skein_512_256
auth       = Key -> ByteString -> Skein_512_256
macKey Key
key ByteString
toBeAuthed
    final :: ByteString
final      = Skein_512_256 -> ByteString
forall a. Serialize a => a -> ByteString
encode Skein_512_256
auth ByteString -> ByteString -> ByteString
`S.append` ByteString
toBeAuthed

-- | Decode (Base64), verify the integrity and authenticity
-- (Skein-MAC-512-256) and decrypt (AES-CTR) the given encoded
-- cookie data.  Returns the original serialized cookie data.
-- Fails if the data is corrupted.
decrypt :: Key                -- ^ Key of the server.
        -> S.ByteString       -- ^ Encoded cookie data given by the browser.
        -> Maybe S.ByteString -- ^ Serialized cookie data.
decrypt :: Key -> ByteString -> Maybe ByteString
decrypt key :: Key
key dataBS64 :: ByteString
dataBS64 = do
    ByteString
dataBS <- ([Char] -> Maybe ByteString)
-> (ByteString -> Maybe ByteString)
-> Either [Char] ByteString
-> Maybe ByteString
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe ByteString -> [Char] -> Maybe ByteString
forall a b. a -> b -> a
const Maybe ByteString
forall a. Maybe a
Nothing) ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (Either [Char] ByteString -> Maybe ByteString)
-> Either [Char] ByteString -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> Either [Char] ByteString
B.decode ByteString
dataBS64
    Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (ByteString -> Int
S.length ByteString
dataBS Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= 48) -- 16 bytes of IV + 32 bytes of Skein-MAC-512-256
    let (auth :: ByteString
auth, toBeAuthed :: ByteString
toBeAuthed) = Int -> ByteString -> (ByteString, ByteString)
S.splitAt 32 ByteString
dataBS
        auth' :: Skein_512_256
auth' = Key -> ByteString -> Skein_512_256
macKey Key
key ByteString
toBeAuthed
    Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Skein_512_256 -> ByteString
forall a. Serialize a => a -> ByteString
encode Skein_512_256
auth' ByteString -> ByteString -> Bool
`constTimeEq` ByteString
auth)
    let (iv :: ByteString
iv, encrypted :: ByteString
encrypted) = Int -> ByteString -> (ByteString, ByteString)
S.splitAt 16 ByteString
toBeAuthed
#if MIN_VERSION_cipher_aes(0, 2, 0)
    let iv' :: ByteString
iv' = ByteString
iv
#else
    let iv' = A.IV iv
#endif
    ByteString -> Maybe ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> Maybe ByteString) -> ByteString -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$! AES -> ByteString -> ByteString -> ByteString
forall iv. Byteable iv => AES -> iv -> ByteString -> ByteString
A.decryptCTR (Key -> AES
aesKey Key
key) ByteString
iv' ByteString
encrypted


-- Significantly more efficient random IV generation. Initial
-- benchmarks placed it at 6.06 us versus 1.69 ms for
-- Crypto.Modes.getIVIO, since it does not require /dev/urandom
-- I/O for every call.

data AESState =
    ASt {-# UNPACK #-} !AESRNG -- Our CPRNG using AES on CTR mode
        {-# UNPACK #-} !Int    -- How many IVs were generated with this
                               -- AESRNG.  Used to control reseeding.

-- | Construct initial state of the CPRNG.
aesSeed :: IO AESState
aesSeed :: IO AESState
aesSeed = do
  AESRNG
rng <- IO AESRNG
makeSystem
  AESState -> IO AESState
forall (m :: * -> *) a. Monad m => a -> m a
return (AESState -> IO AESState) -> AESState -> IO AESState
forall a b. (a -> b) -> a -> b
$! AESRNG -> Int -> AESState
ASt AESRNG
rng 0

-- | Reseed the CPRNG with new entropy from the system pool.
aesReseed :: IO ()
aesReseed :: IO ()
aesReseed = do
  AESRNG
rng' <- IO AESRNG
makeSystem
  IORef AESState -> AESState -> IO ()
forall a. IORef a -> a -> IO ()
I.writeIORef IORef AESState
aesRef (AESState -> IO ()) -> AESState -> IO ()
forall a b. (a -> b) -> a -> b
$ AESRNG -> Int -> AESState
ASt AESRNG
rng' 0

-- | 'IORef' that keeps the current state of the CPRNG.  Yep,
-- global state.  Used in thread-safe was only, though.
aesRef :: I.IORef AESState
aesRef :: IORef AESState
aesRef = IO (IORef AESState) -> IORef AESState
forall a. IO a -> a
unsafePerformIO (IO (IORef AESState) -> IORef AESState)
-> IO (IORef AESState) -> IORef AESState
forall a b. (a -> b) -> a -> b
$ IO AESState
aesSeed IO AESState
-> (AESState -> IO (IORef AESState)) -> IO (IORef AESState)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= AESState -> IO (IORef AESState)
forall a. a -> IO (IORef a)
I.newIORef
{-# NOINLINE aesRef #-}

-- | Construct a new 16-byte IV using our CPRNG.  Forks another
-- thread to reseed the CPRNG should its usage count reach a
-- hardcoded threshold.
aesRNG :: IO IV
aesRNG :: IO IV
aesRNG = do
  (bs :: ByteString
bs, count :: Int
count) <-
      IORef AESState
-> (AESState -> (AESState, (ByteString, Int)))
-> IO (ByteString, Int)
forall a b. IORef a -> (a -> (a, b)) -> IO b
I.atomicModifyIORef IORef AESState
aesRef ((AESState -> (AESState, (ByteString, Int)))
 -> IO (ByteString, Int))
-> (AESState -> (AESState, (ByteString, Int)))
-> IO (ByteString, Int)
forall a b. (a -> b) -> a -> b
$ \(ASt rng :: AESRNG
rng count :: Int
count) ->
#if MIN_VERSION_cprng_aes(0, 5, 0)
          let (bs' :: ByteString
bs', rng' :: AESRNG
rng') = Int -> AESRNG -> (ByteString, AESRNG)
forall gen. CPRG gen => Int -> gen -> (ByteString, gen)
cprgGenerate 16 AESRNG
rng
#elif MIN_VERSION_cprng_aes(0, 3, 2)
          let (bs', rng') = genRandomBytes 16 rng
#else
          let (bs', rng') = genRandomBytes rng 16
#endif
          in (AESRNG -> Int -> AESState
ASt AESRNG
rng' (Int -> Int
forall a. Enum a => a -> a
succ Int
count), (ByteString
bs', Int
count))
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
count Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
threshold) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO ThreadId -> IO ()
forall (m :: * -> *) a. Monad m => m a -> m ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkIO IO ()
aesReseed
  IV -> IO IV
forall (m :: * -> *) a. Monad m => a -> m a
return (IV -> IO IV) -> IV -> IO IV
forall a b. (a -> b) -> a -> b
$! ByteString -> IV
unsafeMkIV ByteString
bs
 where
  void :: m a -> m ()
void f :: m a
f = m a
f m a -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | How many IVs should be generated before reseeding the CPRNG.
-- This number depends basically on how paranoid you are.  We
-- think 100.000 is a good compromise: larger numbers give only a
-- small performance advantage, while it still is a small number
-- since we only generate 1.5 MiB of random data between reseeds.
threshold :: Int
threshold :: Int
threshold = 100000