Discovery.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
module Network.Legion.Discovery (
  main
) where

import Canteven.HTTP (requestLogging, logExceptionsAndContinue,
  DecodeResult(Unsupported, BadEntity, Ok), FromEntity, decodeEntity)
import Canteven.Log.MonadLog (getCantevenOutput)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Aeson (encode, object, (.=), FromJSON, parseJSON,
  Value(Object), (.:), eitherDecode)
import Data.Conduit ((=$=), ($$))
import Data.Map (Map)
import Data.String (IsString)
import Data.Text (pack)
import Data.Time (getCurrentTime)
import Data.Version (showVersion)
import Distribution.Text (simpleParse)
import Network.HTTP.Types (badRequest400, unsupportedMediaType415,
  noContent204)
import Network.Legion (forkLegionary, Runtime, makeRequest, search,
  SearchTag(SearchTag), IndexRecord, irKey, PartitionKey)
import Network.Legion.Config (parseArgs)
import Network.Legion.Discovery.App (newApp, Input(GetRange, GetAll, Ping,
  GetService), Output(InstanceList, PingResponse, ServiceResponse),
  ServiceId(ServiceId), toKey, Time(Time), unServiceAddr, version,
  ServiceAddr(ServiceAddr), InstanceInfo, Service, instances, name,
  unServiceId)
import Network.Legion.Discovery.Config (servicePort)
import Network.Legion.Discovery.HttpError (HttpError)
import Network.Legion.Discovery.LIO (runLIO, LIO)
import Network.Wai (Middleware, modifyResponse)
import Network.Wai.Middleware.AddHeaders (addHeaders)
import Network.Wai.Middleware.StripHeaders (stripHeader)
import Web.Scotty.Resource.Trans (resource, get, post)
import Web.Scotty.Trans (scottyT, middleware, ScottyT, param, setHeader,
  raw, status, text, ActionT, ScottyError, body, header)
import qualified Data.ByteString.Lazy as LBS
import qualified Data.Conduit.List as CL
import qualified Data.Map as Map
import qualified Data.Text.Encoding as TE
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Encoding as TLE
import qualified Network.Legion.Discovery.Config as C
import qualified Paths_legion_discovery as P

main :: IO ()
main = do
  (settings, startupMode, config) <- parseArgs
  logging <- getCantevenOutput (C.logging config)
  app <- newApp

  runLIO logging $ do
    legion <- forkLegionary app settings startupMode
    scottyT (servicePort config) (runLIO logging) $ do
      middleware
        $ requestLogging logging
        . setServer "legion-discovery"
        . logExceptionsAndContinue logging
      webService legion


{- | The web service endpoint definitions.  -}
webService :: Runtime Input Output -> ScottyT HttpError LIO ()
webService runtime = do
  resource "/v1/ping/:serviceId/:version" $
    post $
      simpleParse <$> param "version" >>= \case
        Nothing -> do
          status badRequest400
          text "Invalid version."
        Just ver -> do
          serviceId <- ServiceId <$> param "serviceId"
          now <- Time <$> liftIO getCurrentTime
          withEntity (\ PingRequest {serviceAddress} -> do
              let req = Ping now serviceId ver serviceAddress
              makeRequest runtime (toKey serviceId) req >>= \case
                PingResponse -> status noContent204
                InstanceList _ -> fail "Invalid runtime response."
                ServiceResponse _ -> fail "Invalid runtime response."
            )
  resource "/v1/services" $
    get $ do
      now <- Time <$> liftIO getCurrentTime
      list <-
        search runtime (SearchTag "" Nothing)
        =$= CL.mapMaybeM (fillServiceInfo now)
        $$ CL.consume
      setHeader "content-type" serviceListCT
      raw . encode $ object [
          unServiceId (name service) .= encodeInstances (instances service)
          | service <- list
        ]
  resource "/v1/services/:serviceId" $
    get $ do
      serviceId <- ServiceId <$> param "serviceId"
      now <- Time <$> liftIO getCurrentTime
      respondInstances =<< getInstances runtime (toKey serviceId) (GetAll now)
  resource "/v1/services/:serviceId/:versionRange" $
    get $ do
      serviceId <- ServiceId <$> param "serviceId"
      simpleParse <$> param "versionRange" >>= \case
        Nothing -> do
          status badRequest400
          text "Invalid version range."
        Just range -> do
          now <- Time <$> liftIO getCurrentTime
          respondInstances =<<
            getInstances runtime (toKey serviceId) (GetRange now range)
  where
    {- | Get the service info for the serivce listing.  -}
    fillServiceInfo
      :: Time
      -> IndexRecord
      -> ActionT HttpError LIO (Maybe Service)
    fillServiceInfo now ir =
      makeRequest runtime (irKey ir) (GetService now) >>= \case
        ServiceResponse service -> return service
        PingResponse -> fail "Invalid runtime response."
        InstanceList _ -> fail "Invalid runtime response."

    {- | Send a response containing the service instance list. -}
    respondInstances
      :: Map ServiceAddr InstanceInfo
      -> ActionT HttpError LIO ()
    respondInstances is = do
      setHeader "content-type" instanceListCT
      (raw . encode . encodeInstances) is

    {- | Encode instances into a JSON object -}
    encodeInstances :: Map ServiceAddr InstanceInfo -> Value
    encodeInstances instances = object [
        unServiceAddr addr .= object [
            "version" .= showVersion (version info)
          ]
        | (addr, info) <- Map.toList instances
      ]

{- | Send a legion request that returns an InstanceList response.  -}
getInstances :: (MonadIO io)
  => Runtime Input Output
  -> PartitionKey
  -> Input
  -> io (Map ServiceAddr InstanceInfo)
getInstances runtime key input =
  makeRequest runtime key input >>= \case
    InstanceList instances -> return instances
    PingResponse -> fail "Invalid runtime response."
    ServiceResponse _ -> fail "Invalid runtime response."


{- | Set the server header.  -}
setServer :: String -> Middleware
setServer serviceName = addServerHeader . stripServerHeader
  where
    {- |
      Strip the server header
    -}
    stripServerHeader :: Middleware
    stripServerHeader = modifyResponse (stripHeader "Server")

    {- |
      Add our own server header.
    -}
    addServerHeader :: Middleware
    addServerHeader = addHeaders [("Server", serverValue)]

    {- |
      The value of the @Server:@ header.
    -}
    serverValue =
      TE.encodeUtf8 (pack (serviceName ++ "/" ++ showVersion P.version))


{- | The service instance list content type. -}
instanceListCT :: (IsString a) => a
instanceListCT = "application/vnd.legion-discovery.instance-list+json"


{- | The known service list content type. -}
serviceListCT :: (IsString a) => a
serviceListCT = "application/vnd.legion-discovery.service-list+json"


{- | Scotty shorthand for getting and decoding an entity. -}
withEntity :: (FromEntity a, MonadIO m, ScottyError e)
  => (a -> ActionT e m ())
  -> ActionT e m ()
withEntity f =
  decodeEntity <$> headerLBS "content-type" <*> body >>= \case
    Unsupported -> status unsupportedMediaType415
    BadEntity why -> do
      status badRequest400
      text (TL.pack why)
    Ok b -> f b


{- | Get a header as a 'Date.ByteString.Lazy.ByteString'.  -}
headerLBS :: (ScottyError e, Monad m)
  => TL.Text
  -> ActionT e m (Maybe LBS.ByteString)
headerLBS headerName = fmap TLE.encodeUtf8 <$> header headerName


{- | Decode a ping request entity. -}
newtype PingRequest = PingRequest {
    serviceAddress :: ServiceAddr
  }
instance FromJSON PingRequest where
  parseJSON (Object o) =
    PingRequest . ServiceAddr <$> o .: "serviceAddress"
  parseJSON v = fail
    $ "Can't parse PingRequest from: " ++ show v
instance FromEntity PingRequest where
  decodeEntity (Just PingRequestCT) bytes =
    case eitherDecode bytes of
      Left err -> BadEntity err
      Right req -> Ok req
  decodeEntity _ _ = Unsupported


{- | The content type for ping requests. -}
pattern PingRequestCT = "application/vnd.legion-discovery.ping-request+json"