mongrel2-handler-0.3.2: src/Mongrel2.hs

module Mongrel2
       ( Handler
       , ConnectedHandler
       , mkHandler
       , withConnectedHandler
       , receiveRequest
       , sendReply
       , simpleReply
       -- Re-exported
       , buildResponse
       , ClientID
       , Connection(..)
       , Request(..)
       , UUID
       ) where

import Blaze.ByteString.Builder (Builder)
import Data.Attoparsec
import Mongrel2.Types (Connection(..), Request(..))
import Mongrel2.Parser (messageParser, UUID, ClientID)
import Mongrel2.Response (buildResponse, mkResponse)
import qualified System.ZMQ as ZMQ

data Handler = Handler { handlerPullFrom :: String
                       , handlerPublishTo :: String
                       , handlerId :: Maybe String
                       }

data ConnectedHandler = ConnectedHandler { chPullSocket :: ZMQ.Socket ZMQ.Pull
                                         , chPublishSocket :: ZMQ.Socket ZMQ.Pub
                                         }

-- | 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. The Mongrel2 server will signal a client
-- disconnect by sending a Request with method @JSON@ and a request
-- body containing a JSON object with a key @type@ containing the value
-- @disconnect@. No response should be sent for such a request.
receiveRequest :: ConnectedHandler -> IO (Connection,Request)
receiveRequest handler = do
      msg <- ZMQ.receive (chPullSocket handler) []
      -- 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

-- | Simplied form of @sendReply@ function. This variant sends a reply
-- to the requesting client with no option of sending replies to other
-- clients.
simpleReply :: ConnectedHandler -> Connection -> Builder -> IO ()
simpleReply handler connection =
  sendReply handler serverUUID [clientID]
  where
    serverUUID = connServerUUID connection
    clientID = connClientID connection

-- | Send a reply to the Mongrel2 server.
sendReply :: ConnectedHandler -> UUID -> [ClientID] -> Builder -> IO ()
sendReply handler uuid clientIds replyBuilder = do
  -- Build the response.
  let response = mkResponse uuid clientIds replyBuilder
  -- Send the response.
  ZMQ.send (chPublishSocket handler) response []