Permalink
Fetching contributors…
Cannot retrieve contributors at this time
415 lines (362 sloc) 15.7 KB
{-# 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