-- adopted from Bardur Arantsson's MIT mongrel2-handler
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE OverloadedStrings #-}
module Hack2.Handler.Mongrel2.IO where
import Blaze.ByteString.Builder (Builder,fromByteString)
import Data.Attoparsec
import Hack2.Handler.Mongrel2.ResponseParser (messageParser)
import Hack2.Handler.Mongrel2.Response
import Hack2.Handler.Mongrel2.Types
import qualified System.ZMQ as ZMQ
import Data.ByteString.Char8 (ByteString, unpack)
import Control.Monad (forM_)
import qualified Data.ByteString.Char8 as B
import Data.Enumerator hiding (map)
import Air.Env hiding (log)
import Prelude ()
import Hack2.Handler.Mongrel2.Utils
import Data.Attoparsec
import Data.Aeson (Value(..), json)
import Data.Text.Encoding (encodeUtf8)
import qualified Data.Aeson as Aeson
import Data.Maybe (listToMaybe, fromMaybe, isJust, fromJust)
import qualified Data.Map as Map
import Control.Concurrent.STM
import Data.Int
import qualified Data.Set as Set
-- | Create a new handler. @mkHandler pullFromAddress publishToAddress id@
-- creates a handler which pulls requests from @fromAddress@ and publishes
-- replies to @publishAddress@.
mkHandler :: String -> String -> Maybe String -> Handler
mkHandler = Handler
-- | Run an IO action with a connected handler.
withConnectedHandler :: Handler -> (ConnectedHandler -> IO a) -> IO a
withConnectedHandler h io =
ZMQ.withContext 1 - \c ->
ZMQ.withSocket c ZMQ.Pull - \s -> do
ZMQ.withSocket c ZMQ.Pub - \ps -> do
-- Connect to the poll socket.
ZMQ.connect s - handlerPullFrom h
-- Bind to publish socket.
ZMQ.connect ps - handlerPublishTo h
case handlerId h of
Nothing -> return ()
Just uuid -> ZMQ.setOption ps - ZMQ.Identity uuid
-- Run the action with the connected handler.
io - ConnectedHandler s ps
-- | Receive a parsed request from the Mongrel2 server. Blocks until
-- a message is received.
receiveRequest :: ConnectedHandler -> IO Request
receiveRequest handler = do
msg <- ZMQ.receive (chPullSocket handler) []
log - "ZeroMQ received: " + unpack msg
-- Parse into a structured message.
case parseOnly messageParser msg of
Left errMsg ->
-- Invalid message. This can only happen if there's
-- an error in our parsing code or Mongrel2 itself.
fail - "Couldn't parse message: " + (show errMsg)
Right m -> do
-- Return the parsed message.
return m
jsonToList :: Value -> [(ByteString, ByteString)]
jsonToList value =
let
maybe_string (Aeson.String x) = Just - x
maybe_string _ = Nothing
_headers =
case value of
Aeson.Object headerMap ->
headerMap
.Map.toAscList
.map_snd maybe_string
.select (snd > isJust)
.map (\(x,y) -> (x.encodeUtf8, fromJust y.encodeUtf8))
_ -> []
in
_headers
requestJsonBody :: Request -> Maybe [(ByteString, ByteString)]
requestJsonBody request =
let
_headers = request.requestHeaders.jsonToList
_method = _headers.lookup "METHOD" .fromMaybe "JSON"
in
if _method.is "JSON"
then
let msg = request.requestBody
in
case parseOnly json - msg of
Left errMsg ->
-- Invalid message. This can only happen if there's
-- an error in our parsing code or Mongrel2 itself.
-- log_error - "Couldn't parse message: " + (show errMsg)
Just []
Right m -> do
-- Return the parsed message.
Just - m.jsonToList
else
Nothing
sendResponse :: ConnectedHandler -> TVar (Set.Set Int64) -> Response -> IO ()
sendResponse handler disconnected response =
run_ - _responseEnum $$ _send_iteratee
where
uuid = responseUuid response
clientId = responseClientId response
_responseEnum = responseBody response
_mk_response msg = mkResponse uuid clientId msg
_socket = chPublishSocket handler
-- if put SndMore, memory hogs at socket buffer
_send msg = ZMQ.send _socket (_mk_response msg) [] -- [ZMQ.SndMore]
-- we don't need done, since according to HTTP, ocntent-length should always be set
-- which causes connection to be auto disconnected
_done = ZMQ.send _socket (_mk_response "") []
-- head :: Monad m => Iteratee B.ByteString m (Maybe Word8)
-- head = continue loop where
-- loop (Chunks xs) = case BL.uncons (BL.fromChunks xs) of
-- Just (char, extra) -> yield (Just char) (toChunks extra)
-- Nothing -> head
-- loop EOF = yield Nothing EOF
--
_send_iteratee :: Iteratee ByteString IO ()
_send_iteratee = continue loop where
loop (Chunks xs) = do
_ids <- io - atomically - readTVar disconnected
if _ids.has clientId
then do
io - log - "client disconnected, ignore rest of message"
else do
xs.mapM_ (\msg -> do
io - _send - msg
-- io - log - "ZeroMQ sent: " + unpack msg
)
_send_iteratee
loop EOF = do
yield () EOF
-- let combined = B.append _responseHeader _responseBody
-- putStrLn - "message: " + unpack combined
-- ZMQ.send _socket combined []