{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE StandaloneDeriving        #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeFamilies              #-}

module Pinch.Server
  (
    -- * Thrift Server creation

    ThriftServer (..)
  , createServer
  , Handler(..)
  , Request (..)

    -- * Running a Thrift Server
  , runConnection
  , ThriftError (..)
  , Channel (..)
  , createChannel
  , createChannel1

    -- * Thrift Server context

    -- | The context can be used to pass data from the environment to the thrift server functions.
    -- For example, you could pass the remote host name to the server to use it for logging purposes.
  , Context
  , ContextItem
  , addToContext
  , lookupInContext

    -- * Middlewares
  , multiplex
  , ServiceName (..)
  , onError

    -- * Helper functions

    -- | Functions mostly useful for defining custom `ThriftServer`s.
  , mapRequestMessage
  , getRequestMessage
  , mkApplicationExceptionReply
  ) where

import           Control.Exception        (Exception, SomeException, catchJust,
                                           fromException, throwIO, try)
import           Data.Dynamic             (Dynamic (..), fromDynamic, toDyn)
import           Data.Proxy               (Proxy (..))
import           Data.Typeable            (TypeRep, Typeable, typeOf, typeRep)

import qualified Data.HashMap.Strict      as HM
import qualified Data.Text                as T

import           Pinch.Internal.Exception
import           Pinch.Internal.Message
import           Pinch.Internal.Pinchable
import           Pinch.Internal.RPC
import           Pinch.Internal.TType

import qualified Pinch.Transport          as T

-- | A single request to a thrift server.
data Request out where
  RCall :: !Message -> Request Message
  ROneway :: !Message -> Request ()

deriving instance Show (Request out)

-- | Map the message contained in the request.
mapRequestMessage :: (Message -> Message) -> Request o -> Request o
mapRequestMessage :: forall o. (Message -> Message) -> Request o -> Request o
mapRequestMessage Message -> Message
f (RCall Message
m)   = Message -> Request Message
RCall forall a b. (a -> b) -> a -> b
$ Message -> Message
f Message
m
mapRequestMessage Message -> Message
f (ROneway Message
m) = Message -> Request ()
ROneway forall a b. (a -> b) -> a -> b
$ Message -> Message
f Message
m

-- | Extract the message contained in the request.
getRequestMessage :: Request o -> Message
getRequestMessage :: forall o. Request o -> Message
getRequestMessage (RCall Message
m)   = Message
m
getRequestMessage (ROneway Message
m) = Message
m

-- | A `Thrift` server. Takes the context and the request as input and may produces a reply message.
newtype ThriftServer = ThriftServer { ThriftServer -> forall a. Context -> Request a -> IO a
unThriftServer :: forall a . Context -> Request a -> IO a }

-- | Allows passing context information to a `ThriftServer`.
-- The context is indexed by type.
newtype Context = Context (HM.HashMap TypeRep Dynamic)

instance Semigroup Context where
  (Context HashMap TypeRep Dynamic
a) <> :: Context -> Context -> Context
<> (Context HashMap TypeRep Dynamic
b) = HashMap TypeRep Dynamic -> Context
Context forall a b. (a -> b) -> a -> b
$ HashMap TypeRep Dynamic
a forall a. Semigroup a => a -> a -> a
<> HashMap TypeRep Dynamic
b

instance Monoid Context where
  mempty :: Context
mempty = HashMap TypeRep Dynamic -> Context
Context forall a. Monoid a => a
mempty

class Typeable a => ContextItem a where

instance ContextItem ServiceName


-- | Adds a new item to the context. If an item with the same
-- type is already part of the context, it will be overwritten.
addToContext :: forall i . ContextItem i => i -> Context -> Context
addToContext :: forall i. ContextItem i => i -> Context -> Context
addToContext i
i (Context HashMap TypeRep Dynamic
m) =
  HashMap TypeRep Dynamic -> Context
Context forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert (forall a. Typeable a => a -> TypeRep
typeOf i
i) (forall a. Typeable a => a -> Dynamic
toDyn i
i) HashMap TypeRep Dynamic
m

-- | Lookup a value in the context.
lookupInContext :: forall i . ContextItem i => Context -> Maybe i
lookupInContext :: forall i. ContextItem i => Context -> Maybe i
lookupInContext (Context HashMap TypeRep Dynamic
m) = do
  Dynamic
x <- forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup (forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (forall {k} (t :: k). Proxy t
Proxy :: Proxy i)) HashMap TypeRep Dynamic
m
  case forall a. Typeable a => Dynamic -> Maybe a
fromDynamic @i Dynamic
x of
    Maybe i
Nothing -> forall a. HasCallStack => String -> a
error String
"Impossible!"
    Just i
y  -> forall (f :: * -> *) a. Applicative f => a -> f a
pure i
y

-- | Create a handler for a request type.
data Handler where
  -- | Handle normal call requests. Must return a result.
  CallHandler :: (Pinchable c, Tag c ~ TStruct, Pinchable r, Tag r ~ TStruct) => (Context -> c -> IO r) -> Handler
  -- | Handle oneway requests. Cannot return any result.
  OnewayHandler :: (Pinchable c, Tag c ~ TStruct) => (Context -> c -> IO ()) -> Handler

-- | Creates a new thrift server processing requests with the function `f`.
--
-- By default, if processing a oneway call fails a haskell exception is thrown which will likely
-- terminate the guilty connection. You may use the `onError` combinator to handle this case
-- more gracefully.
createServer :: (T.Text -> Maybe Handler) -> ThriftServer
createServer :: (Text -> Maybe Handler) -> ThriftServer
createServer Text -> Maybe Handler
f = (forall a. Context -> Request a -> IO a) -> ThriftServer
ThriftServer forall a b. (a -> b) -> a -> b
$ \Context
ctx Request a
req ->
  case Request a
req of
    RCall Message
msg ->
      case Text -> Maybe Handler
f forall a b. (a -> b) -> a -> b
$ Message -> Text
messageName Message
msg of
        Just (CallHandler Context -> c -> IO r
f') ->
          case forall a. Parser a -> Either String a
runParser forall a b. (a -> b) -> a -> b
$ forall a. Pinchable a => Value (Tag a) -> Parser a
unpinch forall a b. (a -> b) -> a -> b
$ Message -> Value TStruct
messagePayload Message
msg of
            Right c
args -> do
              r
ret <- Context -> c -> IO r
f' Context
ctx c
args
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Message
                  { messageName :: Text
messageName = Message -> Text
messageName Message
msg
                  , messageType :: MessageType
messageType = MessageType
Reply
                  , messageId :: Int32
messageId   = Message -> Int32
messageId Message
msg
                  , messagePayload :: Value TStruct
messagePayload = forall a. Pinchable a => a -> Value (Tag a)
pinch r
ret
                  }
            Left String
err ->
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Message -> ApplicationException -> Message
mkApplicationExceptionReply Message
msg forall a b. (a -> b) -> a -> b
$
                Text -> ExceptionType -> ApplicationException
ApplicationException (Text
"Unable to parse service arguments: " forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack String
err) ExceptionType
InternalError

        Just (OnewayHandler Context -> c -> IO ()
_) ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Message -> ApplicationException -> Message
mkApplicationExceptionReply Message
msg forall a b. (a -> b) -> a -> b
$
            Text -> ExceptionType -> ApplicationException
ApplicationException Text
"Expected message type Oneway, got Call." ExceptionType
InvalidMessageType
        Maybe Handler
Nothing ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Message -> ApplicationException -> Message
mkApplicationExceptionReply Message
msg forall a b. (a -> b) -> a -> b
$ Text -> ExceptionType -> ApplicationException
ApplicationException Text
"Unknown method name." ExceptionType
WrongMethodName

    ROneway Message
msg ->
      -- we cannot return errors to the client as it is a oneway call.
      -- Instead we just throw an exception, possible terminating
      -- the guilty connection.
      case Text -> Maybe Handler
f forall a b. (a -> b) -> a -> b
$ Message -> Text
messageName Message
msg of
        Just (OnewayHandler Context -> c -> IO ()
f') -> do
          case forall a. Parser a -> Either String a
runParser forall a b. (a -> b) -> a -> b
$ forall a. Pinchable a => Value (Tag a) -> Parser a
unpinch forall a b. (a -> b) -> a -> b
$ Message -> Value TStruct
messagePayload Message
msg of
            Right c
args -> Context -> c -> IO ()
f' Context
ctx c
args
            Left String
err   ->
              forall e a. Exception e => e -> IO a
throwIO forall a b. (a -> b) -> a -> b
$ Text -> ExceptionType -> ApplicationException
ApplicationException (Text
"Unable to parse service arguments: " forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack String
err) ExceptionType
InternalError
        Just (CallHandler Context -> c -> IO r
_) ->
          forall e a. Exception e => e -> IO a
throwIO forall a b. (a -> b) -> a -> b
$ Text -> ExceptionType -> ApplicationException
ApplicationException Text
"Expected message type Call, got Oneway." ExceptionType
InvalidMessageType
        Maybe Handler
Nothing ->
          forall e a. Exception e => e -> IO a
throwIO forall a b. (a -> b) -> a -> b
$ Text -> ExceptionType -> ApplicationException
ApplicationException Text
"Unknown method name." ExceptionType
WrongMethodName

-- | Multiplex multiple services into a single `ThriftServer`.
--
-- The service name is added to the `Context` and may be retrieved using `lookupInContext @ServiceName ctx`.
multiplex :: [(ServiceName, ThriftServer)] -> ThriftServer
multiplex :: [(ServiceName, ThriftServer)] -> ThriftServer
multiplex [(ServiceName, ThriftServer)]
services = (forall a. Context -> Request a -> IO a) -> ThriftServer
ThriftServer forall a b. (a -> b) -> a -> b
$ \Context
ctx Request a
req -> do
  case Request a
req of
    RCall Message
msg -> forall a.
Context -> Request a -> (ApplicationException -> IO a) -> IO a
go Context
ctx Request a
req (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Message -> ApplicationException -> Message
mkApplicationExceptionReply Message
msg)
    -- we cannot send the exception back, because it is a oneway call
    -- instead let's just throw it and crash the server
    ROneway Message
_ -> forall a.
Context -> Request a -> (ApplicationException -> IO a) -> IO a
go Context
ctx Request a
req forall e a. Exception e => e -> IO a
throwIO
  where
    srvMap :: HashMap ServiceName ThriftServer
srvMap = forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HM.fromList [(ServiceName, ThriftServer)]
services

    go :: Context -> Request a -> (ApplicationException -> IO a) -> IO a
    go :: forall a.
Context -> Request a -> (ApplicationException -> IO a) -> IO a
go Context
ctx Request a
req ApplicationException -> IO a
onErr = do
      let (Text
prefix, Text
method) = (Char -> Bool) -> Text -> (Text, Text)
T.span (forall a. Eq a => a -> a -> Bool
/= Char
':') (Message -> Text
messageName forall a b. (a -> b) -> a -> b
$ forall o. Request o -> Message
getRequestMessage Request a
req)
      let prefix' :: ServiceName
prefix' = Text -> ServiceName
ServiceName Text
prefix
      let ctx' :: Context
ctx' = forall i. ContextItem i => i -> Context -> Context
addToContext ServiceName
prefix' Context
ctx
      case ServiceName
prefix' forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
`HM.lookup` HashMap ServiceName ThriftServer
srvMap of
        Maybe ThriftServer
_ | Text -> Bool
T.null Text
method -> ApplicationException -> IO a
onErr forall a b. (a -> b) -> a -> b
$ Text -> ExceptionType -> ApplicationException
ApplicationException Text
"Invalid method name, expecting a colon." ExceptionType
WrongMethodName
        Just ThriftServer
srv -> do

          a
reply <- ThriftServer -> forall a. Context -> Request a -> IO a
unThriftServer ThriftServer
srv Context
ctx' forall a b. (a -> b) -> a -> b
$ forall o. (Message -> Message) -> Request o -> Request o
mapRequestMessage (\Message
msg -> Message
msg { messageName :: Text
messageName = Text -> Text
T.tail Text
method }) Request a
req

          case Request a
req of
            ROneway Message
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            RCall Message
_   -> forall (f :: * -> *) a. Applicative f => a -> f a
pure a
reply
        Maybe ThriftServer
Nothing -> ApplicationException -> IO a
onErr forall a b. (a -> b) -> a -> b
$ Text -> ExceptionType -> ApplicationException
ApplicationException (Text
"No service with name " forall a. Semigroup a => a -> a -> a
<> Text
prefix forall a. Semigroup a => a -> a -> a
<> Text
" available.") ExceptionType
UnknownMethod

-- | Add error handlers to a `ThriftServer`. Exceptions are caught and not re-thrown, but you may do
-- so by calling `ioThrow` yourself.
onError
  :: Exception e
  => (e -> Maybe a) -- ^ Select exceptions to handle.
  -> (a -> IO Message) -- ^ Error handler for normal method calls.
  -> (a -> IO ()) -- ^ Error handler for oneway calls.
  -> ThriftServer -> ThriftServer
onError :: forall e a.
Exception e =>
(e -> Maybe a)
-> (a -> IO Message)
-> (a -> IO ())
-> ThriftServer
-> ThriftServer
onError e -> Maybe a
sel a -> IO Message
callError a -> IO ()
onewayError ThriftServer
srv = (forall a. Context -> Request a -> IO a) -> ThriftServer
ThriftServer forall a b. (a -> b) -> a -> b
$
  \Context
ctx Request a
req ->
    forall e b a.
Exception e =>
(e -> Maybe b) -> IO a -> (b -> IO a) -> IO a
catchJust e -> Maybe a
sel
      (ThriftServer -> forall a. Context -> Request a -> IO a
unThriftServer ThriftServer
srv Context
ctx Request a
req)
      (\a
e -> do
        case Request a
req of
          RCall Message
_   -> a -> IO Message
callError a
e
          ROneway Message
_ -> a -> IO ()
onewayError a
e
      )

-- | Run a Thrift server for a single connection.
runConnection :: Context -> ThriftServer -> Channel -> IO ()
runConnection :: Context -> ThriftServer -> Channel -> IO ()
runConnection Context
ctx ThriftServer
srv Channel
chan = do
  ReadResult Message
msg <- Channel -> IO (ReadResult Message)
readMessage Channel
chan
  case ReadResult Message
msg of
    ReadResult Message
T.RREOF -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    T.RRFailure String
err -> do
      forall e a. Exception e => e -> IO a
throwIO forall a b. (a -> b) -> a -> b
$ Text -> ThriftError
ThriftError forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack String
err
    T.RRSuccess Message
call -> do
      case Message -> MessageType
messageType Message
call of
        MessageType
Call -> do
          Either SomeException Message
r <- forall e a. Exception e => IO a -> IO (Either e a)
try forall a b. (a -> b) -> a -> b
$ ThriftServer -> forall a. Context -> Request a -> IO a
unThriftServer ThriftServer
srv Context
ctx (Message -> Request Message
RCall Message
call)
          case Either SomeException Message
r of
            -- if it is already an ApplicationException, we just send it back
            Left (SomeException
e :: SomeException)
              | Just ApplicationException
appEx <- forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e -> Channel -> Message -> IO ()
writeMessage Channel
chan forall a b. (a -> b) -> a -> b
$ Message -> ApplicationException -> Message
mkApplicationExceptionReply Message
call ApplicationException
appEx
            Left (SomeException
e :: SomeException) -> Channel -> Message -> IO ()
writeMessage Channel
chan forall a b. (a -> b) -> a -> b
$ Message -> ApplicationException -> Message
mkApplicationExceptionReply Message
call forall a b. (a -> b) -> a -> b
$
              Text -> ExceptionType -> ApplicationException
ApplicationException (Text
"Could not process request: " forall a. Semigroup a => a -> a -> a
<> (String -> Text
T.pack forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show SomeException
e)) ExceptionType
InternalError
            Right Message
x -> Channel -> Message -> IO ()
writeMessage Channel
chan Message
x
        MessageType
Oneway -> do
          -- no matter what happens, we can never send back an error because the client is not listening for replies
          -- when doing a oneway calls...
          -- Let's just crash the connection in this case, to avoid silently swallowing errors.
          -- `onError` can be used to handle this more gracefully.
          ThriftServer -> forall a. Context -> Request a -> IO a
unThriftServer ThriftServer
srv Context
ctx (Message -> Request ()
ROneway Message
call)
        -- the client must never send Reply/Exception messages.
        MessageType
t -> forall e a. Exception e => e -> IO a
throwIO forall a b. (a -> b) -> a -> b
$
          Text -> ExceptionType -> ApplicationException
ApplicationException (Text
"Expected call, got " forall a. Semigroup a => a -> a -> a
<> (String -> Text
T.pack forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show MessageType
t)) ExceptionType
InvalidMessageType
      Context -> ThriftServer -> Channel -> IO ()
runConnection Context
ctx ThriftServer
srv Channel
chan

-- | Builds an exception reply given the corresponding request message.
mkApplicationExceptionReply :: Message -> ApplicationException -> Message
mkApplicationExceptionReply :: Message -> ApplicationException -> Message
mkApplicationExceptionReply Message
req ApplicationException
ex = Message
  { messageName :: Text
messageName = Message -> Text
messageName Message
req
  , messageType :: MessageType
messageType = MessageType
Exception
  , messageId :: Int32
messageId = Message -> Int32
messageId Message
req
  , messagePayload :: Value TStruct
messagePayload = forall a. Pinchable a => a -> Value (Tag a)
pinch ApplicationException
ex
  }