CSPRNG.hs

{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}

-- |
-- Module: Crypto.CSPRNG
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- A fast cryptographically-secure pseudorandom number generator:
-- [ChaCha8Rand](https://c2sp.org/chacha8rand), as used by Go's
-- @math/rand/v2@ and runtime.
--
-- The generator draws 992 bytes of output at a time from ChaCha8, and
-- rekeys itself from the next 32 bytes. Its key and output buffer are
-- erased as they are used, and the native ChaCha8 kernels (on aarch64
-- and x86_64) zero the stack they use, so that a later disclosure of
-- the process's memory reveals nothing about prior output. This
-- covers the generator's own state only: copies of seeds passed to it
-- and of output it returns are the caller's to manage, and on other
-- platforms the pure Haskell ChaCha8 fallback makes copies of keys and
-- keystream that the garbage collector manages and can't erase.
--
-- The generator is mutable, living in any 'PrimMonad' (e.g. 'IO' or
-- 'ST s'), and is not thread-safe. Use one generator per thread,
-- deriving more via 'split'.
--
-- A forked process (e.g. via @forkProcess@), or a restored virtual
-- machine snapshot, inherits a copy of the generator's state, and
-- then produces the same output as the original. Create a new
-- generator in such a child, e.g. via 'new_system', rather than using
-- an inherited one.

module Crypto.CSPRNG (
    -- * Generator
    Gen
  , new
  , new_system

    -- * Output
  , word64
  , word64_range
  , int_range
  , bytes

    -- * Splitting
  , split

    -- * Wiping
  , wipe
  ) where

import Control.Monad.Primitive (PrimMonad, PrimState, unsafeIOToPrim)
import Control.Monad.ST (ST)
import qualified Crypto.Cipher.ChaCha as ChaCha
import qualified Crypto.Entropy as Entropy
import qualified Crypto.Range as Range
import Data.Bits ((.&.), complement)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BI
import qualified Data.ByteString.Unsafe as BU
import qualified Data.Primitive.ByteArray as BA
import Data.Word (Word8, Word64, byteSwap64)
import Foreign.ForeignPtr (touchForeignPtr)
import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
import Foreign.Marshal.Utils (copyBytes)
import Foreign.Ptr (Ptr, castPtr, plusPtr)
import GHC.ByteOrder (ByteOrder(..), targetByteOrder)
import GHC.Exts (RealWorld)

-- utils ----------------------------------------------------------------------

fi :: (Integral a, Num b) => a -> b
fi = fromIntegral
{-# INLINE fi #-}

-- convert between host and little-endian byte order
le64 :: Word64 -> Word64
le64 w = case targetByteOrder of
  LittleEndian -> w
  BigEndian    -> byteSwap64 w
{-# INLINE le64 #-}

-- state ----------------------------------------------------------------------

-- a generator's state is a pinned byte array, laid out as:
--
--   offset  size  contents
--   0       32    key for the next refill
--   32      16    zero nonce (12 bytes used)
--   48      8     position of the next unread byte of output
--   64      1024  output (992 bytes) and next key (32 bytes)
--
-- all bytes of output, and all keys, are little-endian.

-- | A ChaCha8Rand generator.
newtype Gen s = Gen (BA.MutableByteArray s)

_KEY, _NONCE, _POS, _OUT, _SIZE, _AVAIL :: Int
_KEY   = 0
_NONCE = 32
_POS   = 48
_OUT   = 64
_SIZE  = 1088
_AVAIL = 992

alloc :: PrimMonad m => m (BA.MutableByteArray (PrimState m))
alloc = do
  st <- BA.newPinnedByteArray _SIZE
  BA.setByteArray st 0 _SIZE (0 :: Word8)
  pure st
{-# INLINE alloc #-}

read_pos :: PrimMonad m => BA.MutableByteArray (PrimState m) -> m Int
read_pos st = do
  !p <- BA.readByteArray st (_POS `quot` 8)
  pure $! fi (p :: Word64)
{-# INLINE read_pos #-}

write_pos :: PrimMonad m => BA.MutableByteArray (PrimState m) -> Int -> m ()
write_pos st p = BA.writeByteArray st (_POS `quot` 8) (fi p :: Word64)
{-# INLINE write_pos #-}

ptr :: BA.MutableByteArray s -> Int -> Ptr Word8
ptr st off = BA.mutableByteArrayContents st `plusPtr` off
{-# INLINE ptr #-}

-- ChaCha8Rand iteration: replace the output buffer with 992 fresh
-- bytes of output, and rekey.
refill :: PrimMonad m => BA.MutableByteArray (PrimState m) -> m ()
refill st = do
  -- sixteen ChaCha8 blocks, with counters 0 to 15, interleaved four
  -- at a time and with the final addition covering only the key
  unsafeIOToPrim $
    ChaCha._interleaved_into ChaCha.ChaCha8 ChaCha.FeedForwardKey
      (ptr st _KEY) 0 (ptr st _NONCE) (ptr st _OUT) 4
  -- the last 32 bytes of output are the next key
  BA.copyMutableByteArray st _KEY st (_OUT + _AVAIL) 32
  BA.setByteArray st (_OUT + _AVAIL) 32 (0 :: Word8)
  write_pos st 0
{-# INLINABLE refill #-}
{-# SPECIALIZE refill :: BA.MutableByteArray RealWorld -> IO () #-}
{-# SPECIALIZE refill :: BA.MutableByteArray s -> ST s () #-}

-- copy the next bytes of output to the destination, erasing them
fill
  :: PrimMonad m
  => BA.MutableByteArray (PrimState m)
  -> Ptr Word8
  -> Int
  -> m ()
fill st dst n = loop 0 where
  loop !off
    | off == n  = pure ()
    | otherwise = do
        !pos <- read_pos st
        if   pos == _AVAIL
        then do
          refill st
          loop off
        else do
          let !k = min (n - off) (_AVAIL - pos)
              !d = dst `plusPtr` off :: Ptr Word8
          BA.copyMutableByteArrayToPtr d st (_OUT + pos) k
          BA.setByteArray st (_OUT + pos) k (0 :: Word8)
          write_pos st (pos + k)
          loop (off + k)
{-# INLINABLE fill #-}

-- generator ------------------------------------------------------------------

-- | Create a generator from a 256-bit seed, returning 'Nothing' if the
--   seed is not exactly 32 bytes.
--
--   The seed must be secret and uniformly random; prefer 'new_system'
--   unless a reproducible stream is required.
--
--   >>> Just gen <- new "ABCDEFGHIJKLMNOPQRSTUVWXYZ123456"
--   >>> word64 gen
--   13219109469176600229
new
  :: PrimMonad m
  => BS.ByteString
  -> m (Maybe (Gen (PrimState m)))
new seed@(BI.PS _ _ l)
  | l /= 32   = pure Nothing
  | otherwise = do
      st <- alloc
      unsafeIOToPrim $ BU.unsafeUseAsCString seed $ \sp ->
        copyBytes (ptr st _KEY) (castPtr sp) 32
      refill st
      pure $! Just (Gen st)
{-# INLINABLE new #-}
{-# SPECIALIZE new :: BS.ByteString -> IO (Maybe (Gen RealWorld)) #-}
{-# SPECIALIZE new :: BS.ByteString -> ST s (Maybe (Gen s)) #-}

-- | Create a generator seeded from the operating system's entropy
--   source, returning 'Nothing' if that fails.
--
--   >>> Just gen <- new_system
new_system :: IO (Maybe (Gen RealWorld))
new_system = do
  st <- alloc
  ok <- Entropy.fill (ptr st _KEY) 32
  if   not ok
  then pure Nothing
  else do
    refill st
    pure $! Just (Gen st)

-- output ---------------------------------------------------------------------

-- | Generate a uniformly-distributed 'Word64' from the next eight
--   bytes of output, read in little-endian order. If a call to 'bytes'
--   has left the output out of eight-byte alignment, the bytes up to
--   the next boundary are discarded.
--
--   >>> Just gen <- new "ABCDEFGHIJKLMNOPQRSTUVWXYZ123456"
--   >>> word64 gen
--   13219109469176600229
word64 :: PrimMonad m => Gen (PrimState m) -> m Word64
word64 (Gen st) = do
  !pos <- read_pos st
  let !aligned = (pos + 7) .&. complement 7
  !p <- if   aligned + 8 > _AVAIL
        then do
          refill st
          pure 0
        else pure aligned
  let !ix = (_OUT + p) `quot` 8
  !w <- BA.readByteArray st ix
  BA.writeByteArray st ix (0 :: Word64)
  write_pos st (p + 8)
  pure $! le64 w
{-# INLINE word64 #-}

-- | Generate a uniformly-distributed 'Word64' in the given inclusive
--   range, whose bounds may be given in either order.
--
--   The result is unbiased; see 'int_range'.
--
--   >>> Just gen <- new "ABCDEFGHIJKLMNOPQRSTUVWXYZ123456"
--   >>> word64_range 1 6 gen
--   5
word64_range
  :: PrimMonad m => Word64 -> Word64 -> Gen (PrimState m) -> m Word64
word64_range a b gen = do
  let !lo = min a b
      !hi = max a b
  !r <- Range.upto (word64 gen) (hi - lo)
  pure $! lo + r
{-# INLINE word64_range #-}

-- | Generate a uniformly-distributed 'Int' in the given inclusive
--   range, whose bounds may be given in either order.
--
--   The result is unbiased, via Lemire's method: an occasional draw is
--   rejected and replaced, so the number of words consumed, and the
--   running time, depend on the (discarded) rejected words.
--
--   >>> Just gen <- new "ABCDEFGHIJKLMNOPQRSTUVWXYZ123456"
--   >>> int_range (-10) 10 gen
--   5
int_range :: PrimMonad m => Int -> Int -> Gen (PrimState m) -> m Int
int_range a b gen = do
  let !lo = min a b
      !hi = max a b
  !r <- Range.upto (word64 gen) (fi hi - fi lo)
  pure $! lo + fi r
{-# INLINE int_range #-}

-- | Generate the given number of bytes. Non-positive counts yield the
--   empty string.
--
--   >>> import qualified Data.ByteString.Base16 as B16
--   >>> Just gen <- new "ABCDEFGHIJKLMNOPQRSTUVWXYZ123456"
--   >>> fmap B16.encode (bytes 8 gen)
--   "a516463d06b673b7"
bytes :: PrimMonad m => Int -> Gen (PrimState m) -> m BS.ByteString
bytes n (Gen st)
  | n <= 0    = pure BS.empty
  | otherwise = do
      fp <- unsafeIOToPrim (BI.mallocByteString n)
      fill st (unsafeForeignPtrToPtr fp) n
      unsafeIOToPrim (touchForeignPtr fp)
      pure $! BI.BS fp n
{-# INLINABLE bytes #-}
{-# SPECIALIZE bytes :: Int -> Gen RealWorld -> IO BS.ByteString #-}
{-# SPECIALIZE bytes :: Int -> Gen s -> ST s BS.ByteString #-}

-- splitting ------------------------------------------------------------------

-- | Derive a new generator, seeded from the next 32 bytes of output.
--
--   The two generators' outputs are independent; use 'split' to give
--   each thread its own generator.
--
--   >>> Just gen <- new_system
--   >>> child <- split gen
split :: PrimMonad m => Gen (PrimState m) -> m (Gen (PrimState m))
split (Gen st) = do
  child <- alloc
  fill st (ptr child _KEY) 32
  refill child
  pure $! Gen child
{-# INLINABLE split #-}

-- wiping ---------------------------------------------------------------------

-- | Erase a generator's state.
--
--   Call this when finished with a generator, to wipe its key and any
--   unused output from memory. The generator must not be used
--   afterwards: it would then produce output from an all-zero key,
--   which is predictable.
--
--   >>> Just gen <- new_system
--   >>> out <- bytes 32 gen
--   >>> wipe gen
wipe :: PrimMonad m => Gen (PrimState m) -> m ()
wipe (Gen st) = do
  BA.setByteArray st 0 _SIZE (0 :: Word8)
  write_pos st _AVAIL
{-# INLINE wipe #-}