    -- |
    -- Module      : Data.SecureMem
    -- License     : BSD-style
    -- Maintainer  : Vincent Hanquez <vincent@snarc.org>
    -- Stability   : Stable
    -- Portability : GHC
    --
    {-# LANGUAGE ForeignFunctionInterface #-}
    {-# LANGUAGE BangPatterns #-}
    {-# LANGUAGE MagicHash #-}
    {-# LANGUAGE CPP #-}
    module Data.SecureMem
        ( SecureMem
        , secureMemGetSize
        , secureMemCopy
        , ToSecureMem(..)
        -- * Allocation and early termination
        , allocateSecureMem
        , createSecureMem
        , unsafeCreateSecureMem
        , finalizeSecureMem
        -- * Pointers manipulation
        , withSecureMemPtr
        , withSecureMemPtrSz
        , withSecureMemCopy
        -- * convertion
        , secureMemFromByteString
        ) where

    import Foreign.C.Types
    import Foreign.ForeignPtr (withForeignPtr, finalizeForeignPtr)
    import Foreign.Ptr
    #if MIN_VERSION_base(4,6,0)
    import GHC.Types (Int(..))
    import GHC.ForeignPtr (ForeignPtr(..), ForeignPtrContents(..), mallocForeignPtrBytes, addForeignPtrConcFinalizer)
    import GHC.Prim (sizeofMutableByteArray#)
    import GHC.Ptr (Ptr(..))
    #else
    import Foreign.ForeignPtr (ForeignPtr, mallocForeignPtrBytes)
    #endif
    import Data.Word (Word8)
    import Data.Monoid
    import Control.Applicative
    import Control.Monad (foldM, void)
    import Data.ByteString (ByteString)
    import Data.Byteable
    import qualified Data.ByteString.Internal as B

    -- | SecureMem is a memory chunk which have the properties of:
    --
    -- * Being scrubbed after its goes out of scope.
    -- * A Show instance that doesn't actually show any content
    -- * A Eq instance that is constant time
    #if MIN_VERSION_base(4,6,0)
    newtype SecureMem = SecureMem { getForeignPtr :: ForeignPtr Word8 }
    #else
    data SecureMem = SecureMem { getForeignPtr :: ForeignPtr Word8, secureMemGetSize :: Int }
    #endif

    #if MIN_VERSION_base(4,6,0)
    -- | Return the size of the memory allocated by this secure mem.
    secureMemGetSize :: SecureMem -> Int
    secureMemGetSize (SecureMem (ForeignPtr _ fpc)) =
        case fpc of
            MallocPtr mba _ -> I# (sizeofMutableByteArray# mba)
            _               -> error "cannot happen"
    #endif

    secureMemEq :: SecureMem -> SecureMem -> Bool
    secureMemEq sm1 sm2 = sz1 == sz2 && B.inlinePerformIO meq
      where meq = withSecureMemPtr sm1 $ \ptr1 ->
                  withSecureMemPtr sm2 $ \ptr2 ->
                  feq ptr1 ptr2
            !sz1 = secureMemGetSize sm1
            !sz2 = secureMemGetSize sm2
            feq | sz1 == 8   = compareEq8
                | sz1 == 16  = compareEq16
                | sz1 == 24  = compareEq24
                | sz1 == 32  = compareEq32
                | sz1 == 64  = compareEq64
                | sz1 == 128 = compareEq128
                | sz1 == 256 = compareEq256
                | otherwise  = compareEq (fromIntegral sz1)

    secureMemAppend :: SecureMem -> SecureMem -> SecureMem
    secureMemAppend s1 s2 =
        unsafeCreateSecureMem (sz1+sz2) $ \dst -> do
            withSecureMemPtr s1 $ \sp1 -> B.memcpy dst sp1 (fromIntegral sz1)
            withSecureMemPtr s2 $ \sp2 -> B.memcpy (dst `plusPtr` sz1) sp2 (fromIntegral sz2)
      where !sz1 = secureMemGetSize s1
            !sz2 = secureMemGetSize s2

    secureMemConcat :: [SecureMem] -> SecureMem
    secureMemConcat l = unsafeCreateSecureMem total $ \dst -> void $ foldM copy dst l
      where total = sum $ map secureMemGetSize l
            copy dst s = withSecureMemPtr s $ \sp -> do
                        B.memcpy dst sp (fromIntegral sz)
                        return (dst `plusPtr` sz)
              where sz = secureMemGetSize s

    secureMemCopy :: SecureMem -> IO SecureMem
    secureMemCopy src =
        createSecureMem sz $ \dst ->
        withSecureMemPtr src $ \s -> B.memcpy dst s (fromIntegral sz)
      where sz = secureMemGetSize src

    withSecureMemCopy :: SecureMem -> (Ptr Word8 -> IO ()) -> IO SecureMem
    withSecureMemCopy sm f = do
        sm2 <- secureMemCopy sm
        withSecureMemPtr sm2 f
        return sm2

    instance Show SecureMem where
        show _ = "<secure-mem>"

    instance Byteable SecureMem where
        toBytes        = secureMemToByteString
        byteableLength = secureMemGetSize
        withBytePtr    = withSecureMemPtr

    instance Eq SecureMem where
        (==) = secureMemEq

    instance Monoid SecureMem where
        mempty  = unsafeCreateSecureMem 0 (\_ -> return ())
        mappend = secureMemAppend
        mconcat = secureMemConcat

    type Finalizer = Ptr Word8 -> IO ()
    type FinalizerWithSize = CInt -> Ptr Word8 -> IO ()

    -- | Types that can be converted to a secure mem object.
    class ToSecureMem a where
        toSecureMem :: a -> SecureMem

    instance ToSecureMem SecureMem where
        toSecureMem a = a
    instance ToSecureMem ByteString where
        toSecureMem bs = secureMemFromByteString bs

    -- | Generic eq function
    foreign import ccall "compare_eq"    compareEq    :: CInt -> Ptr Word8 -> Ptr Word8 -> IO Bool

    -- | Specialized eq functions which are much faster
    foreign import ccall "compare_eq8"   compareEq8   :: Ptr Word8 -> Ptr Word8 -> IO Bool
    foreign import ccall "compare_eq16"  compareEq16  :: Ptr Word8 -> Ptr Word8 -> IO Bool
    foreign import ccall "compare_eq24"  compareEq24  :: Ptr Word8 -> Ptr Word8 -> IO Bool
    foreign import ccall "compare_eq32"  compareEq32  :: Ptr Word8 -> Ptr Word8 -> IO Bool
    foreign import ccall "compare_eq64"  compareEq64  :: Ptr Word8 -> Ptr Word8 -> IO Bool
    foreign import ccall "compare_eq128" compareEq128 :: Ptr Word8 -> Ptr Word8 -> IO Bool
    foreign import ccall "compare_eq256" compareEq256 :: Ptr Word8 -> Ptr Word8 -> IO Bool

    #if MIN_VERSION_base(4,6,0)
    -- | Scruber finalizers
    foreign import ccall "finalizer_scrub8" finalizerScrub8     :: Finalizer
    foreign import ccall "finalizer_scrub16" finalizerScrub16   :: Finalizer
    foreign import ccall "finalizer_scrub24" finalizerScrub24   :: Finalizer
    foreign import ccall "finalizer_scrub32" finalizerScrub32   :: Finalizer
    foreign import ccall "finalizer_scrub64" finalizerScrub64   :: Finalizer
    foreign import ccall "finalizer_scrub128" finalizerScrub128 :: Finalizer
    foreign import ccall "finalizer_scrub256" finalizerScrub256 :: Finalizer
    foreign import ccall "finalizer_scrubvar" finalizerScrubVar :: FinalizerWithSize

    szToScruber :: Int -> Ptr Word8 -> IO ()
    szToScruber 0   = \_ -> return ()
    szToScruber 8   = finalizerScrub8
    szToScruber 16  = finalizerScrub16
    szToScruber 24  = finalizerScrub24
    szToScruber 32  = finalizerScrub32
    szToScruber 64  = finalizerScrub64
    szToScruber 128 = finalizerScrub128
    szToScruber 256 = finalizerScrub256
    szToScruber n   = finalizerScrubVar (fromIntegral n)
    #endif

    -- | Allocate a foreign ptr which will be scrubed before memory free.
    -- the memory is allocated on the haskell heap
    allocateScrubedForeignPtr :: Int -> IO (ForeignPtr a)
    allocateScrubedForeignPtr sz = do
    #if MIN_VERSION_base(4,6,0)
        fptr@(ForeignPtr addr _) <- mallocForeignPtrBytes sz
        addForeignPtrConcFinalizer fptr (scruber (Ptr addr))
        return fptr
      where !scruber = szToScruber sz
    #else
        mallocForeignPtrBytes sz
    #endif

    -- | Allocate a new SecureMem
    --
    -- The memory is allocated on the haskell heap, and will be scrubed
    -- before being released.
    allocateSecureMem :: Int -> IO SecureMem
    #if MIN_VERSION_base(4,6,0)
    allocateSecureMem sz = SecureMem <$> allocateScrubedForeignPtr sz
    #else
    allocateSecureMem sz = SecureMem <$> allocateScrubedForeignPtr sz <*> pure sz
    #endif

    -- | Create a new secure mem and running an initializer function
    createSecureMem :: Int -> (Ptr Word8 -> IO ()) -> IO SecureMem
    createSecureMem sz f = do
        sm <- allocateSecureMem sz
        withSecureMemPtr sm f
        return sm

    -- | Create a new secure mem using inline perform IO to create a pure
    -- result.
    unsafeCreateSecureMem :: Int -> (Ptr Word8 -> IO ()) -> SecureMem
    unsafeCreateSecureMem sz f = B.inlinePerformIO (createSecureMem sz f)

    -- | This is a way to look at the pointer living inside a foreign object. This
    -- function takes a function which is applied to that pointer. The resulting IO
    -- action is then executed
    --
    -- this is similary to withForeignPtr for a ForeignPtr
    withSecureMemPtr :: SecureMem -> (Ptr Word8 -> IO b) -> IO b
    withSecureMemPtr sm f = withForeignPtr (getForeignPtr sm) f

    -- | similar to withSecureMem but also include the size of the pointed memory.
    withSecureMemPtrSz :: SecureMem -> (Int -> Ptr Word8 -> IO b) -> IO b
    withSecureMemPtrSz sm f = withForeignPtr (getForeignPtr sm) (f (secureMemGetSize sm))

    -- | Finalize a SecureMem early
    finalizeSecureMem :: SecureMem -> IO ()
    finalizeSecureMem sm = finalizeForeignPtr $ getForeignPtr sm

    -- | Create a bytestring from a Secure Mem
    secureMemToByteString :: SecureMem -> ByteString
    secureMemToByteString sm =
        B.unsafeCreate sz $ \dst ->
        withSecureMemPtr sm $ \src ->
        B.memcpy dst src (fromIntegral sz)
      where !sz = secureMemGetSize sm

    -- | Create a SecureMem from a bytestring
    secureMemFromByteString :: ByteString -> SecureMem
    secureMemFromByteString b = B.inlinePerformIO $ do
        sm <- allocateSecureMem len
        withSecureMemPtr sm $ \dst -> withBytestringPtr $ \src -> B.memcpy dst src (fromIntegral len)
        return sm
      where (fp, off, !len) = B.toForeignPtr b
            withBytestringPtr f = withForeignPtr fp $ \p -> f (p `plusPtr` off)