Main.hs

{-# LANGUAGE OverloadedStrings #-}

import Control.Applicative ((<|>), liftA2)
import Data.Monoid ((<>))
import Data.Maybe (fromMaybe)
import Data.Text.Lazy (Text)
import qualified Data.Text.Lazy as L
import Text.Read (readMaybe)
import System.Environment (lookupEnv)
import Network.HTTP.Types (StdMethod(HEAD))
import Web.Scotty

getClientIP1 :: ActionM (Maybe Text)
getClientIP1 = header "CLIENT-IP"

getClientIP2 :: ActionM (Maybe Text)
getClientIP2 =
  fmap (fmap firstElement) (header "X-FORWARDED-FOR") where
      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"]

getClientIP :: ActionM ()
getClientIP = do
  ip <- foldl1 (liftA2 (<|>)) [getClientIP1, getClientIP2, getClientIP3]
  case ip of
    Nothing -> text "unknown\n"
    Just rawIp -> text $ rawIp <> "\n"

-- | 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
      Just xs -> fromMaybe defaultValue (stringReader xs)

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

main :: IO ()
main = do
  port <- getListenPort
  scotty port $ do
    get "/" getClientIP
    addroute HEAD "/" $ return ()