@@ -5,26 +5,41 @@ module Network.Transport.TCP
55
66import Network.Transport
77
8+import Control.Applicative
9+
810import Control.Concurrent (forkIO, ThreadId, killThread)
911import Control.Concurrent.Chan
1012import Control.Concurrent.MVar
1113import Control.Monad (forever, forM_)
12−import Data.ByteString.Lazy.Char8 (ByteString)
14+
15+import Data.ByteString (ByteString)
16+import qualified Data.ByteString as BS
17+
18+import Data.Conduit
19+import qualified Data.Conduit.Binary as CB
20+import Data.Conduit.Network
21+
1322import Data.IntMap (IntMap)
23+import qualified Data.IntMap as IntMap
24+
1425import Data.Int
26+import Data.Serialize
1527import Data.Word
28+
1629import Network.Socket
1730 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
1831 , SocketType (Stream), SocketOption (ReuseAddr)
1932 , accept, addrAddress, addrFlags, addrFamily, bindSocket, defaultProtocol
2033 , getAddrInfo, listen, setSocketOption, socket, sClose, withSocketsDo )
34+import qualified Network.Socket as N
35+import qualified Network.Socket.ByteString as NBS
36+
37+import Text.Printf
38+
2139import Safe
2240
2341import qualified Data.Binary as B
24−import qualified Data.ByteString.Lazy.Char8 as BS
25−import qualified Data.IntMap as IntMap
26−import qualified Network.Socket as N
27−import qualified Network.Socket.ByteString.Lazy as NBS
42+
2843
2944type ChanId = Int
3045type Chans = MVar (ChanId, IntMap (Chan [ByteString], [(ThreadId, Socket)]))
@@ -65,7 +80,8 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
6580 return (mkSourceAddr host service chanId, mkTargetEnd chans chanId chan)
6681 , newMulticastWith = error "newMulticastWith: not defined"
6782 , deserialize = {-# SCC "deserialize" #-} \bs ->
68− let (host, service, chanId) = B.decode bs in
83+ let (host, service, chanId) =
84+ either (error $ printf "deserialize: %s") id $ decode bs in
6985 Just $ mkSourceAddr host service chanId
7086 , closeTransport = {-# SCC "closeTransport" #-} do
7187 -- Kill the transport channel process
@@ -82,28 +98,30 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
8298mkSourceAddr :: HostName -> ServiceName -> ChanId -> SourceAddr
8399mkSourceAddr host service chanId = SourceAddr
84100 { connectWith = {-# SCC "connectWith" #-} \_ -> mkSourceEnd host service chanId
85− , serialize = {-# SCC "serialize" #-} B.encode (host, service, chanId)
101+ , serialize = {-# SCC "serialize" #-} encode (host, service, chanId)
86102 }
87103
88104mkSourceEnd :: HostName -> ServiceName -> ChanId -> IO SourceEnd
89105mkSourceEnd host service chanId = withSocketsDo $ do
106+ let err = error $ printf "mkSourceEnd: %s"
90107 serverAddrs <- getAddrInfo Nothing (Just host) (Just service)
91108 let serverAddr = case serverAddrs of
92− [] -> error "mkSourceEnd: getAddrInfo returned []"
109+ [] -> err "getAddrInfo returned []"
93110 as -> head as
94111 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
95112 setSocketOption sock ReuseAddr 1
96113 N.connect sock (addrAddress serverAddr)
97− NBS.sendAll sock $ B.encode (fromIntegral chanId :: Int64)
114+-- let sink = sinkSocket sock :: Sink ByteString IO ()
115+ NBS.sendAll sock $ encode (fromIntegral chanId :: Int64)
98116 return SourceEnd
99117 { send = {-# SCC "send" #-} \bss -> do
100118 let size = fromIntegral (sum . map BS.length $ bss) :: Int64
101119 if size < 255
102120 then
103− NBS.sendAll sock (B.encode (fromIntegral size :: Word8))
121+ NBS.sendAll sock (encode (fromIntegral size :: Word8))
104122 else do
105− NBS.sendAll sock (B.encode (255 :: Word8))
106− NBS.sendAll sock (B.encode size)
123+ NBS.sendAll sock (encode (255 :: Word8))
124+ NBS.sendAll sock (encode size)
107125 mapM_ (NBS.sendAll sock) bss
108126 , closeSourceEnd = {-# SCC "closeSourceEnd" #-} sClose sock
109127 }
@@ -128,19 +146,20 @@ mkTargetEnd chans chanId chan = TargetEnd
128146-- for some reason, an error is raised.
129147procConnections :: Chans -> Socket -> IO ()
130148procConnections chans sock = forever $ do
149+ let err = error $ printf "procConnections: %s"
131150 (clientSock, _clientAddr) <- accept sock
132151 -- decode the first message to find the correct chanId
133152 bss <- recvExact clientSock 8
134153 case bss of
135− [] -> error "procConnections: inbound chanId aborted"
154+ [] -> err "inbound chanId aborted"
136155 bss -> do
137156 let bs = BS.concat bss
138− let chanId = fromIntegral (B.decode bs :: Int64)
157+ let chanId = either err fromIntegral (decode bs :: Either String Int64)
139158 (chanId', chanMap) <- takeMVar chans
140159 case IntMap.lookup chanId chanMap of
141160 Nothing -> do
142161 putMVar chans (chanId', chanMap)
143− error "procConnections: cannot find chanId"
162+ err "cannot find chanId"
144163 Just (chan, socks) -> do
145164 threadId <- forkIO $ procMessages chans chanId chan clientSock
146165 let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
@@ -159,29 +178,30 @@ procMessages chans chanId chan sock = do
159178 case sizeBSs of
160179 [] -> closeSocket
161180 _ -> do
162− let size = fromIntegral (B.decode . BS.concat $ sizeBSs :: Word8)
163− if size == 255
164− then do
181+ case decode . BS.concat $ sizeBSs :: Either String Word8 of
182+ Right 255 -> do
165183 sizeBSs' <- recvExact sock 8
166184 case sizeBSs' of
167185 [] -> closeSocket
168− _ -> procMessage (B.decode . BS.concat $ sizeBSs' :: Int64)
169− else procMessage size
186+ _ -> procMessage (decode . BS.concat $ sizeBSs')
187+ esize -> procMessage (fromIntegral <$> esize)
170188 where
189+ err = error $ printf "procMessages: %s"
171190 closeSocket :: IO ()
172191 closeSocket = do
173192 (chanId', chanMap) <- takeMVar chans
174193 case IntMap.lookup chanId chanMap of
175194 Nothing -> do
176195 putMVar chans (chanId', chanMap)
177− error "procMessages: chanId not found."
196+ err "chanId not found."
178197 Just (chan, socks) -> do
179198 let socks' = filter ((/= sock) . snd) socks
180199 let chanMap' = IntMap.insert chanId (chan, socks') chanMap
181200 putMVar chans (chanId', chanMap')
182201 sClose sock
183− procMessage :: Int64 -> IO ()
184− procMessage size = do
202+ procMessage :: Either String Int64 -> IO ()
203+ procMessage (Left msg) = err msg
204+ procMessage (Right size) = do
185205 bss <- recvExact sock size
186206 case bss of
187207 [] -> closeSocket
@@ -202,13 +222,13 @@ procMessages chans chanId chan sock = do
202222-- indicates success. This hasn't been implemented, since a Transport `receive`
203223-- represents an atomic receipt of a message.
204224recvExact :: Socket -> Int64 -> IO [ByteString]
205−recvExact sock n = go [] sock n
225+recvExact sock n = go [] sock (fromIntegral n)
206226 where
207− go :: [ByteString] -> Socket -> Int64 -> IO [ByteString]
227+ go :: [ByteString] -> Socket -> Int -> IO [ByteString]
208228 go bss _ 0 = return (reverse bss)
209229 go bss sock n = do
210230 bs <- NBS.recv sock n
211231 if BS.null bs
212232 then return []
213− else go (bs:bss) sock (n - BS.length bs)
233+ else go (bs:bss) sock (n - (fromIntegral $ BS.length bs))
214234