    {-|

    The memory subsystem associated with raaz.

    __Warning:__ This module is pretty low level and should not be needed in typical
    use cases. Only developers of protocols and primitives might have a
    reason to look into this module.

    -}

    {-# LANGUAGE TypeFamilies               #-}
    {-# LANGUAGE GADTs                      #-}
    {-# LANGUAGE RankNTypes                 #-}
    {-# LANGUAGE MultiParamTypeClasses      #-}
    {-# LANGUAGE FlexibleInstances          #-}

    module Raaz.Core.Memory
           (
           -- * The Memory subsystem.
           -- $memorysubsystem$

           -- ** Memory elements.
             Memory(..), VoidMemory, copyMemory
           -- *** Initialisation and Extraction.
           -- $init-extract$
           , Initialisable(..), Extractable(..)
           , InitialisableFromBuffer(..), ExtractableToBuffer(..)
           -- *** A basic memory cell.
           , MemoryCell, withCellPointer, getCellPointer

           -- ** Memory threads.
           , MemoryThread(..), doIO , getMemory, modify, execute
           , MT
           -- **** Some low level `MT` actions.

           , liftPointerAction

           -- ** Memory allocation
           ,  Alloc, pointerAlloc
           ) where

    import           Control.Applicative
    import           Control.Monad.IO.Class
    import           Foreign.Storable            ( Storable )
    import           Foreign.Ptr                 ( castPtr, Ptr )
    import           Raaz.Core.MonoidalAction
    import           Raaz.Core.Transfer
    import           Raaz.Core.Types

    -- $memorysubsystem$
    --
    -- Cryptographic operations often need to keep sensitive information
    -- like private keys in its memory space. Such sensitive information
    -- can leak to the external would if the memory where the data is
    -- stored is swapped out to a disk. What makes this particularly
    -- dangerous is that the data can reside on the disk almost
    -- permanently and might even survive when the hardware is
    -- scrapped. The primary purpose of the memory subsystem is to provide
    -- a way to allocate and manage /secure memory/, i.e. memory that will
    -- not be swapped out as long as the memory is used and will be wiped
    -- clean after use. It consists of the following components:
    --
    -- [The `Memory` type class:] A memory element is some type that holds
    -- an internal buffer inside it.
    --
    -- [The `Alloc` type:] Memory elements need to be allocated and this
    -- is involves a lot of low lever pointer arithmetic. The `Alloc`
    -- types gives a high level interface for memory allocation. For a
    -- memory type `mem`, the type `Alloc mem` can be seen as the
    -- _allocation strategy_ for mem. For example, one of the things that
    -- it keeps track of is the space required to create an memory element
    -- of type `mem`. There is a natural applicative instance for `Alloc`
    -- which helps build the allocation strategy for a compound memory
    -- type from its components in a modular fashion _without_ explicit
    -- size calculation or offset computation.
    --
    -- [`MemoryThread`s:] Instances of this class are actions that use
    -- some kind of memory elements inside it. Such a thread can be run
    -- using the combinator `securely` or the combinator `insecurely`. If
    -- one use the combinator `securely`, then the allocation of the
    -- memory element to be used by the action is done using a locked
    -- memory pool which is wiped clean before de-allocation.

    -- $init-extract$
    --
    -- Memory elements often needs to be initialised. Similarly data needs
    -- to be extracted out of memory. An instance declaration
    -- @`Initialisable` mem a@ for the memory type @mem@ indicates that it
    -- can be initialised with the pure value @a@. Similary, if values of
    -- type @b@ can be extracted out of a memory element @mem@, we can
    -- indicate it with an instance of @`Extractable` mem a@.
    --
    -- There is an inherent danger in initialising and extracting pure
    -- values out of memory. Pure values are stored on the Haskell heap
    -- and hence can be swapped out. Consider a memory element @mem@ that
    -- stores some sensitive information, say for example the unencrypted
    -- private key. Suppose we extract this key out of the memory element
    -- as a pure value before its encryption and storage into the key
    -- file. It is likely that the key is swapped out to the disk as the
    -- extracted key is part of the the haskell heap.
    --
    -- The `InitialiseFromBuffer` (`ExtractableToBuffer`) class gives an
    -- interface for reading from (writing to) buffers directly minimising
    -- the chances of inadvertent exposure of sensitive information from
    -- the Haskell heap due to swapping.

    -- | An action of type @`MT` mem a@ is an action that uses internally
    -- a single memory object of type @mem@ and returns a result of type
    -- @a@. All the actions are performed on a single memory element and
    -- hence the side effects persist. It is analogues to the @ST@ monad.
    newtype MT mem a = MT { unMT :: mem -> IO a }

    instance Functor (MT mem) where
      fmap f mst = MT $ \ m -> f <$> unMT mst m

    instance Applicative (MT mem) where
      pure       = MT . const . pure
      mf <*> ma  = MT $ \ m -> unMT mf m <*> unMT ma m

    instance Monad (MT mem) where
      return    =  MT . const . return
      ma >>= f  =  MT runIt
        where runIt mem = unMT ma mem >>= \ a -> unMT (f a) mem

    -- | __WARNING:__ do not lift a secure memory action.
    instance MonadIO (MT mem) where
      liftIO = MT . const

    -- | A class that captures abstract "memory threads". A memory thread
    -- can either be run `securely` or `insecurely`. Pure IO actions can
    -- be run inside a memory thread using the @runIO@.  However, the IO
    -- action that is being run /must not/ directly or indirectly run a
    -- `secure` action ever. In particular, the following code is bad.
    --
    -- > -- BAD EXAMPLE: DO NOT USE.
    -- > runIO $ securely $ foo
    -- >
    --
    -- On the other hand the following code is fine
    --
    -- >
    -- > runIO $ insecurely $ someMemoryAction
    -- >
    --
    -- As to why this is dangerous, it has got to do with the fact that
    -- @mlock@ and @munlock@ do not nest correctly. A single @munlock@ can
    -- unlock multiple calls of @mlock@ on the same page.  Whether a given
    -- @IO@ action unlocks memory is difficult to keep track of; for all
    -- you know, it might be a FFI call that does an @memunlock@. Hence,
    -- currently there is no easy way to enforce this.
    --

    class MemoryThread (mT :: * -> * -> *) where
      -- | Run a memory action with the internal memory allocated from a
      -- locked memory buffer. This memory buffer will never be swapped
      -- out by the operating system and will be wiped clean before
      -- releasing.
      --
      -- Memory locking is an expensive operation and usually there would be
      -- a limit to how much locked memory can be allocated. Nonetheless,
      -- actions that work with sensitive information like passwords should
      -- use this to run an memory action.
      securely   :: Memory mem => mT mem a -> IO a

      -- | Run a memory action with the internal memory used by the action
      -- being allocated from unlocked memory. Use this function when you
      -- work with data that is not sensitive to security considerations
      -- (for example, when you want to verify checksums of files).
      insecurely :: Memory mem => mT mem a -> IO a

      -- | Lift an actual memory thread.
      liftMT :: MT mem a -> mT mem a

      -- | Combinator that allows us to run a memory action on a
      -- sub-memory element. A sub-memory of @submem@ of a memory element
      -- @mem@ is given by a projection @proj : mem -> submem@. The action
      -- @onSubMemory proj@ lifts the a memory thread on the sub element
      -- to the compound element.
      --
      onSubMemory :: (mem -> submem) -> mT submem a -> mT mem a

    instance MemoryThread MT where
      securely               = withSecureMemory . unMT
      insecurely             = withMemory . unMT
      liftMT                 = id
      onSubMemory proj mtsub = MT $ unMT mtsub . proj

    ------------- Lifting pointer actions -----------------------------

    -- | Run a given memory action in the memory thread.
    execute :: MemoryThread mT => (mem -> IO a) -> mT mem a
    execute = liftMT . MT

    -- | Perform an IO action inside the memory thread.
    doIO :: MemoryThread mT => IO a -> mT mem a
    doIO = execute . const

    -- | A pointer action inside a monad @m@ is some function that takes a
    -- pointer action of type @Pointer -> m a@ and supplies it with an
    -- appropriate pointer. In particular, memory allocators are pointer
    -- actions.
    type PointerAction m a b = (Pointer -> m a) -> m b

    -- | An IO allocator can be lifted to the memory thread level as follows.
    liftPointerAction :: PointerAction IO a b -> PointerAction (MT mem) a b
    liftPointerAction allocator mtAction
      = execute $ \ mem -> allocator (\ ptr -> unMT (mtAction ptr) mem)

    -- TODO: This is a very general pattern needs more exploration.

    -- | Get the underlying memory element of the memory thread.
    getMemory :: MemoryThread mT => mT mem mem
    getMemory = execute return

    ------------------------ A memory allocator -----------------------

    type AllocField = Field Pointer

    -- | A memory allocator for the memory type @mem@. The `Applicative`
    -- instance of @Alloc@ can be used to build allocations for
    -- complicated memory elements from simpler ones.
    type Alloc mem = TwistRF AllocField (BYTES Int) mem

    -- | Make an allocator for a given memory type.
    makeAlloc :: LengthUnit l => l -> (Pointer -> mem) -> Alloc mem
    makeAlloc l memCreate = TwistRF (WrapArrow memCreate) $ atLeast l

    -- | Allocates a buffer of size @l@ and returns the pointer to it pointer.
    pointerAlloc :: LengthUnit l => l -> Alloc Pointer
    pointerAlloc l = makeAlloc l id

    ---------------------------------------------------------------------

    -- | Any cryptographic primitives use memory to store stuff. This
    -- class abstracts all types that hold some memory. Cryptographic
    -- application often requires securing the memory from being swapped
    -- out (think of memory used to store private keys or passwords). This
    -- abstraction supports memory securing. If your platform supports
    -- memory locking, then securing a memory will prevent the memory from
    -- being swapped to the disk. Once secured the memory location is
    -- overwritten by nonsense before being freed.
    --
    -- While some basic memory elements like `MemoryCell` are exposed from
    -- the library, often we require compound memory objects built out of
    -- simpler ones. The `Applicative` instance of the `Alloc` can be made
    -- use of in such situation to simplify such instance declaration as
    -- illustrated in the instance declaration for a pair of memory
    -- elements.
    --
    -- > instance (Memory ma, Memory mb) => Memory (ma, mb) where
    -- >
    -- >    memoryAlloc             = (,) <$> memoryAlloc <*> memoryAlloc
    -- >
    -- >    unsafeToPointer (ma, _) =  unsafeToPointer ma
    --
    class Memory m where

      -- | Returns an allocator for this memory.
      memoryAlloc     :: Alloc m

      -- | Returns the pointer to the underlying buffer.
      unsafeToPointer :: m -> Pointer

    -- | A memory element that holds nothing.
    data VoidMemory = VoidMemory { unVoidMemory :: Pointer  }

    instance Memory VoidMemory where
      memoryAlloc      = makeAlloc (0 :: BYTES Int) VoidMemory
      unsafeToPointer  = unVoidMemory

    instance ( Memory ma, Memory mb ) => Memory (ma, mb) where
        memoryAlloc             = (,) <$> memoryAlloc <*> memoryAlloc
        unsafeToPointer (ma, _) =  unsafeToPointer ma

    instance ( Memory ma
             , Memory mb
             , Memory mc
             )
             => Memory (ma, mb, mc) where
      memoryAlloc              = (,,)
                                 <$> memoryAlloc
                                 <*> memoryAlloc
                                 <*> memoryAlloc
      unsafeToPointer (ma,_,_) =  unsafeToPointer ma

    instance ( Memory ma
             , Memory mb
             , Memory mc
             , Memory md
             )
             => Memory (ma, mb, mc, md) where
      memoryAlloc                = (,,,)
                                   <$> memoryAlloc
                                   <*> memoryAlloc
                                   <*> memoryAlloc
                                   <*> memoryAlloc

      unsafeToPointer (ma,_,_,_) =  unsafeToPointer ma

    -- | Copy data from a given memory location to the other. The first
    -- argument is destionation and the second argument is source to match
    -- with the convention followed in memcpy.
    copyMemory :: Memory m => Dest m -- ^ Destination
                           -> Src  m -- ^ Source
                           -> IO ()
    copyMemory dmem smem = memcpy (unsafeToPointer <$> dmem) (unsafeToPointer <$> smem) sz
      where sz       = twistMonoidValue $ getAlloc smem
            getAlloc :: Memory m => Src m -> Alloc m
            getAlloc _ = memoryAlloc

    -- | Perform an action which makes use of this memory. The memory
    -- allocated will automatically be freed when the action finishes
    -- either gracefully or with some exception. Besides being safer,
    -- this method might be more efficient as the memory might be
    -- allocated from the stack directly and will have very little GC
    -- overhead.
    withMemory   :: Memory m => (m -> IO a) -> IO a
    withMemory   = withM memoryAlloc
      where withM :: Alloc m -> (m -> IO a) -> IO a
            withM alctr action = allocaBuffer sz actualAction
              where sz                 = twistMonoidValue alctr
                    getM               = computeField $ twistFunctorValue alctr
                    wipeIt cptr        = memset cptr 0 sz
                    actualAction  cptr = action (getM cptr) <* wipeIt cptr

    -- | Similar to `withMemory` but allocates a secure memory for the
    -- action. Secure memories are never swapped on to disk and will be
    -- wiped clean of sensitive data after use. However, be careful when
    -- using this function in a child thread. Due to the daemonic nature
    -- of Haskell threads, if the main thread exists before the child
    -- thread is done with its job, sensitive data can leak. This is
    -- essentially a limitation of the bracket which is used internally.
    withSecureMemory :: Memory m => (m -> IO a) -> IO a
    withSecureMemory = withSM memoryAlloc
      where -- withSM :: Memory m => Alloc m -> (m -> IO a) -> IO a
            withSM alctr action = allocaSecure sz $ action . getM
              where sz     = twistMonoidValue alctr
                    getM   = computeField $ twistFunctorValue alctr

    ----------------------- Initialising and Extracting stuff ----------------------

    -- | Memories that can be initialised with a pure value. The pure
    -- value resides in the Haskell heap and hence can potentially be
    -- swapped. Therefore, this class should be avoided if compromising
    -- the initialisation value can be dangerous. Consider using
    -- `InitialiseableFromBuffer`
    --

    class Memory m => Initialisable m v where
      initialise :: v -> MT m ()

    -- | Memories from which pure values can be extracted. Once a pure value is
    -- extracted,
    class Memory m => Extractable m v where
      extract  :: MT m v

    -- | Apply the given function to the value in the cell. For a function @f :: b -> a@,
    -- the action @modify f@ first extracts a value of type @b@ from the
    -- memory element, applies @f@ to it and puts the result back into the
    -- memory.
    --
    -- > modify f = do b          <- extract
    -- >               initialise $  f b
    --
    modify :: (Initialisable mem a, Extractable mem b, MemoryThread mT ) =>  (b -> a) -> mT mem ()
    modify f = liftMT $ extract >>= initialise . f

    -- | A memory type that can be initialised from a pointer buffer. The initialisation performs
    -- a direct copy from the input buffer and hence the chances of the
    -- initialisation value ending up in the swap is minimised.
    class Memory m => InitialisableFromBuffer m where
      initialiser :: m -> ReadM (MT m)

    -- | A memory type that can extract bytes into a buffer. The extraction will perform
    -- a direct copy and hence the chances of the extracted value ending
    -- up in the swap space is minimised.
    class Memory m => ExtractableToBuffer m where
      extractor :: m -> WriteM (MT m)

    --------------------- Some instances of Memory --------------------

    -- | A memory location to store a value of type having `Storable`
    -- instance.
    newtype MemoryCell a = MemoryCell { unMemoryCell :: Ptr a }

    instance Storable a => Memory (MemoryCell a) where

      memoryAlloc = allocator undefined
        where allocator :: Storable b => b -> Alloc (MemoryCell b)
              allocator b = makeAlloc (alignedSizeOf b) $ MemoryCell . castPtr

      unsafeToPointer  = castPtr . unMemoryCell

    -- | The location where the actual storing of element happens. This
    -- pointer is guaranteed to be aligned to the alignment restriction of @a@
    actualCellPtr :: Storable a => MemoryCell a -> Ptr a
    actualCellPtr = nextAlignedPtr . unMemoryCell

    -- | Work with the underlying pointer of the memory cell. Useful while
    -- working with ffi functions.
    withCellPointer :: (MemoryThread mT, Storable a) => (Ptr a -> IO b) -> mT (MemoryCell a) b
    {-# INLINE withCellPointer #-}
    withCellPointer action = execute $ action . actualCellPtr

    -- | Get the pointer associated with the given memory cell.
    getCellPointer :: (MemoryThread mT, Storable a) => mT (MemoryCell a) (Ptr a)
    {-# INLINE getCellPointer #-}
    getCellPointer = liftMT $ actualCellPtr <$> getMemory

    instance Storable a => Initialisable (MemoryCell a) a where
      initialise a = execute $ flip pokeAligned a . unMemoryCell
      {-# INLINE initialise #-}

    instance Storable a => Extractable (MemoryCell a) a where
      extract = execute $ peekAligned . unMemoryCell
      {-# INLINE extract #-}

    instance EndianStore a => InitialisableFromBuffer (MemoryCell a) where
      initialiser  = readInto 1 . destination . actualCellPtr

    instance EndianStore a => ExtractableToBuffer (MemoryCell a) where
      extractor  = writeFrom 1 . source . actualCellPtr