{-# LANGUAGE CPP, ForeignFunctionInterface, BangPatterns, ScopedTypeVariables #-}
{-|
 Maintainer: Thomas.DuBuisson@gmail.com
 Stability: beta
 Portability: portable

 Obtain entropy from system sources or x86 RDRAND when available.

-}

module System.EntropyNix
        ( CryptHandle
        , openHandle
        , hGetEntropy
        , closeHandle
        , hardwareRandom
        ) where

import Control.Exception
import Control.Monad (liftM, when)
import Data.ByteString as B
import System.IO.Error (mkIOError, eofErrorType, ioeSetErrorString)
import System.IO.Unsafe
import Data.Bits (xor)

import Foreign (allocaBytes)
import Foreign.Ptr
import Foreign.C.Error
import Foreign.C.Types
import Data.ByteString.Internal as B

#ifdef arch_i386
-- See .cabal wrt GCC 4.8.2 asm compilation bug
#undef HAVE_RDRAND
#endif

import System.Posix (openFd, closeFd, fdReadBuf, OpenMode(..), defaultFileFlags, Fd, OpenFileFlags(..))

source :: FilePath
source :: FilePath
source = FilePath
"/dev/urandom"

-- |Handle for manual resource management
data CryptHandle
    = CH Fd
#ifdef HAVE_GETRANDOM
    | UseGetRandom
#endif

-- | Get random values from the hardware RNG or return Nothing if no
-- supported hardware RNG is available.
--
-- Supported hardware:
--      * RDRAND
--      * Patches welcome
hardwareRandom :: Int -> IO (Maybe B.ByteString)
#ifdef HAVE_RDRAND
hardwareRandom n =
 do b <- cpuHasRdRand
    if b then Just <$> B.create n (\ptr ->
                      do r <- c_get_rand_bytes (castPtr ptr) (fromIntegral n)
                         when (r /= 0) (fail "RDRand failed to gather entropy"))
     else pure Nothing
#else
hardwareRandom :: Int -> IO (Maybe ByteString)
hardwareRandom Int
_ = Maybe ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ByteString
forall a. Maybe a
Nothing
#endif

-- |Open a `CryptHandle`
openHandle :: IO CryptHandle
openHandle :: IO CryptHandle
openHandle =
#ifdef HAVE_GETRANDOM
  if systemHasGetRandom then return UseGetRandom else
#endif
  (Fd -> CryptHandle) -> IO Fd -> IO CryptHandle
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Fd -> CryptHandle
CH IO Fd
openRandomFile

openRandomFile :: IO Fd
openRandomFile :: IO Fd
openRandomFile = do
  CInt -> IO CInt
forall a. a -> IO a
evaluate CInt
ensurePoolInitialized
#if MIN_VERSION_unix(2,8,0)
  openFd source ReadOnly defaultFileFlags { creat = Nothing }
#else
  FilePath -> OpenMode -> Maybe FileMode -> OpenFileFlags -> IO Fd
openFd FilePath
source OpenMode
ReadOnly Maybe FileMode
forall a. Maybe a
Nothing OpenFileFlags
defaultFileFlags
#endif

-- |Close the `CryptHandle`
closeHandle :: CryptHandle -> IO ()
closeHandle :: CryptHandle -> IO ()
closeHandle (CH Fd
h) = Fd -> IO ()
closeFd Fd
h
#ifdef HAVE_GETRANDOM
closeHandle UseGetRandom = return ()
#endif

-- |Read random data from a `CryptHandle`
hGetEntropy :: CryptHandle -> Int -> IO B.ByteString
hGetEntropy :: CryptHandle -> Int -> IO ByteString
hGetEntropy (CH Fd
h) Int
n = Fd -> Int -> IO ByteString
fdReadBS Fd
h Int
n
#ifdef HAVE_GETRANDOM
hGetEntropy UseGetRandom n = do
  bs <- B.createUptoN n (\ptr -> do
    r <- c_entropy_getrandom (castPtr ptr) (fromIntegral n)
    return $ if r == 0 then n else 0)
  if B.length bs == n then return bs
  -- getrandom somehow failed. Fall back on /dev/urandom instead.
  else bracket openRandomFile closeFd $ flip fdReadBS n
#endif

fdReadBS :: Fd -> Int -> IO B.ByteString
fdReadBS :: Fd -> Int -> IO ByteString
fdReadBS Fd
fd Int
n =
    Int -> (Ptr Any -> IO ByteString) -> IO ByteString
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
n ((Ptr Any -> IO ByteString) -> IO ByteString)
-> (Ptr Any -> IO ByteString) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Any
buf -> Ptr Any -> Int -> IO ByteString
forall a. Ptr a -> Int -> IO ByteString
go Ptr Any
buf Int
n
 where
 go :: Ptr a -> Int -> IO ByteString
go Ptr a
buf Int
0   = CStringLen -> IO ByteString
B.packCStringLen (Ptr a -> Ptr CChar
forall a b. Ptr a -> Ptr b
castPtr Ptr a
buf, Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)
 go Ptr a
buf Int
cnt  | Int
cnt Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n = do
        ByteCount
rc <- Fd -> Ptr Word8 -> ByteCount -> IO ByteCount
fdReadBuf Fd
fd (Ptr a -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr a
buf (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cnt)) (Int -> ByteCount
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
cnt)
        case ByteCount
rc of
            ByteCount
0 -> IOError -> IO ByteString
forall a. IOError -> IO a
ioError (IOError -> FilePath -> IOError
ioeSetErrorString (IOErrorType
-> FilePath -> Maybe Handle -> Maybe FilePath -> IOError
mkIOError IOErrorType
eofErrorType FilePath
"fdRead" Maybe Handle
forall a. Maybe a
Nothing Maybe FilePath
forall a. Maybe a
Nothing) FilePath
"EOF")
            ByteCount
n' -> Ptr a -> Int -> IO ByteString
go Ptr a
buf (Int
cnt Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteCount -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ByteCount
n')
 go Ptr a
_ Int
_     = FilePath -> IO ByteString
forall a. HasCallStack => FilePath -> a
error FilePath
"Impossible!  The count of bytes left to read is greater than the request or less than zero!"

#ifdef HAVE_GETRANDOM
foreign import ccall unsafe "system_has_getrandom"
   c_system_has_getrandom :: IO CInt
foreign import ccall safe "entropy_getrandom"
  c_entropy_getrandom :: Ptr CUChar -> CSize -> IO CInt

-- NOINLINE and unsafePerformIO are not totally necessary as getrandom will be
-- consistently either present or not, but it is cheaper not to check multiple
-- times.
systemHasGetRandom :: Bool
{-# NOINLINE systemHasGetRandom #-}
systemHasGetRandom = unsafePerformIO $ fmap (/= 0) c_system_has_getrandom
#endif

foreign import ccall safe "ensure_pool_initialized"
   c_ensure_pool_initialized :: IO CInt

-- Similarly to systemHasGetRandom, NOINLINE is just an optimization.
ensurePoolInitialized :: CInt
{-# NOINLINE ensurePoolInitialized #-}
ensurePoolInitialized :: CInt
ensurePoolInitialized = IO CInt -> CInt
forall a. IO a -> a
unsafePerformIO (IO CInt -> CInt) -> IO CInt -> CInt
forall a b. (a -> b) -> a -> b
$ FilePath -> IO CInt -> IO CInt
forall a. (Eq a, Num a) => FilePath -> IO a -> IO a
throwErrnoIfMinus1 FilePath
"ensurePoolInitialized" (IO CInt -> IO CInt) -> IO CInt -> IO CInt
forall a b. (a -> b) -> a -> b
$ IO CInt
c_ensure_pool_initialized

#ifdef HAVE_RDRAND
foreign import ccall unsafe "cpu_has_rdrand"
   c_cpu_has_rdrand :: IO CInt

foreign import ccall unsafe "get_rand_bytes"
  c_get_rand_bytes :: Ptr CUChar -> CSize -> IO CInt

cpuHasRdRand :: IO Bool
cpuHasRdRand = (/= 0) `fmap` c_cpu_has_rdrand
#endif