{-# LANGUAGE OverloadedStrings #-}

{-|
Module      : AWSLambda.Handler
Stability   : experimental
Portability : POSIX

Entry point for AWS Lambda handlers deployed with @serverless-haskell@ plugin.
-}
module AWSLambda.Handler
  ( lambdaMain
  , lambdaMainRaw
  ) where

import           Control.Exception.Safe (MonadCatch, SomeException(..), displayException, tryAny)
import           Control.Monad (forever, void)
import           Control.Monad.IO.Class

import           Data.Aeson ((.=))
import qualified Data.Aeson as Aeson

import           Data.Typeable (typeOf)

import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Char8 as Char8
import qualified Data.ByteString.Lazy as LBS

import qualified Data.Text.Encoding as Text
import qualified Data.Text.IO as Text

import           GHC.IO.Handle (BufferMode(..), hSetBuffering)

import           Network.HTTP.Client
import           Network.HTTP.Types (HeaderName)

import           System.Environment (lookupEnv)
import           System.IO (stdout)

-- | Process incoming events from @serverless-haskell@ using a provided
-- function.
--
-- The handler receives the input event given to the AWS Lambda function, and
-- its return value is returned from the function.
--
-- This is intended to be used as @main@, for example:
--
-- > import qualified Data.Aeson as Aeson
-- >
-- > import AWSLambda
-- >
-- > main = lambdaMain handler
-- >
-- > handler :: Aeson.Value -> IO [Int]
-- > handler evt = do
-- >   putStrLn "This should go to logs"
-- >   print evt
-- >   pure [1, 2, 3]
--
-- The handler function can receive arbitrary JSON values from custom
-- invocations, or one of the events from the "AWSLambda.Events" module, such as
-- 'AWSLambda.Events.S3Event':
--
-- > import AWSLambda.Events.S3Event
-- >
-- > handler :: S3Event -> IO ()
-- > handler evt = do
-- >   print $ records evt
--
-- If the Lambda function needs to process several types of events, use
-- 'Data.Aeson.Alternative' to combine several handlers:
--
-- > import AWSLambda
-- > import AWSLambda.Events.S3Event
-- > import Data.Aeson
-- > import Data.Aeson.Alternative
-- >
-- > main = lambdaMain $ handlerS3 `alternative` handlerCustom
-- >
-- > handlerS3 :: S3Event -> IO ()
-- > handlerS3 = _
-- >
-- > handlerCustom :: Value -> IO ()
-- > handlerCustom = _
--
-- When run outside the AWS Lambda environment, the input is read as JSON from
-- the command line, and the result of the execution is printed, also as JSON,
-- to the standard output.
lambdaMain ::
     (Aeson.FromJSON event, Aeson.ToJSON res, MonadCatch m, MonadIO m)
  => (event -> m res) -- ^ Function to process the event
  -> m ()
lambdaMain :: (event -> m res) -> m ()
lambdaMain event -> m res
act =
  (ByteString -> m ByteString) -> m ()
forall (m :: * -> *).
(MonadCatch m, MonadIO m) =>
(ByteString -> m ByteString) -> m ()
lambdaMainRaw ((ByteString -> m ByteString) -> m ())
-> (ByteString -> m ByteString) -> m ()
forall a b. (a -> b) -> a -> b
$ \ByteString
input -> do
    case ByteString -> Either String event
forall a. FromJSON a => ByteString -> Either String a
Aeson.eitherDecode ByteString
input of
      Left String
err -> String -> m ByteString
forall a. HasCallStack => String -> a
error String
err
      Right event
event -> do
        res
result <- event -> m res
act event
event
        ByteString -> m ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> m ByteString) -> ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ res -> ByteString
forall a. ToJSON a => a -> ByteString
Aeson.encode res
result

-- | Process the incoming requests (using the AWS Lambda runtime interface or from the standard input).
-- Also set line buffering on standard output for AWS Lambda so the logs are output in a timely manner.
-- This function provides a lower level interface than 'lambdaMain' for users who don't want to use
-- Aeson for encoding and decoding JSON.
lambdaMainRaw :: (MonadCatch m, MonadIO m) => (LBS.ByteString -> m LBS.ByteString) -> m ()
lambdaMainRaw :: (ByteString -> m ByteString) -> m ()
lambdaMainRaw ByteString -> m ByteString
act = do
  Maybe String
lambdaApiAddress <- IO (Maybe String) -> m (Maybe String)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe String) -> m (Maybe String))
-> IO (Maybe String) -> m (Maybe String)
forall a b. (a -> b) -> a -> b
$ String -> IO (Maybe String)
lookupEnv String
lambdaApiAddressEnv
  case Maybe String
lambdaApiAddress of
    Just String
address -> do
      IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Handle -> BufferMode -> IO ()
hSetBuffering Handle
stdout BufferMode
LineBuffering
      Manager
manager <- IO Manager -> m Manager
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Manager -> m Manager) -> IO Manager -> m Manager
forall a b. (a -> b) -> a -> b
$ ManagerSettings -> IO Manager
newManager ManagerSettings
defaultManagerSettings
      m () -> m ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        Response ByteString
invocation <- IO (Response ByteString) -> m (Response ByteString)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Response ByteString) -> m (Response ByteString))
-> IO (Response ByteString) -> m (Response ByteString)
forall a b. (a -> b) -> a -> b
$ Request -> Manager -> IO (Response ByteString)
httpLbs (String -> Request
invocationRequest String
address) Manager
manager
        let input :: ByteString
input = Response ByteString -> ByteString
forall body. Response body -> body
responseBody Response ByteString
invocation
        let requestId :: String
requestId = Response ByteString -> String
forall a. Response a -> String
responseRequestId Response ByteString
invocation
        Either SomeException ByteString
resultOrError <- m ByteString -> m (Either SomeException ByteString)
forall (m :: * -> *) a.
MonadCatch m =>
m a -> m (Either SomeException a)
tryAny (m ByteString -> m (Either SomeException ByteString))
-> m ByteString -> m (Either SomeException ByteString)
forall a b. (a -> b) -> a -> b
$ ByteString -> m ByteString
act ByteString
input
        case Either SomeException ByteString
resultOrError of
          Right ByteString
result   -> IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IO (Response ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Response ()) -> IO ()) -> IO (Response ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ Request -> Manager -> IO (Response ())
httpNoBody (String -> String -> ByteString -> Request
resultRequest String
address String
requestId ByteString
result) Manager
manager
          Left SomeException
exception -> do
            ByteString -> m ()
forall (m :: * -> *). MonadIO m => ByteString -> m ()
putStrLnLBS (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$ Value -> ByteString
forall a. ToJSON a => a -> ByteString
Aeson.encode (Value -> ByteString) -> Value -> ByteString
forall a b. (a -> b) -> a -> b
$ SomeException -> Value
exceptionJSON SomeException
exception
            IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IO (Response ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Response ()) -> IO ()) -> IO (Response ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ Request -> Manager -> IO (Response ())
httpNoBody (String -> String -> SomeException -> Request
errorRequest String
address String
requestId SomeException
exception) Manager
manager
    Maybe String
Nothing -> do
      ByteString
input <- IO ByteString -> m ByteString
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> m ByteString) -> IO ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
LBS.fromStrict (ByteString -> ByteString) -> IO ByteString -> IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ByteString
ByteString.getLine
      ByteString
result <- ByteString -> m ByteString
act ByteString
input
      ByteString -> m ()
forall (m :: * -> *). MonadIO m => ByteString -> m ()
putStrLnLBS ByteString
result

putStrLnLBS :: MonadIO m => LBS.ByteString -> m ()
putStrLnLBS :: ByteString -> m ()
putStrLnLBS = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (ByteString -> IO ()) -> ByteString -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> IO ()
Text.putStrLn (Text -> IO ()) -> (ByteString -> Text) -> ByteString -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
Text.decodeUtf8 (ByteString -> Text)
-> (ByteString -> ByteString) -> ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
LBS.toStrict

lambdaApiAddressEnv :: String
lambdaApiAddressEnv :: String
lambdaApiAddressEnv = String
"AWS_LAMBDA_RUNTIME_API"

lambdaRequest :: String -> String -> Request
lambdaRequest :: String -> String -> Request
lambdaRequest String
apiAddress String
rqPath = String -> Request
parseRequest_ (String -> Request) -> String -> Request
forall a b. (a -> b) -> a -> b
$ String
"http://" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
apiAddress String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"/2018-06-01" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
rqPath

invocationRequest :: String -> Request
invocationRequest :: String -> Request
invocationRequest String
apiAddress = (String -> String -> Request
lambdaRequest String
apiAddress String
"/runtime/invocation/next") { responseTimeout :: ResponseTimeout
responseTimeout = ResponseTimeout
responseTimeoutNone }

resultRequest :: String -> String -> LBS.ByteString -> Request
resultRequest :: String -> String -> ByteString -> Request
resultRequest String
apiAddress String
requestId ByteString
result = (String -> String -> Request
lambdaRequest String
apiAddress (String -> Request) -> String -> Request
forall a b. (a -> b) -> a -> b
$ String
"/runtime/invocation/" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
requestId String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"/response") { method :: ByteString
method = ByteString
"POST", requestBody :: RequestBody
requestBody = ByteString -> RequestBody
RequestBodyLBS ByteString
result }

errorRequest :: String -> String -> SomeException -> Request
errorRequest :: String -> String -> SomeException -> Request
errorRequest String
apiAddress String
requestId SomeException
exception = (String -> String -> Request
lambdaRequest String
apiAddress (String -> Request) -> String -> Request
forall a b. (a -> b) -> a -> b
$ String
"/runtime/invocation/" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
requestId String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"/error") { method :: ByteString
method = ByteString
"POST", requestBody :: RequestBody
requestBody = ByteString -> RequestBody
RequestBodyLBS ByteString
body }
  where
    body :: ByteString
body = Value -> ByteString
forall a. ToJSON a => a -> ByteString
Aeson.encode (Value -> ByteString) -> Value -> ByteString
forall a b. (a -> b) -> a -> b
$ SomeException -> Value
exceptionJSON SomeException
exception

exceptionJSON :: SomeException -> Aeson.Value
exceptionJSON :: SomeException -> Value
exceptionJSON SomeException
exception = [Pair] -> Value
Aeson.object [ Text
"errorMessage" Text -> String -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= SomeException -> String
forall e. Exception e => e -> String
displayException SomeException
exception, Text
"errorType" Text -> String -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= SomeException -> String
exceptionType SomeException
exception]

exceptionType :: SomeException -> String
exceptionType :: SomeException -> String
exceptionType (SomeException e
e) = TypeRep -> String
forall a. Show a => a -> String
show (e -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf e
e)

requestIdHeader :: HeaderName
requestIdHeader :: HeaderName
requestIdHeader = HeaderName
"Lambda-Runtime-Aws-Request-Id"

responseRequestId :: Response a -> String
responseRequestId :: Response a -> String
responseRequestId = ByteString -> String
Char8.unpack (ByteString -> String)
-> (Response a -> ByteString) -> Response a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (HeaderName, ByteString) -> ByteString
forall a b. (a, b) -> b
snd ((HeaderName, ByteString) -> ByteString)
-> (Response a -> (HeaderName, ByteString))
-> Response a
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(HeaderName, ByteString)] -> (HeaderName, ByteString)
forall a. [a] -> a
head ([(HeaderName, ByteString)] -> (HeaderName, ByteString))
-> (Response a -> [(HeaderName, ByteString)])
-> Response a
-> (HeaderName, ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((HeaderName, ByteString) -> Bool)
-> [(HeaderName, ByteString)] -> [(HeaderName, ByteString)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((HeaderName -> ByteString -> Bool)
-> (HeaderName, ByteString) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((HeaderName -> ByteString -> Bool)
 -> (HeaderName, ByteString) -> Bool)
-> (HeaderName -> ByteString -> Bool)
-> (HeaderName, ByteString)
-> Bool
forall a b. (a -> b) -> a -> b
$ \HeaderName
h ByteString
_ -> HeaderName
h HeaderName -> HeaderName -> Bool
forall a. Eq a => a -> a -> Bool
== HeaderName
requestIdHeader) ([(HeaderName, ByteString)] -> [(HeaderName, ByteString)])
-> (Response a -> [(HeaderName, ByteString)])
-> Response a
-> [(HeaderName, ByteString)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Response a -> [(HeaderName, ByteString)]
forall body. Response body -> [(HeaderName, ByteString)]
responseHeaders