{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_HADDOCK prune not-home #-}

{- |
Module : Data.Attoparsec.Framer
Copyright : (c) 2022 Tim Emiola
Maintainer : Tim Emiola <adetokunbo@emio.la>
SPDX-License-Identifier: BSD3

Provides the 'Framer' data type that combines an @Attoparsec 'A.Parser'@ with a
a few additional combinators that allow the parser to be used to process frames
from the framed byte streams commonly used in network protocol implementations.

A @'Framer'@ specifies how the processing function @'runFramer'@ should
parse a byte stream.

Minimally, a @Framer@ specifies

* An @'A.Parser'@, used to extract frames from the byte stream
* a @'FrameHandler'@ responsible using the parsed frames
* the bytestream source, represented a 'ByteSource' 


@'runFramer'@ the 'FrameHandler' is invoked repeatedly; on each
invocation it returns a 'Progression', which indicates if processing should
continue. This makes it possible to terminate for the 'FrameHandler' to signal
that frame processing should terminate.
-}
module Data.Attoparsec.Framer (
  -- * Framer
  ByteSource,
  Framer,
  FrameHandler,
  Progression (..),
  mkFramer,
  mkFramer',

  -- * query/update a  @'Framer'@
  setChunkSize,
  setOnBadParse,
  setOnClosed,
  setOnFrame,
  chunkSize,

  -- * Run the @Framer@
  runFramer,
  runOneFrame,
  
  -- * Exception handling
  -- $exceptions

  -- * exceptions
  BrokenFrame (..),
  NoMoreInput (..),
) where

import Control.Exception (Exception)
import Control.Monad.Catch (MonadThrow (..))
import qualified Data.Attoparsec.ByteString as A
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Word (Word32)


-- | Handles a parsed @frame@, returning a @Progression@ that indicates if further @frames@ should be parsed.
type FrameHandler m frame = frame -> m Progression

-- | A byte stream from which chunks are to be repeatedly retrieved.
type ByteSource m = Word32 -> m ByteString

-- | Used by 'FrameHandler' to indicate if additional frames should be parsed.
data Progression
   = Stop
  | StopUnlessExtra
  | Continue
  deriving (Progression -> Progression -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Progression -> Progression -> Bool
$c/= :: Progression -> Progression -> Bool
== :: Progression -> Progression -> Bool
$c== :: Progression -> Progression -> Bool
Eq, Int -> Progression -> ShowS
[Progression] -> ShowS
Progression -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Progression] -> ShowS
$cshowList :: [Progression] -> ShowS
show :: Progression -> String
$cshow :: Progression -> String
showsPrec :: Int -> Progression -> ShowS
$cshowsPrec :: Int -> Progression -> ShowS
Show)


-- | Use 'A.Parser' to parse a stream of @frames@ from a bytestream
data Framer m frame = Framer
  { forall (m :: * -> *) frame. Framer m frame -> Word32
framerChunkSize :: !Word32
  , forall (m :: * -> *) frame. Framer m frame -> ByteSource m
frameByteSource :: !(ByteSource m)
  , forall (m :: * -> *) frame. Framer m frame -> FrameHandler m frame
framerOnFrame :: !(FrameHandler m frame)
  , forall (m :: * -> *) frame. Framer m frame -> Parser frame
framerParser :: !(A.Parser frame)
  , forall (m :: * -> *) frame. Framer m frame -> m ()
framerOnClosed :: !(m ())
  , forall (m :: * -> *) frame. Framer m frame -> Text -> m ()
framerOnBadParse :: !(Text -> m ())
  }


{- | Construct @'Framer'@ that will handle @frames@ repeatedly until a returned
 @'Progression'@ stops it.
-}
mkFramer' ::
  MonadThrow m =>
  A.Parser frame ->
  FrameHandler m frame ->
  ByteSource m ->
  Framer m frame
mkFramer' :: forall (m :: * -> *) frame.
MonadThrow m =>
Parser frame
-> FrameHandler m frame -> ByteSource m -> Framer m frame
mkFramer' Parser frame
framerParser FrameHandler m frame
framerOnFrame ByteSource m
frameByteSource =
  Framer
    { Parser frame
framerParser :: Parser frame
framerParser :: Parser frame
framerParser
    , FrameHandler m frame
framerOnFrame :: FrameHandler m frame
framerOnFrame :: FrameHandler m frame
framerOnFrame
    , ByteSource m
frameByteSource :: ByteSource m
frameByteSource :: ByteSource m
frameByteSource
    , framerOnBadParse :: Text -> m ()
framerOnBadParse = \Text
_err -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    , framerOnClosed :: m ()
framerOnClosed = forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM NoMoreInput
NoMoreInput
    , framerChunkSize :: Word32
framerChunkSize = Word32
defaultChunkSize
    }


-- | Construct @'Framer'@ that loops continuously.
mkFramer ::
  MonadThrow m =>
  A.Parser a ->
  (a -> m ()) ->
  (Word32 -> m ByteString) ->
  Framer m a
mkFramer :: forall (m :: * -> *) a.
MonadThrow m =>
Parser a -> (a -> m ()) -> (Word32 -> m ByteString) -> Framer m a
mkFramer Parser a
parser a -> m ()
onFrame Word32 -> m ByteString
fetchBytes =
  let onFrameContinue :: a -> m Progression
onFrameContinue a
x = do
        a -> m ()
onFrame a
x
        forall (f :: * -> *) a. Applicative f => a -> f a
pure Progression
Continue
   in forall (m :: * -> *) frame.
MonadThrow m =>
Parser frame
-> FrameHandler m frame -> ByteSource m -> Framer m frame
mkFramer' Parser a
parser a -> m Progression
onFrameContinue Word32 -> m ByteString
fetchBytes


-- | Repeatedly parse and handle frames until the configured @FrameHandler@ ends handling.
runFramer ::
  MonadThrow m =>
  Framer m a ->
  m ()
runFramer :: forall (m :: * -> *) a. MonadThrow m => Framer m a -> m ()
runFramer Framer m a
f =
  let Framer
        { framerChunkSize :: forall (m :: * -> *) frame. Framer m frame -> Word32
framerChunkSize = Word32
fetchSize
        , framerOnBadParse :: forall (m :: * -> *) frame. Framer m frame -> Text -> m ()
framerOnBadParse = Text -> m ()
onErr
        , frameByteSource :: forall (m :: * -> *) frame. Framer m frame -> ByteSource m
frameByteSource = ByteSource m
fetchBytes
        , framerOnFrame :: forall (m :: * -> *) frame. Framer m frame -> FrameHandler m frame
framerOnFrame = FrameHandler m a
onFrame
        , framerParser :: forall (m :: * -> *) frame. Framer m frame -> Parser frame
framerParser = Parser a
parser
        , framerOnClosed :: forall (m :: * -> *) frame. Framer m frame -> m ()
framerOnClosed = m ()
onClosed
        } = Framer m a
f
   in forall (m :: * -> *) a.
MonadThrow m =>
Word32
-> Parser a
-> (Word32 -> m ByteString)
-> (a -> m Progression)
-> (Text -> m ())
-> m ()
-> m ()
runFramer' Word32
fetchSize Parser a
parser ByteSource m
fetchBytes FrameHandler m a
onFrame Text -> m ()
onErr m ()
onClosed


{- | Parse and handle a single frame.

The result is tuple of the outstanding unparsed bytes from the bytestream if
any, and a value indicating if the bytestream has terminated.
-}
runOneFrame ::
  MonadThrow m =>
  Maybe ByteString ->
  Framer m a ->
  m ((Maybe ByteString), Bool)
runOneFrame :: forall (m :: * -> *) a.
MonadThrow m =>
Maybe ByteString -> Framer m a -> m (Maybe ByteString, Bool)
runOneFrame Maybe ByteString
restMb Framer m a
f =
  let Framer
        { framerChunkSize :: forall (m :: * -> *) frame. Framer m frame -> Word32
framerChunkSize = Word32
fetchSize
        , framerOnBadParse :: forall (m :: * -> *) frame. Framer m frame -> Text -> m ()
framerOnBadParse = Text -> m ()
onErr
        , frameByteSource :: forall (m :: * -> *) frame. Framer m frame -> ByteSource m
frameByteSource = ByteSource m
fetchBytes
        , framerOnFrame :: forall (m :: * -> *) frame. Framer m frame -> FrameHandler m frame
framerOnFrame = FrameHandler m a
onFrame
        , framerParser :: forall (m :: * -> *) frame. Framer m frame -> Parser frame
framerParser = Parser a
parser
        , framerOnClosed :: forall (m :: * -> *) frame. Framer m frame -> m ()
framerOnClosed = m ()
onClose
        } = Framer m a
f
   in forall (m :: * -> *) a.
MonadThrow m =>
Maybe ByteString
-> Word32
-> Parser a
-> (Word32 -> m ByteString)
-> (a -> m Progression)
-> (Text -> m ())
-> m ()
-> m (Maybe ByteString, Bool)
runOneFrame' Maybe ByteString
restMb Word32
fetchSize Parser a
parser ByteSource m
fetchBytes FrameHandler m a
onFrame Text -> m ()
onErr m ()
onClose


-- | The chunk size of a @Framer@.
chunkSize :: Framer m a -> Word32
chunkSize :: forall (m :: * -> *) frame. Framer m frame -> Word32
chunkSize = forall (m :: * -> *) frame. Framer m frame -> Word32
framerChunkSize


-- | Update the chunk size of a @Framer@.
setChunkSize :: Word32 -> Framer m a -> Framer m a
setChunkSize :: forall (m :: * -> *) a. Word32 -> Framer m a -> Framer m a
setChunkSize Word32
size Framer m a
f = Framer m a
f {framerChunkSize :: Word32
framerChunkSize = Word32
size}


-- | Update the parse error handler of a @Framer@.
setOnBadParse :: (Text -> m ()) -> Framer m a -> Framer m a
setOnBadParse :: forall (m :: * -> *) a. (Text -> m ()) -> Framer m a -> Framer m a
setOnBadParse Text -> m ()
onErr Framer m a
f = Framer m a
f {framerOnBadParse :: Text -> m ()
framerOnBadParse = Text -> m ()
onErr}


-- | Update the @FrameHandler@ of a @Framer@.
setOnFrame :: FrameHandler m frame -> Framer m frame -> Framer m frame
setOnFrame :: forall (m :: * -> *) frame.
FrameHandler m frame -> Framer m frame -> Framer m frame
setOnFrame FrameHandler m frame
onFrame Framer m frame
f = Framer m frame
f {framerOnFrame :: FrameHandler m frame
framerOnFrame = FrameHandler m frame
onFrame}


-- | Update the end-of-input handler of a @Framer@.
setOnClosed :: (m ()) -> Framer m a -> Framer m a
setOnClosed :: forall (m :: * -> *) a. m () -> Framer m a -> Framer m a
setOnClosed m ()
onClose Framer m a
f = Framer m a
f {framerOnClosed :: m ()
framerOnClosed = m ()
onClose}


runFramer' ::
  MonadThrow m =>
  Word32 ->
  A.Parser a ->
  (Word32 -> m ByteString) ->
  (a -> m Progression) ->
  (Text -> m ()) ->
  m () ->
  m ()
runFramer' :: forall (m :: * -> *) a.
MonadThrow m =>
Word32
-> Parser a
-> (Word32 -> m ByteString)
-> (a -> m Progression)
-> (Text -> m ())
-> m ()
-> m ()
runFramer' Word32
fetchSize Parser a
parser Word32 -> m ByteString
fetchBytes a -> m Progression
handleFrame Text -> m ()
onErr m ()
onClosed = do
  let loop :: Maybe ByteString -> m ()
loop Maybe ByteString
x = do
        (Maybe ByteString
next, Bool
closed) <- forall (m :: * -> *) a.
MonadThrow m =>
Maybe ByteString
-> Word32
-> Parser a
-> (Word32 -> m ByteString)
-> (a -> m Progression)
-> (Text -> m ())
-> m ()
-> m (Maybe ByteString, Bool)
runOneFrame' Maybe ByteString
x Word32
fetchSize Parser a
parser Word32 -> m ByteString
fetchBytes a -> m Progression
handleFrame Text -> m ()
onErr m ()
onClosed
        if Bool -> Bool
not Bool
closed then Maybe ByteString -> m ()
loop Maybe ByteString
next else forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  Maybe ByteString -> m ()
loop forall a. Maybe a
Nothing


runOneFrame' ::
  MonadThrow m =>
  Maybe ByteString ->
  Word32 ->
  A.Parser a ->
  (Word32 -> m ByteString) ->
  (a -> m Progression) ->
  (Text -> m ()) ->
  m () ->
  m ((Maybe ByteString), Bool)
runOneFrame' :: forall (m :: * -> *) a.
MonadThrow m =>
Maybe ByteString
-> Word32
-> Parser a
-> (Word32 -> m ByteString)
-> (a -> m Progression)
-> (Text -> m ())
-> m ()
-> m (Maybe ByteString, Bool)
runOneFrame' Maybe ByteString
restMb Word32
fetchSize Parser a
parser Word32 -> m ByteString
fetchBytes a -> m Progression
handleFrame Text -> m ()
onErr m ()
onClose = do
  let pullChunk :: m ByteString
pullChunk = Word32 -> m ByteString
fetchBytes Word32
fetchSize
      initial :: ByteString
initial = forall a. a -> Maybe a -> a
fromMaybe ByteString
BS.empty Maybe ByteString
restMb
      onParse :: IResult ByteString a -> m (Maybe ByteString, Bool)
onParse (A.Fail ByteString
_ [String]
ctxs String
reason) = do
        let errMessage :: Text
errMessage = [String] -> String -> Text
parsingFailed [String]
ctxs String
reason
        if String
reason forall a. Eq a => a -> a -> Bool
== String
closedReason
          then -- WANTED: a typed way of detecting this condition, i.e,
          -- it is possible not to rely on a  specific error message ?
          do
            m ()
onClose
            forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Maybe a
Nothing, Bool
True)
          else do
            Text -> m ()
onErr Text
errMessage
            forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ String -> BrokenFrame
BrokenFrame String
reason
      onParse (A.Done ByteString
i a
r) = do
        let extraMb :: Maybe ByteString
extraMb = if ByteString -> Bool
BS.null ByteString
i then forall a. Maybe a
Nothing else forall a. a -> Maybe a
Just ByteString
i
        Progression
doMore <- a -> m Progression
handleFrame a
r
        case (Progression
doMore, Maybe ByteString
extraMb) of
          (Progression
Stop, Maybe ByteString
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe ByteString
extraMb, Bool
True)
          (Progression
StopUnlessExtra, Maybe ByteString
Nothing) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe ByteString
extraMb, Bool
True)
          (Progression
_, Maybe ByteString
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe ByteString
extraMb, Bool
False)
      onParse (A.Partial ByteString -> IResult ByteString a
continue) = m ByteString
pullChunk forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IResult ByteString a -> m (Maybe ByteString, Bool)
onParse forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> IResult ByteString a
continue
  forall (m :: * -> *) a.
Monad m =>
m ByteString -> Parser a -> ByteString -> m (Result a)
A.parseWith m ByteString
pullChunk Parser a
parser ByteString
initial forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IResult ByteString a -> m (Maybe ByteString, Bool)
onParse


parsingFailed :: [String] -> String -> Text
parsingFailed :: [String] -> String -> Text
parsingFailed [String]
context String
reason =
  let contexts :: Text
contexts = Text -> [Text] -> Text
Text.intercalate Text
"-" (String -> Text
Text.pack forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String]
context)
      cause :: Text
cause = if forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
reason then Text
Text.empty else Text
":" forall a. Semigroup a => a -> a -> a
<> String -> Text
Text.pack String
reason
   in Text
"bad parse:" forall a. Semigroup a => a -> a -> a
<> Text
contexts forall a. Semigroup a => a -> a -> a
<> Text
cause


{- $exceptions

On failures, @'runFramer'@ throws @'Exception's@ using @'MonadThrow'@ rather
than using an @Either@ or @MonadError@

This is because it is intended to be used to parse framed protocol byte streams;
where parsing or connection errors here are typically not recoverable. In haskell
non-recoverable failures are better modelled using @Exceptions@.

Although it throws 'NoMoreInput' or 'BrokenFrame' when appropriate, it provides
hooks to override these when constructing a 'Framer'.

By use of 'setOnClosed' and 'setOnBadParse', the caller of @runFramer@ can
completely override the exception type that is raised when @runFramer@ encounters
any failure.

-}


-- | Thrown by 'runFramer' or 'runOneFrame' if parsing fails and there is no
-- handler installed using 'setOnBadParse', or it does not throw an exception.
newtype BrokenFrame = BrokenFrame String
  deriving (BrokenFrame -> BrokenFrame -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: BrokenFrame -> BrokenFrame -> Bool
$c/= :: BrokenFrame -> BrokenFrame -> Bool
== :: BrokenFrame -> BrokenFrame -> Bool
$c== :: BrokenFrame -> BrokenFrame -> Bool
Eq, Int -> BrokenFrame -> ShowS
[BrokenFrame] -> ShowS
BrokenFrame -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [BrokenFrame] -> ShowS
$cshowList :: [BrokenFrame] -> ShowS
show :: BrokenFrame -> String
$cshow :: BrokenFrame -> String
showsPrec :: Int -> BrokenFrame -> ShowS
$cshowsPrec :: Int -> BrokenFrame -> ShowS
Show)



instance Exception BrokenFrame


{- | Thrown by 'runFramer' or 'runOneFrame' when no further input is available and
 no end of input handler is set using 'setOnClosed'.
-}
data NoMoreInput = NoMoreInput
  deriving (NoMoreInput -> NoMoreInput -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NoMoreInput -> NoMoreInput -> Bool
$c/= :: NoMoreInput -> NoMoreInput -> Bool
== :: NoMoreInput -> NoMoreInput -> Bool
$c== :: NoMoreInput -> NoMoreInput -> Bool
Eq, Int -> NoMoreInput -> ShowS
[NoMoreInput] -> ShowS
NoMoreInput -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [NoMoreInput] -> ShowS
$cshowList :: [NoMoreInput] -> ShowS
show :: NoMoreInput -> String
$cshow :: NoMoreInput -> String
showsPrec :: Int -> NoMoreInput -> ShowS
$cshowsPrec :: Int -> NoMoreInput -> ShowS
Show)


instance Exception NoMoreInput


closedReason :: String
closedReason :: String
closedReason = String
"not enough input"


defaultChunkSize :: Word32
defaultChunkSize :: Word32
defaultChunkSize = Word32
2048