Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Fetching contributors…

Cannot retrieve contributors at this time

422 lines (369 sloc) 16.002 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 = do
v <- readMVar (getRef conn)
case v of
ConnClosed -> print "aa" >> return False
_
| connEndPoint v == endPoint ->
catchIO (getPeerName (connSock v) >> return True) (const $ return False)
| otherwise -> return False
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
Jump to Line
Something went wrong with that request. Please try again.