@@ -1,3 +1,5 @@
1+{-# LANGUAGE CPP #-}
2+
13module Network.Transport.TCP
24 ( mkTransport
35 , TCPConfig (..)
@@ -5,30 +7,58 @@ module Network.Transport.TCP
57
68import Network.Transport
79
10+import Control.Applicative
811import Control.Concurrent (forkIO, ThreadId, killThread, myThreadId)
912import Control.Concurrent.Chan
1013import Control.Concurrent.MVar
1114import Control.Exception (SomeException, IOException, AsyncException(ThreadKilled),
1215 fromException, throwTo, throw, catch, handle)
1316import Control.Monad (forever, forM_)
14−import Data.ByteString.Lazy.Char8 (ByteString)
17+
18+import qualified Data.ByteString.Internal as BSI
19+
1520import Data.IntMap (IntMap)
16−import Data.Int
17−import Data.Word
21+import qualified Data.IntMap as IntMap
22+
23+import Data.Int (Int32)
24+import qualified Data.Serialize as Ser
25+import Data.Word (Word8)
26+
27+import Foreign.ForeignPtr (withForeignPtr)
28+import Foreign.Ptr (Ptr, plusPtr)
29+import Foreign.Storable (sizeOf)
30+
1831import Network.Socket
1932 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
2033 , SocketType (Stream), SocketOption (ReuseAddr)
2134 , accept, addrAddress, addrFlags, addrFamily, bindSocket, defaultProtocol
2235 , getAddrInfo, listen, setSocketOption, socket, sClose, withSocketsDo )
36+import qualified Network.Socket as N
37+
38+import Text.Printf
39+
2340import Safe
2441import System.IO (stderr, hPutStrLn)
2542
26−import qualified Data.Binary as B
27−import qualified Data.ByteString.Char8 as BSS
43+#ifndef LAZY
44+import Data.ByteString.Char8 (ByteString)
45+import qualified Data.ByteString.Char8 as BS
46+import qualified Network.Socket.ByteString as NBS
47+sendNBS = NBS.sendMany
48+encode = Ser.encode
49+decode = Ser.decode
50+#else
51+import Data.ByteString.Lazy.Char8 (ByteString)
2852import qualified Data.ByteString.Lazy.Char8 as BS
29−import qualified Data.IntMap as IntMap
30−import qualified Network.Socket as N
3153import qualified Network.Socket.ByteString.Lazy as NBS
54+sendNBS sock = NBS.sendAll sock . BS.concat
55+encode = Ser.encodeLazy
56+decode = Ser.decodeLazy
57+#endif
58+{-# INLINE encode #-}
59+{-# INLINE decode #-}
60+encode :: Ser.Serialize a => a -> ByteString
61+decode :: Ser.Serialize a => ByteString -> Either String a
3262
3363type ChanId = Int
3464type Chans = MVar (ChanId, IntMap (Chan [ByteString], [(ThreadId, Socket)]))
@@ -70,7 +100,9 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
70100 return (mkSourceAddr host service chanId, mkTargetEnd chans chanId chan)
71101 , newMulticastWith = error "newMulticastWith: not defined"
72102 , deserialize = {-# SCC "deserialize" #-} \bs ->
73− let (host, service, chanId) = B.decode bs in
103+ let (host, service, chanId) =
104+ either (\m -> error $ printf "deserialize: %s" m)
105+ id $ decode bs in
74106 Just $ mkSourceAddr host service chanId
75107 , closeTransport = {-# SCC "closeTransport" #-} do
76108 -- Kill the transport channel process
@@ -87,29 +119,26 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
87119mkSourceAddr :: HostName -> ServiceName -> ChanId -> SourceAddr
88120mkSourceAddr host service chanId = SourceAddr
89121 { connectWith = {-# SCC "connectWith" #-} \_ -> mkSourceEnd host service chanId
90− , serialize = {-# SCC "serialize" #-} B.encode (host, service, chanId)
122+ , serialize = {-# SCC "serialize" #-} encode (host, service, chanId)
91123 }
92124
93125mkSourceEnd :: HostName -> ServiceName -> ChanId -> IO SourceEnd
94126mkSourceEnd host service chanId = withSocketsDo $ do
127+ let err m = error $ printf "mkSourceEnd: %s" m
95128 serverAddrs <- getAddrInfo Nothing (Just host) (Just service)
96129 let serverAddr = case serverAddrs of
97− [] -> error "mkSourceEnd: getAddrInfo returned []"
130+ [] -> err "getAddrInfo returned []"
98131 as -> head as
99132 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
100133 setSocketOption sock ReuseAddr 1
101134 N.connect sock (addrAddress serverAddr)
102− NBS.sendAll sock $ B.encode (fromIntegral chanId :: Int64)
135+ NBS.sendAll sock $ encode (fromIntegral chanId :: Int32)
103136 return SourceEnd
104137 { send = {-# SCC "send" #-} \bss -> do
105− let size = fromIntegral (sum . map BS.length $ bss) :: Int64
106− if size < 255
107− then
108− NBS.sendAll sock (B.encode (fromIntegral size :: Word8))
109− else do
110− NBS.sendAll sock (B.encode (255 :: Word8))
111− NBS.sendAll sock (B.encode size)
112− mapM_ (NBS.sendAll sock) bss
138+ let size = fromIntegral (sum . map BS.length $ bss) :: Int32
139+ sizeStr | size < 255 = BS.singleton (toEnum $ fromEnum size)
140+ | otherwise = BS.cons (toEnum 255) (encode size)
141+ sendNBS sock (sizeStr:bss)
113142 , closeSourceEnd = {-# SCC "closeSourceEnd" #-} sClose sock
114143 }
115144
@@ -133,19 +162,19 @@ mkTargetEnd chans chanId chan = TargetEnd
133162-- for some reason, an error is raised.
134163procConnections :: Chans -> Socket -> IO ()
135164procConnections chans sock = forever $ do
165+ let err m = error $ printf "procConnections: %s" m
136166 (clientSock, _clientAddr) <- accept sock
137167 -- decode the first message to find the correct chanId
138− bss <- recvExact clientSock 8
168+ bss <- recvExact clientSock (fromIntegral $ sizeOf (0 :: Int32))
139169 case bss of
140− [] -> error "procConnections: inbound chanId aborted"
141− bss -> do
142− let bs = BS.concat bss
143− let chanId = fromIntegral (B.decode bs :: Int64)
170+ Left _ -> err "inbound chanId aborted"
171+ Right bs -> do
172+ let chanId = either err fromIntegral (decode bs :: Either String Int32)
144173 (chanId', chanMap) <- takeMVar chans
145174 case IntMap.lookup chanId chanMap of
146175 Nothing -> do
147176 putMVar chans (chanId', chanMap)
148− error "procConnections: cannot find chanId"
177+ err "cannot find chanId"
149178 Just (chan, socks) -> do
150179 threadId <- forkWithExceptions forkIO "Message Listener" $
151180 procMessages chans chanId chan clientSock
@@ -154,80 +183,94 @@ procConnections chans sock = forever $ do
154183
155184-- | This function first extracts a header of type Word8, which determines
156185-- the size of the ByteString that follows. If this size is 0, this indicates
157−-- that the ByteString is large, so the next value is an Int64, which
186+-- that the ByteString is large, so the next value is an Int32, which
158187-- determines the size of the ByteString that follows. The ByteString is then
159188-- extracted from the socket, and then written to the Chan only when
160189-- complete. If either of the first header size is null this indicates the
161190-- socket has closed.
162191procMessages :: Chans -> ChanId -> Chan [ByteString] -> Socket -> IO ()
163192procMessages chans chanId chan sock = do
164− sizeBSs <- recvExact sock 1
165− case sizeBSs of
166− [] -> closeSocket
167− [onebyte] -> do
168− let size = fromIntegral (B.decode onebyte :: Word8)
169− if size == 255
170− then do
171− sizeBSs' <- recvExact sock 8
172− case sizeBSs' of
173− [] -> closeSocket
174− _ -> procMessage (B.decode . BS.concat$ sizeBSs' :: Int64)
175− else procMessage size
176− ls -> error "Shouldn't receive more than one bytestring when expecting a single byte!"
193+ esizeBS <- recvExact sock 1
194+ case esizeBS of
195+ Left _ -> closeSocket
196+ Right sizeBS -> do
197+ case decode sizeBS :: Either String Word8 of
198+ Right 255 -> do
199+ esizeBS' <- recvExact sock (fromIntegral $ sizeOf (0 :: Int32))
200+ case esizeBS' of
201+ Left _ -> closeSocket
202+ Right sizeBS' -> procMessage $ decode sizeBS'
203+ esize@(Right _) -> procMessage (fromIntegral <$> esize)
204+ Left msg -> error msg
177205 where
206+ err m = error $ printf "procMessages: %s" m
178207 closeSocket :: IO ()
179208 closeSocket = do
180209 (chanId', chanMap) <- takeMVar chans
181210 case IntMap.lookup chanId chanMap of
182211 Nothing -> do
183212 putMVar chans (chanId', chanMap)
184− error "procMessages: chanId not found."
213+ err "chanId not found."
185214 Just (chan, socks) -> do
186215 let socks' = filter ((/= sock) . snd) socks
187216 let chanMap' = IntMap.insert chanId (chan, socks') chanMap
188217 putMVar chans (chanId', chanMap')
189218 sClose sock
190− procMessage :: Int64 -> IO ()
191− procMessage size = do
192− bss <- recvExact sock size
193− case bss of
194− [] -> closeSocket
195− _ -> do
196− writeChan chan bss
219+ procMessage :: Either String Int32 -> IO ()
220+ procMessage (Left msg) = err msg
221+ procMessage (Right size) = do
222+ ebs <- recvExact sock size
223+ case ebs of
224+ Left _ -> closeSocket
225+ Right bs -> do
226+ writeChan chan [bs]
197227 procMessages chans chanId chan sock
198228
199−-- | The result of `recvExact sock n` is a `[ByteString]` whose concatenation
200−-- is of length `n`, received from `sock`. No more bytes than necessary are
201−-- read from the socket.
202−-- NB: This uses Network.Socket.ByteString.recv, which may *discard*
203−-- superfluous input depending on the socket type. Also note that
204−-- if `recv` returns an empty `ByteString` then this means that the socket
205−-- was closed: in this case, we return an empty list.
206−-- NB: It may be appropriate to change the return type to
207−-- IO (Either ByteString ByteString), where `return Left bs` indicates
208−-- a partial retrieval since the socket was closed, and `return Right bs`
209−-- indicates success. This hasn't been implemented, since a Transport `receive`
210−-- represents an atomic receipt of a message.
211−recvExact :: Socket -> Int64 -> IO [ByteString]
229+-- | The normal result of `recvExact sock n` is `Right ByteString`
230+-- whose string of length `n`, received from `sock`. If fewer than `n`
231+-- bytes are read from `sock` before it closes, the result is `Left
232+-- ByteString` whose string is those bytes that were received. No more
233+-- bytes than necessary are read from the socket. NB: This uses
234+-- Network.Socket.ByteString.recv, which may *discard* superfluous
235+-- input depending on the socket type.
236+#ifndef LAZY
237+recvExact :: Socket -> Int32 -> IO (Either ByteString ByteString)
238+recvExact sock l = interceptAllExn "recvExact" $ do
239+ res <- BSI.createAndTrim (fromIntegral l) (go 0)
240+ case BS.length res of
241+ n | n == (fromIntegral l) -> return $ Right res
242+ n -> return $ Left res
243+ where
244+ go :: Int -> Ptr Word8 -> IO Int
245+ go n ptr | n == (fromIntegral l) = return n
246+ go n ptr = do
247+ (p, off, len) <- BSI.toForeignPtr <$> NBS.recv sock (min (fromIntegral l-n) 4096)
248+ if len == 0
249+ then return n
250+ else withForeignPtr p $ \p -> do
251+ BSI.memcpy (ptr `plusPtr` n) (p `plusPtr` off) (fromIntegral len)
252+ go (n+len) ptr
253+#else
254+recvExact :: Socket -> Int32 -> IO (Either ByteString ByteString)
212255recvExact sock n =
213− interceptAllExn "recvExact" $
256+ interceptAllExn "recvExact" $
214257 go [] sock n
215258 where
216− go :: [ByteString] -> Socket -> Int64 -> IO [ByteString]
217− go bss _ 0 = return (reverse bss)
259+ go :: [ByteString] -> Socket -> Int32 -> IO (Either ByteString ByteString)
260+ go bss _ 0 = return (Right $ BS.concat (reverse bss))
218261 go bss sock n = do
219− bs <- NBS.recv sock n
262+ bs <- NBS.recv sock (min (fromIntegral n) 4096)
220263 if BS.null bs
221− then return []
222− else go (bs:bss) sock (n - BS.length bs)
223−
264+ then return (Left $ BS.concat (reverse bss))
265+ else go (bs:bss) sock (n - (fromIntegral $ BS.length bs))
266+#endif
224267
225268interceptAllExn msg =
226269 Control.Exception.handle $ \ e ->
227270 case fromException e of
228271 Just ThreadKilled -> throw e
229272 Nothing -> do
230− BSS.hPutStrLn stderr $ BSS.pack$ "Exception inside "++msg++": "++show e
273+ BS.hPutStrLn stderr $ BS.pack$ "Exception inside "++msg++": "++show e
231274 throw e
232275-- throw (e :: SomeException)
233276
@@ -242,7 +285,7 @@ forkWithExceptions forkit descr action = do
242285-- BSS.hPutStrLn stderr $ BSS.pack$ "Note: Child thread killed: "++descr
243286 return ()
244287 _ -> do
245− BSS.hPutStrLn stderr $ BSS.pack$ "Exception inside child thread "++descr++": "++show e
288+ BS.hPutStrLn stderr $ BS.pack$ "Exception inside child thread "++descr++": "++show e
246289 throwTo parent (e::SomeException)
247290 )
248291