Skip to content
Main.hs 5.29 KiB
Newer Older
Yuanle Song's avatar
Yuanle Song committed
{-# LANGUAGE OverloadedStrings #-}

import Control.Applicative ((<|>), liftA2)
import Data.Monoid ((<>))
Yuanle Song's avatar
Yuanle Song committed
import Data.Maybe (fromMaybe)
Yuanle Song's avatar
Yuanle Song committed
import Data.Text.Lazy (Text)
import qualified Data.Text.Lazy as L
Yuanle Song's avatar
Yuanle Song committed
import Text.Read (readMaybe)
import Data.List.Extra (lower)
import qualified Data.Sequence as S
import Control.Concurrent.MVar
import Control.Monad (when)
import Control.Monad.IO.Class (liftIO)
import qualified Data.Foldable as F
import Data.Time ( getCurrentTime
                 , formatTime
                 , defaultTimeLocale
                 , TimeZone(..)
                 , utcToZonedTime
                 )
import Lucid
Yuanle Song's avatar
Yuanle Song committed
import System.Environment (lookupEnv)
import Network.HTTP.Types (StdMethod(HEAD) , status404)
Yuanle Song's avatar
Yuanle Song committed
import Web.Scotty

-- in RAM call history for most recent 10 calls.
data CallHistory = CallHistory {
      chTime :: Text
    , chClientIP :: Text
    , chUserAgent :: Text
    } deriving (Show)

data RuntimeConfig = RuntimeConfig {
      rcServeHistoryPage :: Bool
    , rcCallHistory :: MVar (S.Seq CallHistory)
    }

Yuanle Song's avatar
Yuanle Song committed
getClientIP1 :: ActionM (Maybe Text)
getClientIP1 = header "CLIENT-IP"
Yuanle Song's avatar
Yuanle Song committed

getClientIP2 :: ActionM (Maybe Text)
getClientIP2 =
Yuanle Song's avatar
Yuanle Song committed
  fmap (fmap firstElement) (header "X-FORWARDED-FOR") where
Yuanle Song's avatar
Yuanle Song committed
      firstElement :: Text -> Text
      firstElement = head . L.splitOn ","

getClientIP3 :: ActionM (Maybe Text)
getClientIP3 = foldl1 (liftA2 (<|>)) $ map header ["X-FORWARDED"
                                                  ,"X-CLUSTER-CLIENT-IP"
                                                  ,"FORWARDED-FOR"
                                                  ,"FORWARDED"
                                                  ,"REMOTE-ADDR"]
-- | get current time in CST timezone.
getCurrentTimeCST :: IO Text
getCurrentTimeCST = do
  now <- getCurrentTime
  let tz = TimeZone { timeZoneMinutes=480
                    , timeZoneSummerOnly=False
                    , timeZoneName="CST"}
  return $ L.pack $ formatTime defaultTimeLocale "%c %z" (utcToZonedTime tz now)

-- | if rcServeHistoryPage is enabled, add call history to rcCallHistory Seq.
writeCallHistoryMaybe :: RuntimeConfig -> Text -> ActionM ()
writeCallHistoryMaybe rc rawIP =
    when (rcServeHistoryPage rc) $ do
      time <- liftIO getCurrentTimeCST
      userAgent <- header "User-Agent"
      callHistory <- liftIO $ takeMVar (rcCallHistory rc)
      let newEntry = CallHistory {
                       chTime = time
                     , chClientIP = rawIP
                     , chUserAgent = fromMaybe "unknown" userAgent}
      liftIO $ putMVar (rcCallHistory rc) ((S.|>) callHistory newEntry)

getClientIP :: RuntimeConfig -> ActionM ()
getClientIP rc = do
Yuanle Song's avatar
Yuanle Song committed
  ip <- foldl1 (liftA2 (<|>)) [getClientIP1, getClientIP2, getClientIP3]
  let rawIP = fromMaybe "unknown" ip
  text $ rawIP <> "\n"
  writeCallHistoryMaybe rc rawIP
Yuanle Song's avatar
Yuanle Song committed
-- | look up an environment variable, if it doesn't exist or is empty, use
-- default value instead. Otherwise, call a reader function on it and use that
-- value if parse is okay, use default value if parse fail.
getEnvDefault :: String -> a -> (String -> Maybe a) -> IO a
getEnvDefault variable defaultValue stringReader = do
    r <- lookupEnv variable
    return $ case r of
      Nothing -> defaultValue
      Just [] -> defaultValue
Yuanle Song's avatar
Yuanle Song committed
      Just xs -> fromMaybe defaultValue (stringReader xs)
Yuanle Song's avatar
Yuanle Song committed

getListenPort :: IO Int
getListenPort = getEnvDefault "PORT" 8081 readMaybe

readBoolMaybe :: String -> Maybe Bool
readBoolMaybe str = Just (lower str `elem` ["true", "yes", "1"])

buildHistoryText :: RuntimeConfig -> IO Text
buildHistoryText rc = do
    callHistory <- readMVar (rcCallHistory rc)
    return $ F.foldl' (\out ch -> out <> chClientIP ch <> ", "
                                      <> chTime ch <> ", "
                                      <> chUserAgent ch <> "\n")
                      "time, client ip, user-agent\n" callHistory

buildHistoryHtml :: RuntimeConfig -> IO Text
buildHistoryHtml rc = do
    callHistory <- readMVar (rcCallHistory rc)
    let thead = tr_ $ do
                  th_ "client ip"
                  th_ "time"
                  th_ "user-agent"
    let tbody = F.mapM_
                (\ch -> tr_ $ do
                           td_ $ toHtml $ chClientIP ch
                           td_ $ toHtml $ chTime ch
                           td_ $ toHtml $ chUserAgent ch) callHistory
    return $ renderText $ do
      doctype_
      html_ [lang_ "en"] $ do
        head_ $ do
          title_ "recent calls"
        body_ $ table_ $ do
                  thead
                  tbody

showHistory :: RuntimeConfig -> ActionM ()
showHistory rc =
  if rcServeHistoryPage rc then
      do
        accept <- header "accept"
        case accept of
          Just "text/plain" -> do
            result <- liftIO $ buildHistoryText rc
            text result
          _ -> do
            history <- liftIO $ buildHistoryHtml rc
            html history
  else
      status status404

Yuanle Song's avatar
Yuanle Song committed
main :: IO ()
Yuanle Song's avatar
Yuanle Song committed
main = do
  port <- getListenPort
  serveHistoryPage <- getEnvDefault "SERVE_HISTORY_PAGE" False readBoolMaybe
  when serveHistoryPage $
    putStrLn "Enabled GET /_calls api"
  callHistory <- newMVar S.empty
  let rc = RuntimeConfig { rcServeHistoryPage = serveHistoryPage
                         , rcCallHistory = callHistory}
Yuanle Song's avatar
Yuanle Song committed
  scotty port $ do
    get "/" (getClientIP rc)
    get "/_calls" (showHistory rc)
    addroute HEAD "/" $ return ()