{-# LANGUAGE RankNTypes #-}
module Pinch.Transport
( Transport(..)
, framedTransport
, unframedTransport
, Connection(..)
, ReadResult(..)
) where
import Data.IORef (newIORef, readIORef, writeIORef)
import Network.Socket (Socket)
import Network.Socket.ByteString (sendAll, recv)
import System.IO (Handle)
import qualified Data.ByteString as BS
import qualified Data.Serialize.Get as G
import qualified Pinch.Internal.Builder as B
class Connection c where
cGetSome :: c -> IO BS.ByteString
cPut :: c -> B.Builder -> IO ()
instance Connection Handle where
cPut :: Handle -> Builder -> IO ()
cPut Handle
c Builder
b = Handle -> ByteString -> IO ()
BS.hPut Handle
c (Builder -> ByteString
B.runBuilder Builder
b)
cGetSome :: Handle -> IO ByteString
cGetSome Handle
h = Handle -> Int -> IO ByteString
BS.hGetSome Handle
h Int
1024
instance Connection Socket where
cPut :: Socket -> Builder -> IO ()
cPut Socket
c Builder
b = Socket -> ByteString -> IO ()
sendAll Socket
c (Builder -> ByteString
B.runBuilder Builder
b)
cGetSome :: Socket -> IO ByteString
cGetSome Socket
s = Socket -> Int -> IO ByteString
recv Socket
s Int
4096
data ReadResult a
= RRSuccess a
| RRFailure String
| RREOF
deriving (ReadResult a -> ReadResult a -> Bool
(ReadResult a -> ReadResult a -> Bool)
-> (ReadResult a -> ReadResult a -> Bool) -> Eq (ReadResult a)
forall a. Eq a => ReadResult a -> ReadResult a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => ReadResult a -> ReadResult a -> Bool
== :: ReadResult a -> ReadResult a -> Bool
$c/= :: forall a. Eq a => ReadResult a -> ReadResult a -> Bool
/= :: ReadResult a -> ReadResult a -> Bool
Eq, Int -> ReadResult a -> ShowS
[ReadResult a] -> ShowS
ReadResult a -> String
(Int -> ReadResult a -> ShowS)
-> (ReadResult a -> String)
-> ([ReadResult a] -> ShowS)
-> Show (ReadResult a)
forall a. Show a => Int -> ReadResult a -> ShowS
forall a. Show a => [ReadResult a] -> ShowS
forall a. Show a => ReadResult a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> ReadResult a -> ShowS
showsPrec :: Int -> ReadResult a -> ShowS
$cshow :: forall a. Show a => ReadResult a -> String
show :: ReadResult a -> String
$cshowList :: forall a. Show a => [ReadResult a] -> ShowS
showList :: [ReadResult a] -> ShowS
Show)
data Transport
= Transport
{ Transport -> Builder -> IO ()
writeMessage :: B.Builder -> IO ()
, Transport -> forall a. Get a -> IO (ReadResult a)
readMessage :: forall a . G.Get a -> IO (ReadResult a)
}
framedTransport :: Connection c => c -> IO Transport
framedTransport :: forall c. Connection c => c -> IO Transport
framedTransport c
c = do
IORef ByteString
readBuffer <- ByteString -> IO (IORef ByteString)
forall a. a -> IO (IORef a)
newIORef ByteString
forall a. Monoid a => a
mempty
Transport -> IO Transport
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Transport -> IO Transport) -> Transport -> IO Transport
forall a b. (a -> b) -> a -> b
$ (Builder -> IO ())
-> (forall a. Get a -> IO (ReadResult a)) -> Transport
Transport Builder -> IO ()
writeMsg (IORef ByteString -> Get a -> IO (ReadResult a)
forall {a}. IORef ByteString -> Get a -> IO (ReadResult a)
readMsg IORef ByteString
readBuffer) where
writeMsg :: Builder -> IO ()
writeMsg Builder
msg = do
c -> Builder -> IO ()
forall c. Connection c => c -> Builder -> IO ()
cPut c
c (Builder -> IO ()) -> Builder -> IO ()
forall a b. (a -> b) -> a -> b
$ Int32 -> Builder
B.int32BE (Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int32) -> Int -> Int32
forall a b. (a -> b) -> a -> b
$ Builder -> Int
B.getSize Builder
msg) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
msg
readMsg :: IORef ByteString -> Get a -> IO (ReadResult a)
readMsg IORef ByteString
readBuffer Get a
parser = do
let
frameParser :: Get a
frameParser = do
Int32
size <- Get Int32
G.getInt32be
Int -> Get a -> Get a
forall a. Int -> Get a -> Get a
G.isolate (Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
size) Get a
parser
ByteString
initial <- IORef ByteString -> IO ByteString
forall a. IORef a -> IO a
readIORef IORef ByteString
readBuffer
(ByteString
leftovers, ReadResult a
r) <- IO ByteString
-> Get a -> ByteString -> IO (ByteString, ReadResult a)
forall a.
IO ByteString
-> Get a -> ByteString -> IO (ByteString, ReadResult a)
runGetWith (c -> IO ByteString
forall c. Connection c => c -> IO ByteString
cGetSome c
c) Get a
frameParser ByteString
initial
IORef ByteString -> ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef ByteString
readBuffer (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$! ByteString
leftovers
ReadResult a -> IO (ReadResult a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ReadResult a
r
unframedTransport :: Connection c => c -> IO Transport
unframedTransport :: forall c. Connection c => c -> IO Transport
unframedTransport c
c = do
IORef ByteString
readBuffer <- ByteString -> IO (IORef ByteString)
forall a. a -> IO (IORef a)
newIORef ByteString
forall a. Monoid a => a
mempty
Transport -> IO Transport
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Transport -> IO Transport) -> Transport -> IO Transport
forall a b. (a -> b) -> a -> b
$ (Builder -> IO ())
-> (forall a. Get a -> IO (ReadResult a)) -> Transport
Transport Builder -> IO ()
writeMsg (IORef ByteString -> Get a -> IO (ReadResult a)
forall {a}. IORef ByteString -> Get a -> IO (ReadResult a)
readMsg IORef ByteString
readBuffer)
where
writeMsg :: Builder -> IO ()
writeMsg = c -> Builder -> IO ()
forall c. Connection c => c -> Builder -> IO ()
cPut c
c
readMsg :: IORef ByteString -> Get a -> IO (ReadResult a)
readMsg IORef ByteString
buf Get a
p = do
ByteString
initial <- IORef ByteString -> IO ByteString
forall a. IORef a -> IO a
readIORef IORef ByteString
buf
(ByteString
leftovers, ReadResult a
r) <- IO ByteString
-> Get a -> ByteString -> IO (ByteString, ReadResult a)
forall a.
IO ByteString
-> Get a -> ByteString -> IO (ByteString, ReadResult a)
runGetWith (c -> IO ByteString
forall c. Connection c => c -> IO ByteString
cGetSome c
c) Get a
p ByteString
initial
IORef ByteString -> ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef ByteString
buf (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$! ByteString
leftovers
ReadResult a -> IO (ReadResult a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ReadResult a
r
runGetWith :: IO BS.ByteString -> G.Get a -> BS.ByteString -> IO (BS.ByteString, ReadResult a)
runGetWith :: forall a.
IO ByteString
-> Get a -> ByteString -> IO (ByteString, ReadResult a)
runGetWith IO ByteString
getBs Get a
p ByteString
initial = Result a -> IO (ByteString, ReadResult a)
forall {a}. Result a -> IO (ByteString, ReadResult a)
go (Get a -> ByteString -> Result a
forall a. Get a -> ByteString -> Result a
G.runGetPartial Get a
p ByteString
initial)
where
go :: Result a -> IO (ByteString, ReadResult a)
go Result a
r = case Result a
r of
G.Fail String
err ByteString
bs -> do
(ByteString, ReadResult a) -> IO (ByteString, ReadResult a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString
bs, String -> ReadResult a
forall a. String -> ReadResult a
RRFailure String
err)
G.Done a
a ByteString
bs -> do
(ByteString, ReadResult a) -> IO (ByteString, ReadResult a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString
bs, a -> ReadResult a
forall a. a -> ReadResult a
RRSuccess a
a)
G.Partial ByteString -> Result a
cont -> do
ByteString
bs <- IO ByteString
getBs
if ByteString -> Bool
BS.null ByteString
bs
then
(ByteString, ReadResult a) -> IO (ByteString, ReadResult a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString
bs, ReadResult a
forall a. ReadResult a
RREOF)
else
Result a -> IO (ByteString, ReadResult a)
go (Result a -> IO (ByteString, ReadResult a))
-> Result a -> IO (ByteString, ReadResult a)
forall a b. (a -> b) -> a -> b
$ ByteString -> Result a
cont ByteString
bs