Server.hs

{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{- | Contains the servant implentation of the 'DiscoveryApi'. -}
module Network.Legion.Discovery.Server (
  discoveryServer,
) where

import Control.Monad.Except (throwError)
import Control.Monad.IO.Class (liftIO)
import Data.ByteString (hGetContents)
import Data.GraphViz (graphvizWithHandle, GraphvizCommand(Dot),
  GraphvizOutput(Svg))
import Data.Monoid ((<>))
import Data.Text (Text)
import Data.Text.Encoding (encodeUtf8)
import Data.Version (Version)
import Distribution.Version (anyVersion)
import Network.Legion (Runtime)
import Network.Legion.Discovery.Api(DiscoveryApi, V1Api, ClientList,
  DiscoveryApi, PingRequest(PingRequest), Range(Range), serviceAddress,
  metadata, PingApi, InstanceList(InstanceList), GraphApi, Graph(Graph),
  SvgGraph(SvgGraph), ServicesApi, unClientList)
import Network.Legion.Discovery.Graphviz (toDotGraph)
import Network.Legion.Discovery.LegionApp (Input, Output, unEntityName,
  Client(Client), cName, cVersion, EntityName, Client, Output, getRange,
  getAllServices)
import Servant ((:<|>)((:<|>)), Server, ServantErr, errBody,
  errHeaders, err400, NoContent(NoContent), Handler)
import qualified Data.ByteString.Lazy as BSL
import qualified Network.Legion.Discovery.LegionApp as App


{- | The implementation of the 'DiscoveryApi' api.  -}
discoveryServer :: Runtime Input Output -> Server DiscoveryApi
discoveryServer = v1Server


{- | version 1 of the api. -}
v1Server :: Runtime Input Output -> Server V1Api
v1Server runtime = pingServer :<|> servicesServer :<|> graphServer
  where
    pingServer :: Server PingApi
    pingServer Nothing _ = throwError missingUserAgent
    pingServer (Just clients) PingRequest {serviceAddress, metadata} = do
        strictClients <- mapM toStrict (unClientList clients)
        liftIO $ do
          sequence_ [
              App.ping runtime name version serviceAddress metadata
              | (name, version) <- strictClients
            ]
          return NoContent
      where
        {- |
          Return the client name and version, or else throw an error if
          the client doesn't provide a version.
        -}
        toStrict :: Client -> Handler (EntityName, Version)
        toStrict Client {cName, cVersion = Just version} =
          return (cName, version)
        toStrict Client {cName, cVersion = Nothing} =
          throwError $ badRequest (
              "Missing version for client product: "
              <> unEntityName cName
              <> ". Please fix your User-Agent header."
            )

    servicesServer :: Server ServicesApi
    servicesServer =
          allInstances
        :<|>
          rangedInstances
      where
        allInstances
          :: Maybe ClientList
          -> EntityName
          -> Handler InstanceList
        allInstances clients name = 
          rangedInstances clients name (Range anyVersion)

        rangedInstances
          :: Maybe ClientList
          -> EntityName
          -> Range
          -> Handler InstanceList
        rangedInstances Nothing _ _ = throwError missingUserAgent
        rangedInstances (Just clients) name (Range range) = liftIO $
          InstanceList <$> getRange runtime (unClientList clients) name range


    graphServer :: Server GraphApi
    graphServer = svgGraph :<|> dotGraph
      where
        dotGraph = Graph . toDotGraph <$> getAllServices runtime
        svgGraph = do
          g <- toDotGraph <$> getAllServices runtime
          bytes <- liftIO $ graphvizWithHandle Dot g Svg hGetContents
          return (SvgGraph (BSL.fromStrict bytes))


{- | 400 BadRequest servant error helper. -}
badRequest :: Text -> ServantErr
badRequest msg = err400 {
    errBody = BSL.fromStrict (encodeUtf8 msg),
    errHeaders = [("content-type", "text/plain")]
  }


{- | Errror for when the user-agent header is missing. -}
missingUserAgent :: ServantErr
missingUserAgent = badRequest "Missing User-Agent request header."