module Thrift.Transport.Wai (
RequestTransport,
StreamTransport,
fromRequest,
toStreamTransport,
thriftWaiApp,
thriftMiddleware
) where
import Thrift.Transport
import Thrift.Protocol
import Thrift.Transport.IOBuffer
import Network.Wai as Wai
import Network.Wai.Internal
import Data.IORef
import Blaze.ByteString.Builder
import Data.Monoid
import Network.HTTP.Types (status200)
import Network.HTTP.Types.Method
data RequestTransport = RequestTransport Request ReadBuffer
fromRequest :: Request -> IO RequestTransport
fromRequest req = RequestTransport
<$> return req
<*> (Wai.lazyRequestBody req >>= newIORef)
instance Transport RequestTransport where
tIsOpen = const $ return True
tClose = const $ return ()
tRead (RequestTransport _ b) n = readBuf b (fromIntegral n)
tPeek (RequestTransport _ b) = peekBuf b
tWrite _ _ = fail "RequestTransport does not support write"
tFlush _ = fail "RequestTransport does not support flush"
type WriteBuilder = IORef Builder
newtype ResponseTransport = ResponseTransport WriteBuilder
newResponseTransport :: IO ResponseTransport
newResponseTransport = ResponseTransport <$> (newIORef mempty)
instance Transport ResponseTransport where
tIsOpen = const $ return False
tClose = const $ return ()
tRead _ _ = fail "Read operation is not supported for response"
tPeek _ = fail "Peek is not allowed for response buffers"
tWrite (ResponseTransport b) bs = modifyIORef b (<> fromLazyByteString bs)
tFlush (ResponseTransport b) = modifyIORef b (<> flush)
data StreamTransport = StreamTransport { writer :: Builder -> IO ()
, flusher :: IO ()
}
toStreamTransport :: (Builder -> IO () ) -> IO () -> StreamTransport
toStreamTransport w f = StreamTransport w f
instance Transport StreamTransport where
tIsOpen = const $ return True
tClose = const $ return ()
tRead _ _ = fail "Read operation is not supported for response"
tPeek _ = fail "Peek is not allowed for response buffers"
tWrite st bs = writer st $ fromLazyByteString bs
tFlush st = flusher st
thriftWaiApp ::
(Protocol ip, Protocol op)
=> h
-> (RequestTransport -> ip RequestTransport)
-> (StreamTransport -> op StreamTransport)
-> (h -> (ip RequestTransport, op StreamTransport) -> IO Bool)
-> Application
thriftWaiApp h isp osp proc_ req responder = do
inp <- isp <$> fromRequest req
responder $ Wai.responseStream status200 [] $ \write flushstream -> do
let out = osp (StreamTransport write flushstream)
_ <- proc_ h (inp, out)
return ()
thriftMiddleware :: (Protocol ip, Protocol op)
=> h
-> (RequestTransport -> ip RequestTransport)
-> (StreamTransport -> op StreamTransport)
-> (h -> (ip RequestTransport, op StreamTransport) -> IO Bool)
-> Application
-> Application
thriftMiddleware h isp osp proc_ app req responder = app req $ \res ->
case res of
ResponseStream {} ->
if methodPost == requestMethod req then
do
inp <- isp <$> fromRequest req
responder $ Wai.responseStream status200 (responseHeaders res) $ \write flushstream -> do
let out = osp (StreamTransport write flushstream)
_ <- proc_ h (inp, out)
return ()
else
responder res
_ -> responder res