{-# 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 :: (Message -> Message) -> Request o -> Request o
mapRequestMessage Message -> Message
f (RCall Message
m)   = Message -> Request Message
RCall (Message -> Request Message) -> Message -> Request Message
forall a b. (a -> b) -> a -> b
$ Message -> Message
f Message
m
mapRequestMessage Message -> Message
f (ROneway Message
m) = Message -> Request ()
ROneway (Message -> Request ()) -> Message -> Request ()
forall a b. (a -> b) -> a -> b
$ Message -> Message
f Message
m

-- | Extract the message contained in the request.
getRequestMessage :: Request o -> Message
getRequestMessage :: 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 (HashMap TypeRep Dynamic -> Context)
-> HashMap TypeRep Dynamic -> Context
forall a b. (a -> b) -> a -> b
$ HashMap TypeRep Dynamic
a HashMap TypeRep Dynamic
-> HashMap TypeRep Dynamic -> HashMap TypeRep Dynamic
forall a. Semigroup a => a -> a -> a
<> HashMap TypeRep Dynamic
b

instance Monoid Context where
  mempty :: Context
mempty = HashMap TypeRep Dynamic -> Context
Context HashMap TypeRep Dynamic
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 :: i -> Context -> Context
addToContext i
i (Context HashMap TypeRep Dynamic
m) =
  HashMap TypeRep Dynamic -> Context
Context (HashMap TypeRep Dynamic -> Context)
-> HashMap TypeRep Dynamic -> Context
forall a b. (a -> b) -> a -> b
$ TypeRep
-> Dynamic -> HashMap TypeRep Dynamic -> HashMap TypeRep Dynamic
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert (i -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf i
i) (i -> Dynamic
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 :: Context -> Maybe i
lookupInContext (Context HashMap TypeRep Dynamic
m) = do
  Dynamic
x <- TypeRep -> HashMap TypeRep Dynamic -> Maybe Dynamic
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup (Proxy i -> TypeRep
forall k (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy i
forall k (t :: k). Proxy t
Proxy :: Proxy i)) HashMap TypeRep Dynamic
m
  case Dynamic -> Maybe i
forall a. Typeable a => Dynamic -> Maybe a
fromDynamic @i Dynamic
x of
    Maybe i
Nothing -> String -> Maybe i
forall a. HasCallStack => String -> a
error String
"Impossible!"
    Just i
y  -> i -> Maybe i
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. Context -> Request a -> IO a) -> ThriftServer)
-> (forall a. Context -> Request a -> IO a) -> 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 (Text -> Maybe Handler) -> Text -> Maybe Handler
forall a b. (a -> b) -> a -> b
$ Message -> Text
messageName Message
msg of
        Just (CallHandler Context -> c -> IO r
f') ->
          case Parser c -> Either String c
forall a. Parser a -> Either String a
runParser (Parser c -> Either String c) -> Parser c -> Either String c
forall a b. (a -> b) -> a -> b
$ Value (Tag c) -> Parser c
forall a. Pinchable a => Value (Tag a) -> Parser a
unpinch (Value (Tag c) -> Parser c) -> Value (Tag c) -> Parser c
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
              Message -> IO Message
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Message -> IO Message) -> Message -> IO Message
forall a b. (a -> b) -> a -> b
$ Message :: Text -> MessageType -> Int32 -> Value TStruct -> Message
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 = r -> Value (Tag r)
forall a. Pinchable a => a -> Value (Tag a)
pinch r
ret
                  }
            Left String
err ->
              Message -> IO Message
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Message -> IO Message) -> Message -> IO Message
forall a b. (a -> b) -> a -> b
$ Message -> ApplicationException -> Message
mkApplicationExceptionReply Message
msg (ApplicationException -> Message)
-> ApplicationException -> Message
forall a b. (a -> b) -> a -> b
$
                Text -> ExceptionType -> ApplicationException
ApplicationException (Text
"Unable to parse service arguments: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack String
err) ExceptionType
InternalError

        Just (OnewayHandler Context -> c -> IO ()
_) ->
          Message -> IO Message
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Message -> IO Message) -> Message -> IO Message
forall a b. (a -> b) -> a -> b
$ Message -> ApplicationException -> Message
mkApplicationExceptionReply Message
msg (ApplicationException -> Message)
-> ApplicationException -> Message
forall a b. (a -> b) -> a -> b
$
            Text -> ExceptionType -> ApplicationException
ApplicationException Text
"Expected message type Oneway, got Call." ExceptionType
InvalidMessageType
        Maybe Handler
Nothing ->
          Message -> IO Message
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Message -> IO Message) -> Message -> IO Message
forall a b. (a -> b) -> a -> b
$ Message -> ApplicationException -> Message
mkApplicationExceptionReply Message
msg (ApplicationException -> Message)
-> ApplicationException -> Message
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 (Text -> Maybe Handler) -> Text -> Maybe Handler
forall a b. (a -> b) -> a -> b
$ Message -> Text
messageName Message
msg of
        Just (OnewayHandler Context -> c -> IO ()
f') -> do
          case Parser c -> Either String c
forall a. Parser a -> Either String a
runParser (Parser c -> Either String c) -> Parser c -> Either String c
forall a b. (a -> b) -> a -> b
$ Value (Tag c) -> Parser c
forall a. Pinchable a => Value (Tag a) -> Parser a
unpinch (Value (Tag c) -> Parser c) -> Value (Tag c) -> Parser c
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   ->
              ApplicationException -> IO a
forall e a. Exception e => e -> IO a
throwIO (ApplicationException -> IO a) -> ApplicationException -> IO a
forall a b. (a -> b) -> a -> b
$ Text -> ExceptionType -> ApplicationException
ApplicationException (Text
"Unable to parse service arguments: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack String
err) ExceptionType
InternalError
        Just (CallHandler Context -> c -> IO r
_) ->
          ApplicationException -> IO a
forall e a. Exception e => e -> IO a
throwIO (ApplicationException -> IO a) -> ApplicationException -> IO a
forall a b. (a -> b) -> a -> b
$ Text -> ExceptionType -> ApplicationException
ApplicationException Text
"Expected message type Call, got Oneway." ExceptionType
InvalidMessageType
        Maybe Handler
Nothing ->
          ApplicationException -> IO a
forall e a. Exception e => e -> IO a
throwIO (ApplicationException -> IO a) -> ApplicationException -> IO a
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. Context -> Request a -> IO a) -> ThriftServer)
-> (forall a. Context -> Request a -> IO a) -> ThriftServer
forall a b. (a -> b) -> a -> b
$ \Context
ctx Request a
req -> do
  case Request a
req of
    RCall Message
msg -> Context -> Request a -> (ApplicationException -> IO a) -> IO a
forall a.
Context -> Request a -> (ApplicationException -> IO a) -> IO a
go Context
ctx Request a
req (Message -> IO Message
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Message -> IO Message)
-> (ApplicationException -> Message)
-> ApplicationException
-> IO Message
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
_ -> Context -> Request a -> (ApplicationException -> IO a) -> IO a
forall a.
Context -> Request a -> (ApplicationException -> IO a) -> IO a
go Context
ctx Request a
req ApplicationException -> IO a
forall e a. Exception e => e -> IO a
throwIO
  where
    srvMap :: HashMap ServiceName ThriftServer
srvMap = [(ServiceName, ThriftServer)] -> HashMap ServiceName ThriftServer
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 :: 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 (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
':') (Message -> Text
messageName (Message -> Text) -> Message -> Text
forall a b. (a -> b) -> a -> b
$ Request a -> Message
forall o. Request o -> Message
getRequestMessage Request a
req)
      let prefix' :: ServiceName
prefix' = Text -> ServiceName
ServiceName Text
prefix
      let ctx' :: Context
ctx' = ServiceName -> Context -> Context
forall i. ContextItem i => i -> Context -> Context
addToContext ServiceName
prefix' Context
ctx
      case ServiceName
prefix' ServiceName
-> HashMap ServiceName ThriftServer -> Maybe ThriftServer
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 (ApplicationException -> IO a) -> ApplicationException -> IO a
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 -> Context -> Request a -> IO a
ThriftServer -> forall a. Context -> Request a -> IO a
unThriftServer ThriftServer
srv Context
ctx' (Request a -> IO a) -> Request a -> IO a
forall a b. (a -> b) -> a -> b
$ (Message -> Message) -> Request a -> Request a
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
_ -> () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            RCall Message
_   -> a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
reply
        Maybe ThriftServer
Nothing -> ApplicationException -> IO a
onErr (ApplicationException -> IO a) -> ApplicationException -> IO a
forall a b. (a -> b) -> a -> b
$ Text -> ExceptionType -> ApplicationException
ApplicationException (Text
"No service with name " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
prefix Text -> Text -> Text
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 :: (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. Context -> Request a -> IO a) -> ThriftServer)
-> (forall a. Context -> Request a -> IO a) -> ThriftServer
forall a b. (a -> b) -> a -> b
$
  \Context
ctx Request a
req ->
    (e -> Maybe a) -> IO a -> (a -> IO a) -> IO a
forall e b a.
Exception e =>
(e -> Maybe b) -> IO a -> (b -> IO a) -> IO a
catchJust e -> Maybe a
sel
      (ThriftServer -> Context -> Request a -> IO a
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 -> () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    T.RRFailure String
err -> do
      ThriftError -> IO ()
forall e a. Exception e => e -> IO a
throwIO (ThriftError -> IO ()) -> ThriftError -> IO ()
forall a b. (a -> b) -> a -> b
$ Text -> ThriftError
ThriftError (Text -> ThriftError) -> Text -> 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 <- IO Message -> IO (Either SomeException Message)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO Message -> IO (Either SomeException Message))
-> IO Message -> IO (Either SomeException Message)
forall a b. (a -> b) -> a -> b
$ ThriftServer -> Context -> Request Message -> IO Message
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 <- SomeException -> Maybe ApplicationException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e -> Channel -> Message -> IO ()
writeMessage Channel
chan (Message -> IO ()) -> Message -> IO ()
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 (Message -> IO ()) -> Message -> IO ()
forall a b. (a -> b) -> a -> b
$ Message -> ApplicationException -> Message
mkApplicationExceptionReply Message
call (ApplicationException -> Message)
-> ApplicationException -> Message
forall a b. (a -> b) -> a -> b
$
              Text -> ExceptionType -> ApplicationException
ApplicationException (Text
"Could not process request: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> (String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ SomeException -> String
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 -> Context -> Request () -> IO ()
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 -> ApplicationException -> IO ()
forall e a. Exception e => e -> IO a
throwIO (ApplicationException -> IO ()) -> ApplicationException -> IO ()
forall a b. (a -> b) -> a -> b
$
          Text -> ExceptionType -> ApplicationException
ApplicationException (Text
"Expected call, got " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> (String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ MessageType -> String
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 :: Text -> MessageType -> Int32 -> Value TStruct -> Message
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 = ApplicationException -> Value (Tag ApplicationException)
forall a. Pinchable a => a -> Value (Tag a)
pinch ApplicationException
ex
  }