-- |
-- Module      : Crypto.Random.EntropyPool
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : experimental
-- Portability : Good
--
module Crypto.Random.EntropyPool
    ( EntropyPool
    , createEntropyPool
    , createEntropyPoolWith
    , getEntropyFrom
    ) where

import           Control.Concurrent.MVar
import           Crypto.Random.Entropy.Unsafe
import           Crypto.Internal.ByteArray (ByteArray, ScrubbedBytes)
import qualified Crypto.Internal.ByteArray as B
import           Data.Word (Word8)
import           Data.Maybe (catMaybes)
import           Foreign.Marshal.Utils (copyBytes)
import           Foreign.Ptr (plusPtr, Ptr)

-- | Pool of Entropy. Contains a self-mutating pool of entropy,
-- that is always guaranteed to contain data.
data EntropyPool = EntropyPool [EntropyBackend] (MVar Int) ScrubbedBytes

-- size of entropy pool by default
defaultPoolSize :: Int
defaultPoolSize :: Int
defaultPoolSize = Int
4096

-- | Create a new entropy pool of a specific size
--
-- While you can create as many entropy pools as you want,
-- the pool can be shared between multiples RNGs.
createEntropyPoolWith :: Int -> [EntropyBackend] -> IO EntropyPool
createEntropyPoolWith :: Int -> [EntropyBackend] -> IO EntropyPool
createEntropyPoolWith Int
poolSize [EntropyBackend]
backends = do
    MVar Int
m  <- Int -> IO (MVar Int)
forall a. a -> IO (MVar a)
newMVar Int
0
    ScrubbedBytes
sm <- Int -> (Ptr Word8 -> IO ()) -> IO ScrubbedBytes
forall ba p. ByteArray ba => Int -> (Ptr p -> IO ()) -> IO ba
B.alloc Int
poolSize (Int -> [EntropyBackend] -> Ptr Word8 -> IO ()
replenish Int
poolSize [EntropyBackend]
backends)
    EntropyPool -> IO EntropyPool
forall (m :: * -> *) a. Monad m => a -> m a
return (EntropyPool -> IO EntropyPool) -> EntropyPool -> IO EntropyPool
forall a b. (a -> b) -> a -> b
$ [EntropyBackend] -> MVar Int -> ScrubbedBytes -> EntropyPool
EntropyPool [EntropyBackend]
backends MVar Int
m ScrubbedBytes
sm

-- | Create a new entropy pool with a default size.
--
-- While you can create as many entropy pools as you want,
-- the pool can be shared between multiples RNGs.
createEntropyPool :: IO EntropyPool
createEntropyPool :: IO EntropyPool
createEntropyPool = do
    [EntropyBackend]
backends <- [Maybe EntropyBackend] -> [EntropyBackend]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe EntropyBackend] -> [EntropyBackend])
-> IO [Maybe EntropyBackend] -> IO [EntropyBackend]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` [IO (Maybe EntropyBackend)] -> IO [Maybe EntropyBackend]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [IO (Maybe EntropyBackend)]
supportedBackends
    Int -> [EntropyBackend] -> IO EntropyPool
createEntropyPoolWith Int
defaultPoolSize [EntropyBackend]
backends

-- | Put a chunk of the entropy pool into a buffer
getEntropyPtr :: EntropyPool -> Int -> Ptr Word8 -> IO ()
getEntropyPtr :: EntropyPool -> Int -> Ptr Word8 -> IO ()
getEntropyPtr (EntropyPool [EntropyBackend]
backends MVar Int
posM ScrubbedBytes
sm) Int
n Ptr Word8
outPtr =
    ScrubbedBytes -> (Ptr Word8 -> IO ()) -> IO ()
forall ba p a. ByteArrayAccess ba => ba -> (Ptr p -> IO a) -> IO a
B.withByteArray ScrubbedBytes
sm ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
entropyPoolPtr ->
        MVar Int -> (Int -> IO Int) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar Int
posM ((Int -> IO Int) -> IO ()) -> (Int -> IO Int) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
pos ->
            Ptr Word8 -> Ptr Word8 -> Int -> Int -> IO Int
forall b. Ptr b -> Ptr Word8 -> Int -> Int -> IO Int
copyLoop Ptr Word8
outPtr Ptr Word8
entropyPoolPtr Int
pos Int
n
  where poolSize :: Int
poolSize = ScrubbedBytes -> Int
forall ba. ByteArrayAccess ba => ba -> Int
B.length ScrubbedBytes
sm
        copyLoop :: Ptr b -> Ptr Word8 -> Int -> Int -> IO Int
copyLoop Ptr b
d Ptr Word8
s Int
pos Int
left
            | Int
left Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
pos
            | Bool
otherwise = do
                Int
wrappedPos <-
                    if Int
pos Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
poolSize
                        then Int -> [EntropyBackend] -> Ptr Word8 -> IO ()
replenish Int
poolSize [EntropyBackend]
backends Ptr Word8
s IO () -> IO Int -> IO Int
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0
                        else Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
pos
                let m :: Int
m = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
poolSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
wrappedPos) Int
left
                Ptr b -> Ptr b -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr b
d (Ptr Word8
s Ptr Word8 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
wrappedPos) Int
m
                Ptr b -> Ptr Word8 -> Int -> Int -> IO Int
copyLoop (Ptr b
d Ptr b -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
m) Ptr Word8
s (Int
wrappedPos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
m) (Int
left Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
m)

-- | Grab a chunk of entropy from the entropy pool.
getEntropyFrom :: ByteArray byteArray => EntropyPool -> Int -> IO byteArray
getEntropyFrom :: EntropyPool -> Int -> IO byteArray
getEntropyFrom EntropyPool
pool Int
n = Int -> (Ptr Word8 -> IO ()) -> IO byteArray
forall ba p. ByteArray ba => Int -> (Ptr p -> IO ()) -> IO ba
B.alloc Int
n (EntropyPool -> Int -> Ptr Word8 -> IO ()
getEntropyPtr EntropyPool
pool Int
n)