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