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