/src/Main.hs

  1{-# LANGUAGE OverloadedStrings #-}
  2module Main where
  3
  4import           Control.Applicative ((<|>))
  5import           Control.Monad.Trans (liftIO)
  6import           Control.Concurrent (forkIO, threadDelay)
  7import           Control.Concurrent.Chan (Chan, readChan, dupChan)
  8import           Control.Exception (bracket)
  9
 10import           Snap.Types
 11import           Snap.Util.FileServe (serveFile, serveDirectory)
 12import           Snap.Http.Server( quickHttpServe)
 13
 14import           Data.ByteString(ByteString)
 15import qualified Data.ByteString.Char8 as BS
 16import           Data.UString (UString, u)
 17import qualified Data.UString as US
 18import           Data.Digest.Pure.SHA (sha1, bytestringDigest)
 19import           Data.Time.Clock.POSIX (POSIXTime)
 20import           Blaze.ByteString.Builder(fromByteString)
 21
 22import qualified System.UUID.V4 as UUID
 23
 24import           AMQPEvents(AMQPEvent(..), Channel, openEventChannel, publishEvent)
 25import           EventStream(ServerEvent(..), eventSourceStream, eventSourceResponse)
 26
 27import           DB
 28
 29import qualified Models.Connection as Conn
 30import qualified Models.User as User
 31
 32import           System.Posix.Env(getEnvDefault)
 33import           Data.Time.Clock.POSIX (getPOSIXTime)
 34
 35import           Text.StringTemplate
 36
 37
 38-- |Setup a channel listening to an AMQP exchange and start Snap
 39main :: IO ()
 40main = do
 41    uuid      <- fmap (u . show) UUID.uuid
 42    origin    <- getEnvDefault "ORIGIN" "http://127.0.0.1"
 43    templates <- directoryGroup "templates" :: IO (STGroup ByteString)
 44
 45    let queue = US.append "eventsource." uuid
 46    let Just js = fmap (render . (setAttribute "origin" origin)) (getStringTemplate "eshq.js" templates)
 47
 48    (publisher, listener) <- openEventChannel (show queue)
 49
 50    bracket openDB (\db -> Conn.remove db uuid >> closeDB db) $ \db -> do
 51        forkIO $ connectionSweeper db uuid
 52        quickHttpServe $
 53            ifTop (serveFile "static/index.html") <|>
 54            path "iframe" (serveFile "static/iframe.html") <|>
 55            path "es.js" (serveJS js) <|>
 56            dir "static" (serveDirectory "static") <|>
 57            method POST (route [ 
 58                ("event", postEvent db publisher queue),
 59                ("socket", createSocket db uuid),
 60                ("socket/:socket", postEventFromSocket db publisher queue)
 61            ]) <|>
 62            method GET (route [
 63                ("broker", brokerInfo db uuid),
 64                ("eventsource", eventSource db uuid listener) 
 65            ])
 66
 67
 68-- |Clean up disconnected connections for this broker at regular intervals
 69connectionSweeper :: DB -> UString -> IO ()
 70connectionSweeper db uuid = do
 71    threadDelay 15000000
 72    Conn.sweep db uuid
 73    connectionSweeper db uuid
 74
 75
 76brokerInfo :: DB -> UString -> Snap ()
 77brokerInfo db uuid = do
 78    result <- liftIO $ Conn.count db uuid
 79    case result of
 80        Right count ->
 81            sendJSON $ BS.pack $ "{\"brokerId\": " ++ (show uuid) ++ ", \"connections\": " ++ (show count) ++ "}"
 82        Left e -> do
 83            modifyResponse $ setResponseCode 500
 84            writeBS $ BS.pack $ "Database Connection Problem: " ++ (show e)
 85
 86-- |Create a new socket and return the ID
 87createSocket :: DB -> UString -> Snap ()
 88createSocket db uuid = do
 89    withAuth db $ \user -> do
 90      withParam "channel" $ \channel -> do
 91        socketId   <- liftIO $ fmap show UUID.uuid
 92        presenceId <- getParam "presence_id"
 93        result     <- liftIO $ Conn.store db Conn.Connection {
 94              Conn.socketId     = u socketId
 95            , Conn.brokerId     = uuid
 96            , Conn.userId       = User.apiKey user
 97            , Conn.channel      = channel
 98            , Conn.presenceId   = fmap ufrombs presenceId
 99            , Conn.disconnectAt = Just 10
100        }
101        case result of
102          Left failure -> do
103              logError (BS.pack $ show failure)
104              showError 500 "Database Connection Error"
105          Right _ ->
106              sendJSON $ BS.pack ("{\"socket\": \"" ++ socketId ++ "\"}")
107
108
109postEvent :: DB -> Channel -> UString -> Snap ()
110postEvent db chan queue =
111    withAuth db $ \user ->
112      withParam "channel" $ \channel ->
113          withParam "data" $ \dataParam -> do
114              liftIO $ publishEvent chan (show queue) $
115                  AMQPEvent (utobs channel) (utobs $ User.apiKey user) (utobs dataParam) Nothing Nothing
116              writeBS "Ok"
117
118-- |Post a new event from a socket.
119postEventFromSocket :: DB -> Channel -> UString -> Snap ()
120postEventFromSocket db chan queue =
121    withConnection db $ \conn ->
122        withParam "data" $ \dataParam -> do
123            liftIO $ publishEvent chan (show queue) $ 
124                AMQPEvent (utobs $ Conn.channel conn) (utobs $ Conn.userId conn) (utobs dataParam) Nothing Nothing
125            writeBS "Ok"
126
127
128-- |Stream events from a channel of AMQPEvents to EventSource
129eventSource :: DB -> UString -> Chan AMQPEvent -> Snap ()
130eventSource db uuid chan = do
131    chan'   <- liftIO $ dupChan chan
132    withConnection db $ \conn -> do
133      liftIO $ before conn
134      transport <- getTransport
135      transport (filterEvents conn chan') (after conn)
136  where
137    before conn = Conn.store db conn { Conn.brokerId = uuid } >> return ()
138    after conn = Conn.mark db (conn { Conn.disconnectAt = Just 10 } ) >> return ()
139
140serveJS :: ByteString -> Snap ()
141serveJS js = do
142    modifyResponse $ setContentType "text/javascript; charset=UTF-8"
143    writeBS js
144
145
146withParam :: UString -> (UString -> Snap ()) -> Snap ()
147withParam param fn = do
148    param' <- getParam (utobs param)
149    case param' of
150        Just value -> fn (ufrombs value)
151        Nothing    -> showError 400 $ BS.concat ["Missing param: ", utobs param]
152
153
154withConnection :: DB -> (Conn.Connection -> Snap ()) -> Snap ()
155withConnection db fn = do
156    withParam "socket" $ \sid -> do
157        withDBResult (Conn.get db sid) (showError 404 "Socket Not Found") fn
158
159
160withAuth :: DB -> (User.User -> Snap ()) -> Snap ()
161withAuth db handler = do
162  key       <- getParam "key"
163  token     <- getParam "token"
164  timestamp <- getParam "timestamp"
165  case (key, token, timestamp) of
166    (Just key', Just token', Just timestamp') -> do
167      currentTime <- liftIO getPOSIXTime
168      withDBResult (User.get db (ufrombs key')) (showError 404 "User not found") $ \user ->
169          if validTime timestamp' currentTime && User.authenticate user token' timestamp'
170            then handler user
171            else showError 401 "Access Denied"
172
173
174withDBResult :: IO (Either Failure (Maybe a)) -> Snap () -> (a -> Snap ()) -> Snap ()
175withDBResult f notFound found= do
176    result <- liftIO f
177    case result of
178      Right (Just model) -> found model
179      Right Nothing      -> notFound
180      Left  failure      -> do
181          logError (BS.pack $ show failure)
182          showError 500 "Database Connection Error"
183
184
185validTime :: ByteString -> POSIXTime -> Bool
186validTime timestamp currentTime =
187    let t1 = read $ BS.unpack timestamp
188        t2 = floor currentTime in
189        abs (t1 - t2) < 5 * 60
190
191
192showError :: Int -> ByteString -> Snap ()
193showError code msg = do
194    modifyResponse $ setResponseCode code
195    writeBS msg
196    r <- getResponse
197    finishWith r
198
199
200sendJSON :: ByteString -> Snap ()
201sendJSON json = do
202    modifyResponse $ setContentType "application/json"
203    writeBS json
204
205
206-- |Returns the transport method to use for this request
207getTransport :: Snap (IO ServerEvent -> IO () -> Snap ())
208getTransport = withRequest $ \request ->
209    case getHeader "X-Requested-With" request of
210      Just "XMLHttpRequest" -> return eventSourceResponse
211      _                     -> return eventSourceStream
212
213
214-- |Filter AMQPEvents by channelId
215filterEvents :: Conn.Connection -> Chan AMQPEvent -> IO ServerEvent
216filterEvents conn chan = do
217    event <- readChan chan
218    if amqpUser event == userId && amqpChannel event == channel
219        then return $ ServerEvent (toBS $ amqpName event) (toBS $ amqpId event) [fromByteString $ amqpData event]
220        else filterEvents conn chan
221  where
222    toBS    = fmap fromByteString
223    userId  = utobs $ Conn.userId conn
224    channel = utobs $ Conn.channel conn
225
226
227ufrombs :: ByteString -> UString
228ufrombs = US.fromByteString_
229
230
231utobs :: UString -> ByteString
232utobs = US.toByteString