{-# LANGUAGE OverloadedStrings #-}

module System.Remote.Snap
    ( startServer
    ) where

import Control.Applicative ((<$>), (<|>))
import Control.Exception (throwIO)
import Control.Monad.IO.Class (liftIO)
import qualified Data.ByteString as S
import qualified Data.ByteString.Char8 as S8
import Data.Function (on)
import qualified Data.HashMap.Strict as M
import qualified Data.List as List
import qualified Data.Text.Encoding as T
import Data.Word (Word8)
import Network.Socket (NameInfoFlag(NI_NUMERICHOST), addrAddress, getAddrInfo,
                       getNameInfo)
import Paths_ekg (getDataDir)
import Prelude hiding (read)
import Snap.Core (MonadSnap, Request, Snap, finishWith, getHeader, getRequest,
                  getResponse, method, Method(GET), modifyResponse, pass,
                  rqPathInfo, setContentType, setResponseStatus,
                  writeLBS)
import Snap.Http.Server (httpServe)
import qualified Snap.Http.Server.Config as Config
import Snap.Util.FileServe (serveDirectory)
import System.FilePath ((</>))

import System.Metrics
import System.Remote.Json

------------------------------------------------------------------------

-- | Convert a host name (e.g. \"localhost\" or \"127.0.0.1\") to a
-- numeric host address (e.g. \"127.0.0.1\").
getNumericHostAddress :: S.ByteString -> IO S.ByteString
getNumericHostAddress :: ByteString -> IO ByteString
getNumericHostAddress host :: ByteString
host = do
    [AddrInfo]
ais <- Maybe AddrInfo -> Maybe HostName -> Maybe HostName -> IO [AddrInfo]
getAddrInfo Maybe AddrInfo
forall a. Maybe a
Nothing (HostName -> Maybe HostName
forall a. a -> Maybe a
Just (ByteString -> HostName
S8.unpack ByteString
host)) Maybe HostName
forall a. Maybe a
Nothing
    case [AddrInfo]
ais of
        [] -> IO ByteString
forall a. IO a
unsupportedAddressError
        (ai :: AddrInfo
ai:_) -> do
            (Maybe HostName, Maybe HostName)
ni <- [NameInfoFlag]
-> Bool -> Bool -> SockAddr -> IO (Maybe HostName, Maybe HostName)
getNameInfo [NameInfoFlag
NI_NUMERICHOST] Bool
True Bool
False (AddrInfo -> SockAddr
addrAddress AddrInfo
ai)
            case (Maybe HostName, Maybe HostName)
ni of
                (Just numericHost :: HostName
numericHost, _) -> ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$! HostName -> ByteString
S8.pack HostName
numericHost
                _ -> IO ByteString
forall a. IO a
unsupportedAddressError
  where
    unsupportedAddressError :: IO a
unsupportedAddressError = IOError -> IO a
forall e a. Exception e => e -> IO a
throwIO (IOError -> IO a) -> IOError -> IO a
forall a b. (a -> b) -> a -> b
$
        HostName -> IOError
userError (HostName -> IOError) -> HostName -> IOError
forall a b. (a -> b) -> a -> b
$ "unsupported address: " HostName -> HostName -> HostName
forall a. [a] -> [a] -> [a]
++ ByteString -> HostName
S8.unpack ByteString
host

startServer :: Store
            -> S.ByteString  -- ^ Host to listen on (e.g. \"localhost\")
            -> Int           -- ^ Port to listen on (e.g. 8000)
            -> IO ()
startServer :: Store -> ByteString -> Int -> IO ()
startServer store :: Store
store host :: ByteString
host port :: Int
port = do
    -- Snap doesn't allow for non-numeric host names in
    -- 'Snap.setBind'. We work around that limitation by converting a
    -- possible non-numeric host name to a numeric address.
    ByteString
numericHost <- ByteString -> IO ByteString
getNumericHostAddress ByteString
host
    let conf :: Config Snap a
conf = Bool -> Config Snap a -> Config Snap a
forall (m :: * -> *) a. Bool -> Config m a -> Config m a
Config.setVerbose Bool
False (Config Snap a -> Config Snap a) -> Config Snap a -> Config Snap a
forall a b. (a -> b) -> a -> b
$
               ConfigLog -> Config Snap a -> Config Snap a
forall (m :: * -> *) a. ConfigLog -> Config m a -> Config m a
Config.setErrorLog ConfigLog
Config.ConfigNoLog (Config Snap a -> Config Snap a) -> Config Snap a -> Config Snap a
forall a b. (a -> b) -> a -> b
$
               ConfigLog -> Config Snap a -> Config Snap a
forall (m :: * -> *) a. ConfigLog -> Config m a -> Config m a
Config.setAccessLog ConfigLog
Config.ConfigNoLog (Config Snap a -> Config Snap a) -> Config Snap a -> Config Snap a
forall a b. (a -> b) -> a -> b
$
               Int -> Config Snap a -> Config Snap a
forall (m :: * -> *) a. Int -> Config m a -> Config m a
Config.setPort Int
port (Config Snap a -> Config Snap a) -> Config Snap a -> Config Snap a
forall a b. (a -> b) -> a -> b
$
               ByteString -> Config Snap a -> Config Snap a
forall (m :: * -> *) a. ByteString -> Config m a -> Config m a
Config.setHostname ByteString
host (Config Snap a -> Config Snap a) -> Config Snap a -> Config Snap a
forall a b. (a -> b) -> a -> b
$
               ByteString -> Config Snap a -> Config Snap a
forall (m :: * -> *) a. ByteString -> Config m a -> Config m a
Config.setBind ByteString
numericHost (Config Snap a -> Config Snap a) -> Config Snap a -> Config Snap a
forall a b. (a -> b) -> a -> b
$
               Config Snap a
forall (m :: * -> *) a. MonadSnap m => Config m a
Config.defaultConfig
    Config Snap Any -> Snap () -> IO ()
forall a. Config Snap a -> Snap () -> IO ()
httpServe Config Snap Any
forall a. Config Snap a
conf (Store -> Snap ()
monitor Store
store)

-- | A handler that can be installed into an existing Snap application.
monitor :: Store -> Snap ()
monitor :: Store -> Snap ()
monitor store :: Store
store = do
    HostName
dataDir <- IO HostName -> Snap HostName
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO HostName
getDataDir
    (Snap () -> Snap ()
forall a. Snap a -> Snap a
jsonHandler (Snap () -> Snap ()) -> Snap () -> Snap ()
forall a b. (a -> b) -> a -> b
$ Store -> Snap ()
forall (m :: * -> *). MonadSnap m => Store -> m ()
serve Store
store)
        Snap () -> Snap () -> Snap ()
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> HostName -> Snap ()
forall (m :: * -> *). MonadSnap m => HostName -> m ()
serveDirectory (HostName
dataDir HostName -> HostName -> HostName
</> "assets")
  where
    jsonHandler :: Snap a -> Snap a
jsonHandler = ByteString -> Snap a -> Snap a
forall (m :: * -> *) a. MonadSnap m => ByteString -> m a -> m a
wrapHandler "application/json"
    wrapHandler :: ByteString -> m a -> m a
wrapHandler fmt :: ByteString
fmt handler :: m a
handler = Method -> m a -> m a
forall (m :: * -> *) a. MonadSnap m => Method -> m a -> m a
method Method
GET (m a -> m a) -> m a -> m a
forall a b. (a -> b) -> a -> b
$ ByteString -> m a -> m a
forall (m :: * -> *) a. MonadSnap m => ByteString -> m a -> m a
format ByteString
fmt (m a -> m a) -> m a -> m a
forall a b. (a -> b) -> a -> b
$ m a
handler

-- | The Accept header of the request.
acceptHeader :: Request -> Maybe S.ByteString
acceptHeader :: Request -> Maybe ByteString
acceptHeader req :: Request
req = CI ByteString -> Request -> Maybe ByteString
forall a. HasHeaders a => CI ByteString -> a -> Maybe ByteString
getHeader "Accept" Request
req

-- | Runs a Snap monad action only if the request's Accept header
-- matches the given MIME type.
format :: MonadSnap m => S.ByteString -> m a -> m a
format :: ByteString -> m a -> m a
format fmt :: ByteString
fmt action :: m a
action = do
    Request
req <- m Request
forall (m :: * -> *). MonadSnap m => m Request
getRequest
    let acceptHdr :: Maybe ByteString
acceptHdr = ([ByteString] -> ByteString
forall a. [a] -> a
List.head ([ByteString] -> ByteString)
-> (ByteString -> [ByteString]) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> [ByteString]
parseHttpAccept) (ByteString -> ByteString) -> Maybe ByteString -> Maybe ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Request -> Maybe ByteString
acceptHeader Request
req
    case Maybe ByteString
acceptHdr of
        Just hdr :: ByteString
hdr | ByteString
hdr ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
fmt -> m a
action
        _ -> m a
forall (m :: * -> *) a. MonadSnap m => m a
pass

-- | Serve all counter, gauges and labels, built-in or not, as a
-- nested JSON object.
serve :: MonadSnap m => Store -> m ()
serve :: Store -> m ()
serve store :: Store
store = do
    Request
req <- m Request
forall (m :: * -> *). MonadSnap m => m Request
getRequest
    (Response -> Response) -> m ()
forall (m :: * -> *). MonadSnap m => (Response -> Response) -> m ()
modifyResponse ((Response -> Response) -> m ()) -> (Response -> Response) -> m ()
forall a b. (a -> b) -> a -> b
$ ByteString -> Response -> Response
setContentType "application/json"
    if ByteString -> Bool
S.null (Request -> ByteString
rqPathInfo Request
req)
        then m ()
serveAll
        else ByteString -> m ()
forall (m :: * -> *). MonadSnap m => ByteString -> m ()
serveOne (Request -> ByteString
rqPathInfo Request
req)
  where
    serveAll :: m ()
serveAll = do
        Sample
metrics <- IO Sample -> m Sample
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Sample -> m Sample) -> IO Sample -> m Sample
forall a b. (a -> b) -> a -> b
$ Store -> IO Sample
sampleAll Store
store
        ByteString -> m ()
forall (m :: * -> *). MonadSnap m => ByteString -> m ()
writeLBS (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$ Sample -> ByteString
encodeAll Sample
metrics
    serveOne :: ByteString -> m ()
serveOne pathInfo :: ByteString
pathInfo = do
        let segments :: [ByteString]
segments  = Char -> ByteString -> [ByteString]
S8.split '/' ByteString
pathInfo
            nameBytes :: ByteString
nameBytes = ByteString -> [ByteString] -> ByteString
S8.intercalate "." [ByteString]
segments
        case ByteString -> Either UnicodeException Text
T.decodeUtf8' ByteString
nameBytes of
            Left _ -> do
                (Response -> Response) -> m ()
forall (m :: * -> *). MonadSnap m => (Response -> Response) -> m ()
modifyResponse ((Response -> Response) -> m ()) -> (Response -> Response) -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> Response -> Response
setResponseStatus 400 "Bad Request"
                Response
r <- m Response
forall (m :: * -> *). MonadSnap m => m Response
getResponse
                Response -> m ()
forall (m :: * -> *) a. MonadSnap m => Response -> m a
finishWith Response
r
            Right name :: Text
name -> do
                Sample
metrics <- IO Sample -> m Sample
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Sample -> m Sample) -> IO Sample -> m Sample
forall a b. (a -> b) -> a -> b
$ Store -> IO Sample
sampleAll Store
store
                case Text -> Sample -> Maybe Value
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup Text
name Sample
metrics of
                    Nothing -> m ()
forall (m :: * -> *) a. MonadSnap m => m a
pass
                    Just metric :: Value
metric -> ByteString -> m ()
forall (m :: * -> *). MonadSnap m => ByteString -> m ()
writeLBS (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$ Value -> ByteString
encodeOne Value
metric

------------------------------------------------------------------------
-- Utilities for working with accept headers

-- | Parse the HTTP accept string to determine supported content types.
parseHttpAccept :: S.ByteString -> [S.ByteString]
parseHttpAccept :: ByteString -> [ByteString]
parseHttpAccept = ((ByteString, Double) -> ByteString)
-> [(ByteString, Double)] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
List.map (ByteString, Double) -> ByteString
forall a b. (a, b) -> a
fst
                ([(ByteString, Double)] -> [ByteString])
-> (ByteString -> [(ByteString, Double)])
-> ByteString
-> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((ByteString, Double) -> (ByteString, Double) -> Ordering)
-> [(ByteString, Double)] -> [(ByteString, Double)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
List.sortBy (Double -> Double -> Ordering
rcompare (Double -> Double -> Ordering)
-> ((ByteString, Double) -> Double)
-> (ByteString, Double)
-> (ByteString, Double)
-> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (ByteString, Double) -> Double
forall a b. (a, b) -> b
snd)
                ([(ByteString, Double)] -> [(ByteString, Double)])
-> (ByteString -> [(ByteString, Double)])
-> ByteString
-> [(ByteString, Double)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString -> (ByteString, Double))
-> [ByteString] -> [(ByteString, Double)]
forall a b. (a -> b) -> [a] -> [b]
List.map ByteString -> (ByteString, Double)
forall b. (Read b, Fractional b) => ByteString -> (ByteString, b)
grabQ
                ([ByteString] -> [(ByteString, Double)])
-> (ByteString -> [ByteString])
-> ByteString
-> [(ByteString, Double)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word8 -> ByteString -> [ByteString]
S.split 44 -- comma
  where
    rcompare :: Double -> Double -> Ordering
    rcompare :: Double -> Double -> Ordering
rcompare = (Double -> Double -> Ordering) -> Double -> Double -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip Double -> Double -> Ordering
forall a. Ord a => a -> a -> Ordering
compare
    grabQ :: ByteString -> (ByteString, b)
grabQ s :: ByteString
s =
        let (s' :: ByteString
s', q :: ByteString
q) = Word8 -> ByteString -> (ByteString, ByteString)
breakDiscard 59 ByteString
s -- semicolon
            (_, q' :: ByteString
q') = Word8 -> ByteString -> (ByteString, ByteString)
breakDiscard 61 ByteString
q -- equals sign
         in (ByteString -> ByteString
trimWhite ByteString
s', ByteString -> b
forall p. (Read p, Fractional p) => ByteString -> p
readQ (ByteString -> b) -> ByteString -> b
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
trimWhite ByteString
q')
    readQ :: ByteString -> p
readQ s :: ByteString
s = case ReadS p
forall a. Read a => ReadS a
reads ReadS p -> ReadS p
forall a b. (a -> b) -> a -> b
$ ByteString -> HostName
S8.unpack ByteString
s of
                (x :: p
x, _):_ -> p
x
                _ -> 1.0
    trimWhite :: ByteString -> ByteString
trimWhite = (Word8 -> Bool) -> ByteString -> ByteString
S.dropWhile (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== 32) -- space

breakDiscard :: Word8 -> S.ByteString -> (S.ByteString, S.ByteString)
breakDiscard :: Word8 -> ByteString -> (ByteString, ByteString)
breakDiscard w :: Word8
w s :: ByteString
s =
    let (x :: ByteString
x, y :: ByteString
y) = (Word8 -> Bool) -> ByteString -> (ByteString, ByteString)
S.break (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
w) ByteString
s
    in (ByteString
x, Int -> ByteString -> ByteString
S.drop 1 ByteString
y)