-- |
-- Module      : Crypto.Random.Entropy
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : experimental
-- Portability : Good
--
{-# LANGUAGE CPP #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE DeriveDataTypeable #-}
module Crypto.Random.Entropy
    ( EntropyPool
    , createEntropyPool
    , createTestEntropyPool
    , grabEntropyPtr
    , grabEntropy
    , grabEntropyIO
    ) where

import Control.Monad (when)
import Control.Concurrent.MVar
import System.IO.Unsafe (unsafePerformIO)
import Data.Maybe (catMaybes)
import Data.SecureMem
import Data.Typeable (Typeable)
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import qualified Data.ByteString.Internal as B
import Data.Word (Word8)
import Foreign.Marshal.Utils (copyBytes)
import Foreign.Ptr (plusPtr, Ptr)
import Foreign.ForeignPtr (withForeignPtr)

import Crypto.Random.Entropy.Source
#ifdef SUPPORT_RDRAND
import Crypto.Random.Entropy.RDRand
#endif
#ifdef WINDOWS
import Crypto.Random.Entropy.Windows
#else
import Crypto.Random.Entropy.Unix
#endif

supportedBackends :: [IO (Maybe EntropyBackend)]
supportedBackends :: [IO (Maybe EntropyBackend)]
supportedBackends =
    [
#ifdef SUPPORT_RDRAND
    openBackend (undefined :: RDRand),
#endif
#ifdef WINDOWS
    openBackend (undefined :: WinCryptoAPI)
#else
    DevRandom -> IO (Maybe EntropyBackend)
forall b. EntropySource b => b -> IO (Maybe EntropyBackend)
openBackend (DevRandom
forall a. HasCallStack => a
undefined :: DevRandom), DevURandom -> IO (Maybe EntropyBackend)
forall b. EntropySource b => b -> IO (Maybe EntropyBackend)
openBackend (DevURandom
forall a. HasCallStack => a
undefined :: DevURandom)
#endif
    ]

data EntropyBackend = forall b . EntropySource b => EntropyBackend b

newtype TestEntropySource = TestEntropySource ByteString

instance EntropySource TestEntropySource where
    entropyOpen :: IO (Maybe TestEntropySource)
entropyOpen    = Maybe TestEntropySource -> IO (Maybe TestEntropySource)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe TestEntropySource
forall a. Maybe a
Nothing
    entropyGather :: TestEntropySource -> Ptr Word8 -> Int -> IO Int
entropyGather (TestEntropySource ByteString
bs) Ptr Word8
dst Int
n
        | Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1  = Ptr Word8 -> Word8 -> CSize -> IO (Ptr Word8)
B.memset Ptr Word8
dst (ByteString -> Int -> Word8
B.index ByteString
bs Int
0) (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) IO (Ptr Word8) -> 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
n
        | Bool
otherwise = do ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fptr ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
loop Ptr Word8
dst (Ptr Word8
ptr Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
o) Int
n
                         Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
n
      where (B.PS ForeignPtr Word8
fptr Int
o Int
len) = ByteString
bs
            loop :: Ptr Word8 -> Ptr Word8 -> Int -> IO ()
loop Ptr Word8
d Ptr Word8
s Int
i
                | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0    = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
len  = Ptr Word8 -> Ptr Word8 -> Int -> IO ()
B.memcpy Ptr Word8
d Ptr Word8
s (Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i)
                | Bool
otherwise = Ptr Word8 -> Ptr Word8 -> Int -> IO ()
B.memcpy Ptr Word8
d Ptr Word8
s (Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len) IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
loop (Ptr Word8
d Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
len) Ptr Word8
s (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
len)
    entropyClose :: TestEntropySource -> IO ()
entropyClose TestEntropySource
_ = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

openBackend :: EntropySource b => b -> IO (Maybe EntropyBackend)
openBackend :: b -> IO (Maybe EntropyBackend)
openBackend b
b = (b -> EntropyBackend) -> Maybe b -> Maybe EntropyBackend
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> EntropyBackend
forall b. EntropySource b => b -> EntropyBackend
EntropyBackend (Maybe b -> Maybe EntropyBackend)
-> IO (Maybe b) -> IO (Maybe EntropyBackend)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` b -> IO (Maybe b)
forall b. EntropySource b => b -> IO (Maybe b)
callOpen b
b
  where callOpen :: EntropySource b => b -> IO (Maybe b)
        callOpen :: b -> IO (Maybe b)
callOpen b
_ = IO (Maybe b)
forall a. EntropySource a => IO (Maybe a)
entropyOpen

gatherBackend :: EntropyBackend -> Ptr Word8 -> Int -> IO Int
gatherBackend :: EntropyBackend -> Ptr Word8 -> Int -> IO Int
gatherBackend (EntropyBackend b
backend) Ptr Word8
ptr Int
n = b -> Ptr Word8 -> Int -> IO Int
forall a. EntropySource a => a -> Ptr Word8 -> Int -> IO Int
entropyGather b
backend Ptr Word8
ptr Int
n

-- | Pool of Entropy. contains a self mutating pool of entropy,
-- that is always guarantee to contains data.
data EntropyPool = EntropyPool [EntropyBackend] (MVar Int) SecureMem
  deriving Typeable

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

-- | Create a new entropy pool of a specific size
--
-- You can create as many entropy pools as you want, and a given pool can be shared between multiples RNGs.
createEntropyPoolWith :: Int -> [EntropyBackend] -> IO EntropyPool
createEntropyPoolWith :: Int -> [EntropyBackend] -> IO EntropyPool
createEntropyPoolWith Int
poolSize [EntropyBackend]
backends = do
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([EntropyBackend] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [EntropyBackend]
backends) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"cannot get any source of entropy on this system"
    SecureMem
sm <- Int -> IO SecureMem
allocateSecureMem Int
poolSize
    MVar Int
m  <- Int -> IO (MVar Int)
forall a. a -> IO (MVar a)
newMVar Int
0
    SecureMem -> (Ptr Word8 -> IO ()) -> IO ()
forall b. SecureMem -> (Ptr Word8 -> IO b) -> IO b
withSecureMemPtr SecureMem
sm ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ 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 -> SecureMem -> EntropyPool
EntropyPool [EntropyBackend]
backends MVar Int
m SecureMem
sm

-- | Create a new entropy pool with a default size.
--
-- While you can create as many entropy pool 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

-- | Create a dummy entropy pool that is deterministic, and
-- dependant on the input bytestring only.
--
-- This is stricly reserved for testing purpose when a deterministic seed need
-- to be generated with deterministic RNGs.
--
-- Do not use in production code.
createTestEntropyPool :: ByteString -> EntropyPool
createTestEntropyPool :: ByteString -> EntropyPool
createTestEntropyPool ByteString
bs
    | ByteString -> Bool
B.null ByteString
bs = String -> EntropyPool
forall a. HasCallStack => String -> a
error String
"cannot create entropy pool from an empty bytestring"
    | Bool
otherwise = IO EntropyPool -> EntropyPool
forall a. IO a -> a
unsafePerformIO (IO EntropyPool -> EntropyPool) -> IO EntropyPool -> EntropyPool
forall a b. (a -> b) -> a -> b
$ Int -> [EntropyBackend] -> IO EntropyPool
createEntropyPoolWith Int
defaultPoolSize [TestEntropySource -> EntropyBackend
forall b. EntropySource b => b -> EntropyBackend
EntropyBackend (TestEntropySource -> EntropyBackend)
-> TestEntropySource -> EntropyBackend
forall a b. (a -> b) -> a -> b
$ ByteString -> TestEntropySource
TestEntropySource ByteString
bs]

-- | Put a chunk of the entropy pool into a buffer
grabEntropyPtr :: Int -> EntropyPool -> Ptr Word8 -> IO ()
grabEntropyPtr :: Int -> EntropyPool -> Ptr Word8 -> IO ()
grabEntropyPtr Int
n (EntropyPool [EntropyBackend]
backends MVar Int
posM SecureMem
sm) Ptr Word8
outPtr =
    SecureMem -> (Ptr Word8 -> IO ()) -> IO ()
forall b. SecureMem -> (Ptr Word8 -> IO b) -> IO b
withSecureMemPtr SecureMem
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 = SecureMem -> Int
secureMemGetSize SecureMem
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.
grabEntropyIO :: Int -> EntropyPool -> IO SecureMem
grabEntropyIO :: Int -> EntropyPool -> IO SecureMem
grabEntropyIO Int
n EntropyPool
pool = do
    SecureMem
out <- Int -> IO SecureMem
allocateSecureMem Int
n
    SecureMem -> (Ptr Word8 -> IO ()) -> IO ()
forall b. SecureMem -> (Ptr Word8 -> IO b) -> IO b
withSecureMemPtr SecureMem
out ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> EntropyPool -> Ptr Word8 -> IO ()
grabEntropyPtr Int
n EntropyPool
pool
    SecureMem -> IO SecureMem
forall (m :: * -> *) a. Monad m => a -> m a
return (SecureMem -> IO SecureMem) -> SecureMem -> IO SecureMem
forall a b. (a -> b) -> a -> b
$ SecureMem
out

-- | Grab a chunk of entropy from the entropy pool.
--
-- Great care need to be taken here when using the output,
-- as this use unsafePerformIO to actually get entropy.
--
-- Use grabEntropyIO if unsure.
{-# NOINLINE grabEntropy #-}
grabEntropy :: Int -> EntropyPool -> SecureMem
grabEntropy :: Int -> EntropyPool -> SecureMem
grabEntropy Int
n EntropyPool
pool = IO SecureMem -> SecureMem
forall a. IO a -> a
unsafePerformIO (IO SecureMem -> SecureMem) -> IO SecureMem -> SecureMem
forall a b. (a -> b) -> a -> b
$ Int -> EntropyPool -> IO SecureMem
grabEntropyIO Int
n EntropyPool
pool

replenish :: Int -> [EntropyBackend] -> Ptr Word8 -> IO ()
replenish :: Int -> [EntropyBackend] -> Ptr Word8 -> IO ()
replenish Int
poolSize [EntropyBackend]
backends Ptr Word8
ptr = Int -> [EntropyBackend] -> Ptr Word8 -> Int -> IO ()
loop Int
0 [EntropyBackend]
backends Ptr Word8
ptr Int
poolSize
  where loop :: Int -> [EntropyBackend] -> Ptr Word8 -> Int -> IO ()
        loop :: Int -> [EntropyBackend] -> Ptr Word8 -> Int -> IO ()
loop Int
retry [] Ptr Word8
p Int
n | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0     = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                          | Int
retry Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3 = String -> IO ()
forall a. HasCallStack => String -> a
error String
"cannot fully replenish"
                          | Bool
otherwise  = Int -> [EntropyBackend] -> Ptr Word8 -> Int -> IO ()
loop (Int
retryInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [EntropyBackend]
backends Ptr Word8
p Int
n
        loop Int
_     (EntropyBackend
_:[EntropyBackend]
_)  Ptr Word8
_ Int
0 = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        loop Int
retry (EntropyBackend
b:[EntropyBackend]
bs) Ptr Word8
p Int
n = do
            Int
r <- EntropyBackend -> Ptr Word8 -> Int -> IO Int
gatherBackend EntropyBackend
b Ptr Word8
p Int
n
            Int -> [EntropyBackend] -> Ptr Word8 -> Int -> IO ()
loop Int
retry [EntropyBackend]
bs (Ptr Word8
p Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
r) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
r)