Admin.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}

{- | This module contains the admin interface code. -}
module Network.Legion.Admin (
  runAdmin,
  forkAdmin,
) where


import Canteven.HTTP (requestLogging, logExceptionsAndContinue)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Logger (askLoggerIO, runLoggingT, logDebug,
  MonadLoggerIO)
import Control.Monad.Trans.Class (lift)
import Data.Default.Class (def)
import Data.Text.Encoding (encodeUtf8)
import Data.Text.Lazy (Text)
import Data.Version (showVersion)
import Network.HTTP.Types (notFound404)
import Network.Legion.Application (LegionConstraints)
import Network.Legion.Fork (forkC)
import Network.Legion.LIO (LIO)
import Network.Legion.PartitionKey (PartitionKey(K), unKey)
import Network.Legion.Runtime (Runtime, debugRuntimeState, getDivergent,
  debugLocalPartitions, debugPartition, eject, debugIndex)
import Network.Legion.Settings (RuntimeSettings, adminPort, adminHost)
import Network.Wai (Middleware, modifyResponse)
import Network.Wai.Handler.Warp (HostPreference, defaultSettings, Port,
  setHost, setPort)
import Network.Wai.Middleware.AddHeaders (addHeaders)
import Network.Wai.Middleware.StripHeaders (stripHeader)
import Paths_legion (version)
import Text.Read (readMaybe)
import Web.Scotty.Resource.Trans (resource, get, delete)
import Web.Scotty.Trans (Options, scottyOptsT, settings, ScottyT, param,
  middleware, status, json, text)
import qualified Data.Map as Map
import qualified Data.Text as T
import qualified Data.Text.Lazy as TL

{- | Fork the admin website in a background thread. -}
forkAdmin :: (LegionConstraints e o s, MonadLoggerIO m)
  => RuntimeSettings
    {- ^ Settings and configuration of the legion framework. -}
  -> Runtime e o s
  -> m ()
forkAdmin legionSettings runtime = do
  logging <- askLoggerIO
  liftIO . (`runLoggingT` logging) . forkC "admin thread" $
    runAdmin (adminPort legionSettings) (adminHost legionSettings) runtime


{- | Run the admin service. Does not return. -}
runAdmin :: (LegionConstraints e o s, MonadLoggerIO io)
  => Port
  -> HostPreference
  -> Runtime e o s
  -> io ()
runAdmin addr host runtime = do
  logging <- askLoggerIO
  let
    website :: ScottyT Text LIO ()
    website = do
      middleware
        $ requestLogging logging
        . setServer
        . logExceptionsAndContinue logging

      resource "/rts" $
        get $ text . TL.pack . show =<< debugRuntimeState runtime
      resource "/index" $
        get $ json =<< debugIndex runtime
      resource "/divergent" $
        get $
          json . Map.mapKeys show =<< getDivergent runtime
      resource "/partitions" $
        get $
          json . Map.mapKeys (show . toInteger . unKey)
            =<< debugLocalPartitions runtime
      resource "/partitions/:key" $
        get $ do
          key <- K . read <$> param "key"
          json =<< debugPartition runtime key
      resource "/peers/:peer" $
        delete $
          readMaybe <$> param "peer" >>= \case
            Nothing -> status notFound404
            Just peer -> do
              lift . $(logDebug) . T.pack $ "Ejecting peer: " ++ show peer
              eject runtime peer

  scottyOptsT (options addr host) (`runLoggingT` logging) website


{- | Build some warp settings based on the configured socket address. -}
options :: Port -> HostPreference -> Options
options port host = def {
    settings =
      setPort port
      . setHost host
      $ defaultSettings
  }


setServer :: Middleware
setServer = 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 =
      encodeUtf8 (T.pack ("legion-admin/" ++ showVersion version))