ByteSource.hs

{-# LANGUAGE FlexibleContexts  #-}
-- |
--
-- Module      : Raaz.Core.ByteSource
-- Copyright   : (c) Piyush P Kurur, 2019
-- License     : Apache-2.0 OR BSD-3-Clause
-- Maintainer  : Piyush P Kurur <ppk@iitpkd.ac.in>
-- Stability   : experimental
--

module Raaz.Core.ByteSource
       ( -- * Byte sources.
         -- $bytesource$

         ByteSource(..), PureByteSource
       --    InfiniteSource(..)
       , FillResult(..)
       , fill, processChunks
       , withFillResult
       ) where

import           Control.Monad.IO.Class
import qualified Data.ByteString      as B
import qualified Data.ByteString.Lazy as L

import           Raaz.Core.Prelude
import           Raaz.Core.Types      (BYTES, Ptr, LengthUnit (..), movePtr)
import           Raaz.Core.Util.ByteString( unsafeCopyToPointer
                                          , unsafeNCopyToPointer
                                          , length
                                          )
import           Raaz.Core.Types.Pointer  (hFillBuf, Pointer, unsafeWithPointer)

-- $bytesource$
--
-- Cryptographic input come from various sources; they can come from
-- network sockets or might be just a string in the Haskell. To give a
-- uniform interfaces for all such inputs, we define the abstract
-- concept of a /byte source/. Essentially a byte source is one from
-- which we can fill a buffer with bytes.
--
-- Among instances of `ByteSource`, some like for example
-- `B.ByteString` are /pure/ in the sense filling a buffer with bytes
-- from such a source has no other side-effects. This is in contrast
-- to a source like a sockets. The type class `PureByteSource`
-- captures such byte sources.
--

-- | This type captures the result of a fill operation.
data FillResult a = Remaining a           -- ^ There is still bytes left.
                  | Exhausted (BYTES Int) -- ^ source exhausted with so much
                                          -- bytes read.
                    deriving (Show, Eq)

instance Functor FillResult where
  fmap f (Remaining a ) = Remaining $ f a
  fmap _ (Exhausted sz) = Exhausted sz

-- | Combinator to handle a fill result.
withFillResult :: (a -> b)          -- ^ stuff to do when filled
               -> (BYTES Int -> b)  -- ^ stuff to do when exhausted
               -> FillResult a      -- ^ the fill result to process
               -> b
withFillResult continueWith _     (Remaining a)  = continueWith a
withFillResult _            endBy (Exhausted sz) = endBy sz

------------------------ Byte sources ----------------------------------

-- | Abstract byte sources. A bytesource is something that you can use
-- to fill a buffer.
--
--  __WARNING:__ The source is required to return `Exhausted` in the
-- boundary case where it has exactly the number of bytes
-- requested. In other words, if the source returns @Remaining@ on any
-- particular request, there should be at least 1 additional byte left
-- on the source for the next request. Cryptographic block primitives
-- perform certain special processing, like padding for example, for
-- the last block and it is required to know whether the last block
-- has been read or not.
class ByteSource src where
  -- | Fills a buffer from the source.
  fillBytes :: BYTES Int  -- ^ Buffer size
            -> src        -- ^ The source to fill.
            -> Ptr a      -- ^ Buffer pointer
            -> IO (FillResult src)

-- | A version of fillBytes that takes type safe lengths as input.
fill :: ( Pointer ptr
        , LengthUnit len
        , ByteSource src
        )
     => len
     -> src
     -> ptr a
     -> IO (FillResult src)
fill len src = unsafeWithPointer $ fillBytes (inBytes len) src
{-# INLINE fill #-}

-- | Process data from a source in chunks of a particular size.
processChunks :: ( Pointer ptr, MonadIO m, LengthUnit chunkSize, ByteSource src)
              => m a                 -- ^ action on a complete chunk,
              -> (BYTES Int -> m b)  -- ^ action on the last partial chunk,
              -> src                 -- ^ the source
              -> ptr something       -- ^ buffer to fill the chunk in
              -> chunkSize           -- ^ size of the chunksize
              -> m b
processChunks mid end source ptr csz = go source
  where fillChunk src = liftIO $ fill csz src ptr
        step src      = mid >> go src
        go src        = fillChunk src >>= withFillResult step end


-- | A byte source src is pure if filling from it does not have any
-- other side effect on the state of the byte source. Formally, two
-- different fills form the same source should fill the buffer with
-- the same bytes.  This additional constraint on the source helps to
-- /purify/ certain crypto computations like computing the hash or mac
-- of the source. Usualy sources like `B.ByteString` etc are pure byte
-- sources. A file handle is a byte source that is /not/ a pure
-- source.
class ByteSource src => PureByteSource src where

----------------------- Instances of byte source -----------------------

-- | __WARNING:__ The `fillBytes` may block.
instance ByteSource Handle where
  {-# INLINE fillBytes #-}
  fillBytes sz hand cptr = do
    count <- hFillBuf hand cptr sz
    eof   <- hIsEOF hand
    if eof then return $ Exhausted count
      else return $ Remaining hand

instance ByteSource B.ByteString where
  {-# INLINE fillBytes #-}
  fillBytes sz bs cptr | l <= sz    = do unsafeCopyToPointer bs cptr
                                         return $ Exhausted l
                       | otherwise = do unsafeNCopyToPointer sz bs cptr
                                        return $ Remaining rest
       where l    = length bs
             rest = B.drop (fromIntegral sz) bs

instance ByteSource L.ByteString where
  {-# INLINE fillBytes #-}
  fillBytes sz bs = fmap (fmap L.fromChunks) . fillBytes sz (L.toChunks bs)

instance ByteSource src => ByteSource (Maybe src) where
  {-# INLINE fillBytes #-}
  fillBytes sz ma cptr = maybe exhausted fillIt ma
          where exhausted = return $ Exhausted 0
                fillIt a  = fmap Just <$> fillBytes sz a cptr

instance ByteSource src => ByteSource [src] where
  fillBytes _  []     _    = return $ Exhausted 0
  fillBytes sz (x:xs) cptr = do
    result <- fillBytes sz x cptr
    case result of
      Remaining nx     -> return $ Remaining $ nx:xs
      Exhausted bytesX -> let nptr              = movePtr cptr bytesX
                              whenXSExhausted bytesXS = return $ Exhausted $ bytesX + bytesXS
                              whenXSRemains           = return . Remaining
                           in fillBytes (sz - bytesX) xs nptr
                              >>= withFillResult whenXSRemains whenXSExhausted


--------------------- Instances of pure byte source --------------------

instance PureByteSource B.ByteString where
instance PureByteSource L.ByteString where
instance PureByteSource src => PureByteSource [src]
instance PureByteSource src => PureByteSource (Maybe src)