-
Notifications
You must be signed in to change notification settings - Fork 59
/
TCP.hs
414 lines (362 loc) · 15.7 KB
/
TCP.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
{-# LANGUAGE TypeSynonymInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module : Network.TCP
-- Copyright : See LICENSE file
-- License : BSD
--
-- Maintainer : Ganesh Sittampalam <ganesh@earth.li>
-- Stability : experimental
-- Portability : non-portable (not tested)
--
-- Some utility functions for working with the Haskell @network@ package. Mostly
-- for internal use by the @Network.HTTP@ code.
--
-----------------------------------------------------------------------------
module Network.TCP
( Connection
, EndPoint(..)
, openTCPPort
, isConnectedTo
, openTCPConnection
, socketConnection
, isTCPConnectedTo
, HandleStream
, HStream(..)
, StreamHooks(..)
, nullHooks
, setStreamHooks
, getStreamHooks
, hstreamToConnection
) where
import Network.Socket
( Socket, SocketOption(KeepAlive)
, SocketType(Stream), connect
, shutdown, ShutdownCmd(..)
, sClose, setSocketOption, getPeerName
, socket, Family(AF_UNSPEC), defaultProtocol, getAddrInfo
, defaultHints, addrFamily, withSocketsDo
, addrSocketType, addrAddress
)
import qualified Network.Stream as Stream
( Stream(readBlock, readLine, writeBlock, close, closeOnEnd) )
import Network.Stream
( ConnError(..)
, Result
, failWith
, failMisc
)
import Network.BufferType
import Network.HTTP.Base ( catchIO )
import Network.Socket ( socketToHandle )
import Data.Char ( toLower )
import Data.Word ( Word8 )
import Control.Concurrent
import Control.Exception ( onException )
import Control.Monad ( liftM, when )
import System.IO ( Handle, hFlush, IOMode(..), hClose )
import System.IO.Error ( isEOFError )
import qualified Data.ByteString as Strict
import qualified Data.ByteString.Lazy as Lazy
-----------------------------------------------------------------
------------------ TCP Connections ------------------------------
-----------------------------------------------------------------
-- | The 'Connection' newtype is a wrapper that allows us to make
-- connections an instance of the Stream class, without GHC extensions.
-- While this looks sort of like a generic reference to the transport
-- layer it is actually TCP specific, which can be seen in the
-- implementation of the 'Stream Connection' instance.
newtype Connection = Connection (HandleStream String)
newtype HandleStream a = HandleStream {getRef :: MVar (Conn a)}
data EndPoint = EndPoint { epHost :: String, epPort :: Int }
instance Eq EndPoint where
EndPoint host1 port1 == EndPoint host2 port2 =
map toLower host1 == map toLower host2 && port1 == port2
data Conn a
= MkConn { connSock :: ! Socket
, connHandle :: Handle
, connBuffer :: BufferOp a
, connInput :: Maybe a
, connEndPoint :: EndPoint
, connHooks :: Maybe (StreamHooks a)
, connCloseEOF :: Bool -- True => close socket upon reaching end-of-stream.
}
| ConnClosed
deriving(Eq)
hstreamToConnection :: HandleStream String -> Connection
hstreamToConnection h = Connection h
connHooks' :: Conn a -> Maybe (StreamHooks a)
connHooks' ConnClosed{} = Nothing
connHooks' x = connHooks x
-- all of these are post-op hooks
data StreamHooks ty
= StreamHooks
{ hook_readLine :: (ty -> String) -> Result ty -> IO ()
, hook_readBlock :: (ty -> String) -> Int -> Result ty -> IO ()
, hook_writeBlock :: (ty -> String) -> ty -> Result () -> IO ()
, hook_close :: IO ()
, hook_name :: String -- hack alert: name of the hook itself.
}
instance Eq ty => Eq (StreamHooks ty) where
(==) _ _ = True
nullHooks :: StreamHooks ty
nullHooks = StreamHooks
{ hook_readLine = \ _ _ -> return ()
, hook_readBlock = \ _ _ _ -> return ()
, hook_writeBlock = \ _ _ _ -> return ()
, hook_close = return ()
, hook_name = ""
}
setStreamHooks :: HandleStream ty -> StreamHooks ty -> IO ()
setStreamHooks h sh = modifyMVar_ (getRef h) (\ c -> return c{connHooks=Just sh})
getStreamHooks :: HandleStream ty -> IO (Maybe (StreamHooks ty))
getStreamHooks h = readMVar (getRef h) >>= return.connHooks
-- | @HStream@ overloads the use of 'HandleStream's, letting you
-- overload the handle operations over the type that is communicated
-- across the handle. It comes in handy for @Network.HTTP@ 'Request'
-- and 'Response's as the payload representation isn't fixed, but overloaded.
--
-- The library comes with instances for @ByteString@s and @String@, but
-- should you want to plug in your own payload representation, defining
-- your own @HStream@ instance _should_ be all that it takes.
--
class BufferType bufType => HStream bufType where
openStream :: String -> Int -> IO (HandleStream bufType)
openSocketStream :: String -> Int -> Socket -> IO (HandleStream bufType)
readLine :: HandleStream bufType -> IO (Result bufType)
readBlock :: HandleStream bufType -> Int -> IO (Result bufType)
writeBlock :: HandleStream bufType -> bufType -> IO (Result ())
close :: HandleStream bufType -> IO ()
closeQuick :: HandleStream bufType -> IO ()
closeOnEnd :: HandleStream bufType -> Bool -> IO ()
instance HStream Strict.ByteString where
openStream = openTCPConnection
openSocketStream = socketConnection
readBlock c n = readBlockBS c n
readLine c = readLineBS c
writeBlock c str = writeBlockBS c str
close c = closeIt c Strict.null True
closeQuick c = closeIt c Strict.null False
closeOnEnd c f = closeEOF c f
instance HStream Lazy.ByteString where
openStream = \ a b -> openTCPConnection_ a b True
openSocketStream = \ a b c -> socketConnection_ a b c True
readBlock c n = readBlockBS c n
readLine c = readLineBS c
writeBlock c str = writeBlockBS c str
close c = closeIt c Lazy.null True
closeQuick c = closeIt c Lazy.null False
closeOnEnd c f = closeEOF c f
instance Stream.Stream Connection where
readBlock (Connection c) = Network.TCP.readBlock c
readLine (Connection c) = Network.TCP.readLine c
writeBlock (Connection c) = Network.TCP.writeBlock c
close (Connection c) = Network.TCP.close c
closeOnEnd (Connection c) f = Network.TCP.closeEOF c f
instance HStream String where
openStream = openTCPConnection
openSocketStream = socketConnection
readBlock ref n = readBlockBS ref n
-- This function uses a buffer, at this time the buffer is just 1000 characters.
-- (however many bytes this is is left to the user to decypher)
readLine ref = readLineBS ref
-- The 'Connection' object allows no outward buffering,
-- since in general messages are serialised in their entirety.
writeBlock ref str = writeBlockBS ref str -- (stringToBuf str)
-- Closes a Connection. Connection will no longer
-- allow any of the other Stream functions. Notice that a Connection may close
-- at any time before a call to this function. This function is idempotent.
-- (I think the behaviour here is TCP specific)
close c = closeIt c null True
-- Closes a Connection without munching the rest of the stream.
closeQuick c = closeIt c null False
closeOnEnd c f = closeEOF c f
-- | @openTCPPort uri port@ establishes a connection to a remote
-- host, using 'getHostByName' which possibly queries the DNS system, hence
-- may trigger a network connection.
openTCPPort :: String -> Int -> IO Connection
openTCPPort uri port = openTCPConnection uri port >>= return.Connection
-- Add a "persistent" option? Current persistent is default.
-- Use "Result" type for synchronous exception reporting?
openTCPConnection :: BufferType ty => String -> Int -> IO (HandleStream ty)
openTCPConnection uri port = openTCPConnection_ uri port False
openTCPConnection_ :: BufferType ty => String -> Int -> Bool -> IO (HandleStream ty)
openTCPConnection_ uri port stashInput = do
-- HACK: uri is sometimes obtained by calling Network.URI.uriRegName, and this includes
-- the surrounding square brackets for an RFC 2732 host like [::1]. It's not clear whether
-- it should, or whether all call sites should be using something different instead, but
-- the simplest short-term fix is to strip any surrounding square brackets here.
-- It shouldn't affect any as this is the only situation they can occur - see RFC 3986.
let fixedUri =
case uri of
'[':(rest@(c:_)) | last rest == ']'
-> if c == 'v' || c == 'V'
then error $ "Unsupported post-IPv6 address " ++ uri
else init rest
_ -> uri
-- use withSocketsDo here in case the caller hasn't used it, which would make getAddrInfo fail on Windows
-- although withSocketsDo is supposed to wrap the entire program, in practice it is safe to use it locally
-- like this as it just does a once-only installation of a shutdown handler to run at program exit,
-- rather than actually shutting down after the action
addrinfos <- withSocketsDo $ getAddrInfo (Just $ defaultHints { addrFamily = AF_UNSPEC, addrSocketType = Stream }) (Just fixedUri) (Just . show $ port)
case addrinfos of
[] -> fail "openTCPConnection: getAddrInfo returned no address information"
(a:_) -> do
s <- socket (addrFamily a) Stream defaultProtocol
onException (do
setSocketOption s KeepAlive 1
connect s (addrAddress a)
socketConnection_ fixedUri port s stashInput
) (sClose s)
-- | @socketConnection@, like @openConnection@ but using a pre-existing 'Socket'.
socketConnection :: BufferType ty
=> String
-> Int
-> Socket
-> IO (HandleStream ty)
socketConnection hst port sock = socketConnection_ hst port sock False
-- Internal function used to control the on-demand streaming of input
-- for /lazy/ streams.
socketConnection_ :: BufferType ty
=> String
-> Int
-> Socket
-> Bool
-> IO (HandleStream ty)
socketConnection_ hst port sock stashInput = do
h <- socketToHandle sock ReadWriteMode
mb <- case stashInput of { True -> liftM Just $ buf_hGetContents bufferOps h; _ -> return Nothing }
let conn = MkConn
{ connSock = sock
, connHandle = h
, connBuffer = bufferOps
, connInput = mb
, connEndPoint = EndPoint hst port
, connHooks = Nothing
, connCloseEOF = False
}
v <- newMVar conn
return (HandleStream v)
closeConnection :: HStream a => HandleStream a -> IO Bool -> IO ()
closeConnection ref readL = do
-- won't hold onto the lock for the duration
-- we are draining it...ToDo: have Connection
-- into a shutting-down state so that other
-- threads will simply back off if/when attempting
-- to also close it.
c <- readMVar (getRef ref)
closeConn c `catchIO` (\_ -> return ())
modifyMVar_ (getRef ref) (\ _ -> return ConnClosed)
where
-- Be kind to peer & close gracefully.
closeConn ConnClosed = return ()
closeConn conn = do
let sk = connSock conn
hFlush (connHandle conn)
shutdown sk ShutdownSend
suck readL
hClose (connHandle conn)
shutdown sk ShutdownReceive
sClose sk
suck :: IO Bool -> IO ()
suck rd = do
f <- rd
if f then return () else suck rd
-- | Checks both that the underlying Socket is connected
-- and that the connection peer matches the given
-- host name (which is recorded locally).
isConnectedTo :: Connection -> EndPoint -> IO Bool
isConnectedTo (Connection conn) endPoint = isTCPConnectedTo conn endPoint
isTCPConnectedTo :: HandleStream ty -> EndPoint -> IO Bool
isTCPConnectedTo conn endPoint = do
v <- readMVar (getRef conn)
case v of
ConnClosed -> return False
_
| connEndPoint v == endPoint ->
catchIO (getPeerName (connSock v) >> return True) (const $ return False)
| otherwise -> return False
readBlockBS :: HStream a => HandleStream a -> Int -> IO (Result a)
readBlockBS ref n = onNonClosedDo ref $ \ conn -> do
x <- bufferGetBlock ref n
maybe (return ())
(\ h -> hook_readBlock h (buf_toStr $ connBuffer conn) n x)
(connHooks' conn)
return x
-- This function uses a buffer, at this time the buffer is just 1000 characters.
-- (however many bytes this is is left for the user to decipher)
readLineBS :: HStream a => HandleStream a -> IO (Result a)
readLineBS ref = onNonClosedDo ref $ \ conn -> do
x <- bufferReadLine ref
maybe (return ())
(\ h -> hook_readLine h (buf_toStr $ connBuffer conn) x)
(connHooks' conn)
return x
-- The 'Connection' object allows no outward buffering,
-- since in general messages are serialised in their entirety.
writeBlockBS :: HandleStream a -> a -> IO (Result ())
writeBlockBS ref b = onNonClosedDo ref $ \ conn -> do
x <- bufferPutBlock (connBuffer conn) (connHandle conn) b
maybe (return ())
(\ h -> hook_writeBlock h (buf_toStr $ connBuffer conn) b x)
(connHooks' conn)
return x
closeIt :: HStream ty => HandleStream ty -> (ty -> Bool) -> Bool -> IO ()
closeIt c p b = do
closeConnection c (if b
then readLineBS c >>= \ x -> case x of { Right xs -> return (p xs); _ -> return True}
else return True)
conn <- readMVar (getRef c)
maybe (return ())
(hook_close)
(connHooks' conn)
closeEOF :: HandleStream ty -> Bool -> IO ()
closeEOF c flg = modifyMVar_ (getRef c) (\ co -> return co{connCloseEOF=flg})
bufferGetBlock :: HStream a => HandleStream a -> Int -> IO (Result a)
bufferGetBlock ref n = onNonClosedDo ref $ \ conn -> do
case connInput conn of
Just c -> do
let (a,b) = buf_splitAt (connBuffer conn) n c
modifyMVar_ (getRef ref) (\ co -> return co{connInput=Just b})
return (return a)
_ -> do
catchIO (buf_hGet (connBuffer conn) (connHandle conn) n >>= return.return)
(\ e ->
if isEOFError e
then do
when (connCloseEOF conn) $ catchIO (closeQuick ref) (\ _ -> return ())
return (return (buf_empty (connBuffer conn)))
else return (failMisc (show e)))
bufferPutBlock :: BufferOp a -> Handle -> a -> IO (Result ())
bufferPutBlock ops h b =
catchIO (buf_hPut ops h b >> hFlush h >> return (return ()))
(\ e -> return (failMisc (show e)))
bufferReadLine :: HStream a => HandleStream a -> IO (Result a)
bufferReadLine ref = onNonClosedDo ref $ \ conn -> do
case connInput conn of
Just c -> do
let (a,b0) = buf_span (connBuffer conn) (/='\n') c
let (newl,b1) = buf_splitAt (connBuffer conn) 1 b0
modifyMVar_ (getRef ref) (\ co -> return co{connInput=Just b1})
return (return (buf_append (connBuffer conn) a newl))
_ -> catchIO
(buf_hGetLine (connBuffer conn) (connHandle conn) >>=
return . return . appendNL (connBuffer conn))
(\ e ->
if isEOFError e
then do
when (connCloseEOF conn) $ catchIO (closeQuick ref) (\ _ -> return ())
return (return (buf_empty (connBuffer conn)))
else return (failMisc (show e)))
where
-- yes, this s**ks.. _may_ have to be addressed if perf
-- suggests worthiness.
appendNL ops b = buf_snoc ops b nl
nl :: Word8
nl = fromIntegral (fromEnum '\n')
onNonClosedDo :: HandleStream a -> (Conn a -> IO (Result b)) -> IO (Result b)
onNonClosedDo h act = do
x <- readMVar (getRef h)
case x of
ConnClosed{} -> return (failWith ErrorClosed)
_ -> act x