{-# 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 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
        (AddrInfo
ai:[AddrInfo]
_) -> 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 HostName
numericHost, Maybe HostName
_) -> 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
                (Maybe HostName, Maybe HostName)
_ -> 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
$ HostName
"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 ByteString
host 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 = 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
</> 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 ByteString
"application/json"
    wrapHandler :: ByteString -> m a -> m a
wrapHandler ByteString
fmt 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 Request
req = CI ByteString -> Request -> Maybe ByteString
forall a. HasHeaders a => CI ByteString -> a -> Maybe ByteString
getHeader CI ByteString
"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 ByteString
fmt 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 ByteString
hdr | ByteString
hdr ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
fmt -> m a
action
        Maybe ByteString
_ -> 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 = 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 ByteString
"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 ByteString
pathInfo = do
        let segments :: [ByteString]
segments  = Char -> ByteString -> [ByteString]
S8.split Char
'/' ByteString
pathInfo
            nameBytes :: ByteString
nameBytes = ByteString -> [ByteString] -> ByteString
S8.intercalate ByteString
"." [ByteString]
segments
        case ByteString -> Either UnicodeException Text
T.decodeUtf8' ByteString
nameBytes of
            Left UnicodeException
_ -> 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 Int
400 ByteString
"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 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
                    Maybe Value
Nothing -> m ()
forall (m :: * -> *) a. MonadSnap m => m a
pass
                    Just 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 Word8
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 ByteString
s =
        let (ByteString
s', ByteString
q) = Word8 -> ByteString -> (ByteString, ByteString)
breakDiscard Word8
59 ByteString
s -- semicolon
            (ByteString
_, ByteString
q') = Word8 -> ByteString -> (ByteString, ByteString)
breakDiscard Word8
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 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
                (p
x, HostName
_):[(p, HostName)]
_ -> p
x
                [(p, HostName)]
_ -> p
1.0
    trimWhite :: ByteString -> ByteString
trimWhite = (Word8 -> Bool) -> ByteString -> ByteString
S.dropWhile (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
32) -- space

breakDiscard :: Word8 -> S.ByteString -> (S.ByteString, S.ByteString)
breakDiscard :: Word8 -> ByteString -> (ByteString, ByteString)
breakDiscard Word8
w ByteString
s =
    let (ByteString
x, 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 Int
1 ByteString
y)