module Randomization
  ( Randomization (..)
  , deterministicallyRandom
  , DeterministicRandomization (..)
  , hoistRandomization
  , defaultRandomization
  ) where

import Internal.Prelude

import Crypto.Random (ChaChaDRG, DRG (randomBytesGenerate), drgNew)
import Data.IORef (atomicModifyIORef', newIORef)

-- | General means of obtaining randomness
newtype Randomization m = Randomization
  { forall (m :: * -> *). Randomization m -> Natural -> m ByteString
getRandomBytes :: Natural -> m ByteString
  -- ^ Given a requested number of bytes, this action
  --   should produce a 'ByteString' of that length.
  }

hoistRandomization
  :: (forall a. m a -> m' a) -> Randomization m -> Randomization m'
hoistRandomization :: forall (m :: * -> *) (m' :: * -> *).
(forall a. m a -> m' a) -> Randomization m -> Randomization m'
hoistRandomization forall a. m a -> m' a
f (Randomization Natural -> m ByteString
g) = (Natural -> m' ByteString) -> Randomization m'
forall (m :: * -> *). (Natural -> m ByteString) -> Randomization m
Randomization (m ByteString -> m' ByteString
forall a. m a -> m' a
f (m ByteString -> m' ByteString)
-> (Natural -> m ByteString) -> Natural -> m' ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Natural -> m ByteString
g)

-- | Convert from a deterministic generator to an effectful one
deterministicallyRandom
  :: DeterministicRandomization -> IO (Randomization IO)
deterministicallyRandom :: DeterministicRandomization -> IO (Randomization IO)
deterministicallyRandom =
  IO (IORef DeterministicRandomization)
-> IO (IORef DeterministicRandomization)
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef DeterministicRandomization)
 -> IO (IORef DeterministicRandomization))
-> (DeterministicRandomization
    -> IO (IORef DeterministicRandomization))
-> DeterministicRandomization
-> IO (IORef DeterministicRandomization)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DeterministicRandomization -> IO (IORef DeterministicRandomization)
forall a. a -> IO (IORef a)
newIORef (DeterministicRandomization
 -> IO (IORef DeterministicRandomization))
-> (IORef DeterministicRandomization -> IO (Randomization IO))
-> DeterministicRandomization
-> IO (Randomization IO)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Randomization IO -> IO (Randomization IO)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Randomization IO -> IO (Randomization IO))
-> (IORef DeterministicRandomization -> Randomization IO)
-> IORef DeterministicRandomization
-> IO (Randomization IO)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. \IORef DeterministicRandomization
ref ->
    (Natural -> IO ByteString) -> Randomization IO
forall (m :: * -> *). (Natural -> m ByteString) -> Randomization m
Randomization ((Natural -> IO ByteString) -> Randomization IO)
-> (Natural -> IO ByteString) -> Randomization IO
forall a b. (a -> b) -> a -> b
$ \Natural
n ->
      IO ByteString -> IO ByteString
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> IO ByteString) -> IO ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ IORef DeterministicRandomization
-> (DeterministicRandomization
    -> (DeterministicRandomization, ByteString))
-> IO ByteString
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef DeterministicRandomization
ref ((DeterministicRandomization
  -> (DeterministicRandomization, ByteString))
 -> IO ByteString)
-> (DeterministicRandomization
    -> (DeterministicRandomization, ByteString))
-> IO ByteString
forall a b. (a -> b) -> a -> b
$ \(DeterministicRandomization Natural -> (ByteString, DeterministicRandomization)
gen) ->
        (ByteString, DeterministicRandomization)
-> (DeterministicRandomization, ByteString)
forall a b. (a, b) -> (b, a)
swap ((ByteString, DeterministicRandomization)
 -> (DeterministicRandomization, ByteString))
-> (ByteString, DeterministicRandomization)
-> (DeterministicRandomization, ByteString)
forall a b. (a -> b) -> a -> b
$ Natural -> (ByteString, DeterministicRandomization)
gen Natural
n

-- | A deterministic random generator
newtype DeterministicRandomization = DeterministicRandomization
  { DeterministicRandomization
-> Natural -> (ByteString, DeterministicRandomization)
nextRandomBytes :: Natural -> (ByteString, DeterministicRandomization)
  -- ^ Given a requested number of bytes, this function should give a
  --   'ByteString' of that length and a new deterministic generator.
  }

-- | Cryptographically secure deterministic randomization seeded from
--   system entropy using @ChaChaDRG@ from the @crypton@ package
defaultRandomization :: IO (Randomization IO)
defaultRandomization :: IO (Randomization IO)
defaultRandomization =
  DeterministicRandomization -> IO (Randomization IO)
deterministicallyRandom (DeterministicRandomization -> IO (Randomization IO))
-> (ChaChaDRG -> DeterministicRandomization)
-> ChaChaDRG
-> IO (Randomization IO)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ChaChaDRG -> DeterministicRandomization
makeDeterministicRandomization (ChaChaDRG -> IO (Randomization IO))
-> IO ChaChaDRG -> IO (Randomization IO)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO ChaChaDRG -> IO ChaChaDRG
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ChaChaDRG
forall (randomly :: * -> *).
MonadRandom randomly =>
randomly ChaChaDRG
drgNew
 where
  makeDeterministicRandomization :: ChaChaDRG -> DeterministicRandomization
  makeDeterministicRandomization :: ChaChaDRG -> DeterministicRandomization
makeDeterministicRandomization ChaChaDRG
drg =
    (Natural -> (ByteString, DeterministicRandomization))
-> DeterministicRandomization
DeterministicRandomization ((Natural -> (ByteString, DeterministicRandomization))
 -> DeterministicRandomization)
-> (Natural -> (ByteString, DeterministicRandomization))
-> DeterministicRandomization
forall a b. (a -> b) -> a -> b
$ \Natural
n ->
      let (ByteString
bs, ChaChaDRG
drg') = Int -> ChaChaDRG -> (ByteString, ChaChaDRG)
forall gen byteArray.
(DRG gen, ByteArray byteArray) =>
Int -> gen -> (byteArray, gen)
forall byteArray.
ByteArray byteArray =>
Int -> ChaChaDRG -> (byteArray, ChaChaDRG)
randomBytesGenerate (Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
n) ChaChaDRG
drg
      in  (ByteString
bs, ChaChaDRG -> DeterministicRandomization
makeDeterministicRandomization ChaChaDRG
drg')