diff --git a/Network/Connection.hs b/Network/Connection.hs --- a/Network/Connection.hs +++ b/Network/Connection.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DeriveDataTypeable #-} -- | -- Module : Network.Connection -- License : BSD-style @@ -16,6 +18,9 @@ , TLSSettings(..) , SockSettings(..) + -- * Exceptions + , LineTooLong(..) + -- * Library initialization , initConnectionContext , ConnectionContext @@ -29,6 +34,7 @@ , connectionGet , connectionGetChunk , connectionGetChunk' + , connectionGetLine , connectionPut -- * TLS related operation @@ -37,8 +43,8 @@ ) where import Control.Applicative -import Control.Monad ((>=>), when) import Control.Concurrent.MVar +import Control.Monad (join) import qualified Control.Exception as E import qualified System.IO.Error as E @@ -50,6 +56,7 @@ import Network.Socks5 import qualified Network as N +import Data.Data import Data.ByteString (ByteString) import qualified Data.ByteString as B import qualified Data.ByteString.Lazy as L @@ -64,6 +71,12 @@ type Manager = MVar (M.Map TLS.SessionID TLS.SessionData) data ConnectionSessionManager = ConnectionSessionManager Manager +-- | This is the exception raised if we reached the user specified limit for +-- the line in ConnectionGetLine. +data LineTooLong = LineTooLong deriving (Show,Typeable) + +instance E.Exception LineTooLong + instance TLS.SessionManager ConnectionSessionManager where sessionResume (ConnectionSessionManager mvar) sessionID = withMVar mvar (return . M.lookup sessionID) @@ -92,17 +105,11 @@ withBackend :: (ConnectionBackend -> IO a) -> Connection -> IO a withBackend f conn = modifyMVar (connectionBackend conn) (\b -> f b >>= \a -> return (b,a)) -withBuffer :: (ByteString -> IO (ByteString, b)) -> Connection -> IO b -withBuffer f conn = modifyMVar (connectionBuffer conn) f - -setEOF :: Connection -> IO () -setEOF conn = modifyMVar_ (connectionEOF conn) (const (return True)) - -checkEOF :: Connection -> IO () -checkEOF conn = readMVar (connectionEOF conn) >>= flip when (E.ioError $ E.mkIOError E.eofErrorType "" Nothing Nothing) - connectionNew :: ConnectionParams -> ConnectionBackend -> IO Connection -connectionNew p backend = Connection <$> newMVar backend <*> newMVar B.empty <*> pure (connectionHostname p, connectionPort p) <*> newMVar False +connectionNew p backend = + Connection <$> newMVar backend + <*> newMVar (Just B.empty) + <*> pure (connectionHostname p, connectionPort p) -- | Use an already established handle to create a connection object. -- @@ -121,7 +128,7 @@ -> ConnectionParams -- ^ The parameters for this connection (where to connect, and such). -> IO Connection -- ^ The new established connection on success. connectTo cg cParams = do - h <- conFct (connectionHostname cParams) (N.PortNumber $ connectionPort cParams) + h <- conFct (connectionHostname cParams) (N.PortNumber $ connectionPort cParams) connectFromHandle cg h cParams where conFct = case connectionUseSocks cParams of @@ -130,42 +137,109 @@ -- | Put a block of data in the connection. connectionPut :: Connection -> ByteString -> IO () -connectionPut connection content = checkEOF connection >> withBackend doWrite connection +connectionPut connection content = withBackend doWrite connection where doWrite (ConnectionStream h) = B.hPut h content >> hFlush h doWrite (ConnectionTLS ctx) = TLS.sendData ctx $ L.fromChunks [content] -- | Get some bytes from a connection. -- --- The size argument is just the maximum that could be returned to the user, --- however the call will returns as soon as there's data, even if there's less --- data than expected. +-- The size argument is just the maximum that could be returned to the user. +-- The call will return as soon as there's data, even if there's less +-- than requested. Hence, it behaves like 'B.hGetSome'. +-- +-- On end of input, 'connectionGet' returns 0, but subsequent calls will throw +-- an 'E.isEOFError' exception. connectionGet :: Connection -> Int -> IO ByteString -connectionGet conn size = connectionGetChunk' conn $ B.splitAt size +connectionGet conn size + | size < 0 = fail "Network.Connection.connectionGet: size < 0" + | size == 0 = return B.empty + | otherwise = connectionGetChunkBase "connectionGet" conn $ B.splitAt size -- | Get the next block of data from the connection. connectionGetChunk :: Connection -> IO ByteString -connectionGetChunk conn = connectionGetChunk' conn $ \s -> (s, B.empty) +connectionGetChunk conn = + connectionGetChunkBase "connectionGetChunk" conn $ \s -> (s, B.empty) -- | Like 'connectionGetChunk', but return the unused portion to the buffer, -- where it will be the next chunk read. connectionGetChunk' :: Connection -> (ByteString -> (a, ByteString)) -> IO a -connectionGetChunk' conn f = checkEOF conn >> withBuffer getData conn - where getData buf - | B.null buf = do - chunk <- withBackend (getMoreData >=> markEOF) conn - return $ swap $ f chunk - | otherwise = - return $ swap $ f buf +connectionGetChunk' = connectionGetChunkBase "connectionGetChunk'" - markEOF bs - | B.null bs = setEOF conn >> return bs - | otherwise = return bs +connectionGetChunkBase :: String -> Connection -> (ByteString -> (a, ByteString)) -> IO a +connectionGetChunkBase loc conn f = + modifyMVar (connectionBuffer conn) $ \m -> + case m of + Nothing -> throwEOF conn loc + Just buf + | B.null buf -> do + chunk <- withBackend getMoreData conn + if B.null chunk + then closeBuf chunk + else updateBuf chunk + | otherwise -> + updateBuf buf + where + getMoreData (ConnectionTLS tlsctx) = TLS.recvData tlsctx + getMoreData (ConnectionStream h) = B.hGetSome h (16 * 1024) - getMoreData (ConnectionTLS tlsctx) = TLS.recvData tlsctx - getMoreData (ConnectionStream h) = B.hGetSome h (16 * 1024) + updateBuf buf = case f buf of (a, !buf') -> return (Just buf', a) + closeBuf buf = case f buf of (a, _buf') -> return (Nothing, a) - swap (a, b) = (b, a) +-- | Get the next line, using ASCII LF as the line terminator. +-- +-- This throws an 'isEOFError' exception on end of input, and LineTooLong when +-- the number of bytes gathered is over the limit without a line terminator. +-- +-- The actual line returned can be bigger than the limit specified, provided +-- that the last chunk returned by the underlaying backend contains a LF. +-- In another world only when we need more input and limit is reached that the +-- LineTooLong exception will be raised. +-- +-- An end of file will be considered as a line terminator too, if line is +-- not empty. +connectionGetLine :: Int -- ^ Maximum number of bytes before raising a LineTooLong exception + -> Connection -- ^ Connection + -> IO ByteString -- ^ The received line with the LF trimmed +connectionGetLine limit conn = more (throwEOF conn loc) 0 id + where + loc = "connectionGetLine" + lineTooLong = E.throwIO LineTooLong + -- Accumulate chunks using a difference list, and concatenate them + -- when an end-of-line indicator is reached. + more eofK !currentSz !dl = + getChunk (\s -> let len = B.length s + in if currentSz + len > limit + then lineTooLong + else more eofK (currentSz + len) (dl . (s:))) + (\s -> done (dl . (s:))) + (done dl) + + done :: ([ByteString] -> [ByteString]) -> IO ByteString + done dl = return $! B.concat $ dl [] + + -- Get another chunk, and call one of the continuations + getChunk :: (ByteString -> IO r) -- moreK: need more input + -> (ByteString -> IO r) -- doneK: end of line (line terminator found) + -> IO r -- eofK: end of file + -> IO r + getChunk moreK doneK eofK = + join $ connectionGetChunkBase loc conn $ \s -> + if B.null s + then (eofK, B.empty) + else case B.breakByte 10 s of + (a, b) + | B.null b -> (moreK a, B.empty) + | otherwise -> (doneK a, B.tail b) + +throwEOF :: Connection -> String -> IO a +throwEOF conn loc = + E.throwIO $ E.mkIOError E.eofErrorType loc' Nothing (Just path) + where + loc' = "Network.Connection." ++ loc + path = let (host, port) = connectionID conn + in host ++ ":" ++ show port + -- | Close a connection. connectionClose :: Connection -> IO () connectionClose = withBackend backendClose @@ -173,12 +247,12 @@ backendClose (ConnectionStream h) = hClose h -- | Activate secure layer using the parameters specified. --- +-- -- This is typically used to negociate a TLS channel on an already -- establish channel, e.g. supporting a STARTTLS command. it also -- flush the received buffer to prevent application confusing -- received data before and after the setSecure call. --- +-- -- If the connection is already using TLS, nothing else happens. connectionSetSecure :: ConnectionContext -> Connection @@ -189,7 +263,7 @@ modifyMVar (connectionBackend connection) $ \backend -> case backend of (ConnectionStream h) -> do ctx <- tlsEstablish h (makeTLSParams cg params) - return (ConnectionTLS ctx, B.empty) + return (ConnectionTLS ctx, Just B.empty) (ConnectionTLS _) -> return (backend, b) -- | Returns if the connection is establish securely or not. diff --git a/Network/Connection/Types.hs b/Network/Connection/Types.hs --- a/Network/Connection/Types.hs +++ b/Network/Connection/Types.hs @@ -47,7 +47,7 @@ -- The simple settings is just the hostname and portnumber of the proxy server. -- -- That's for now the only settings in the SOCKS package, --- socks password, or authentication is not yet implemented. +-- socks password, or any sort of other authentications is not yet implemented. data SockSettings = SockSettingsSimple HostName PortNumber -- | TLS Settings that can be either expressed as simple settings, @@ -74,9 +74,8 @@ -- | This opaque type represent a connection to a destination. data Connection = Connection { connectionBackend :: MVar ConnectionBackend - , connectionBuffer :: MVar ByteString + , connectionBuffer :: MVar (Maybe ByteString) -- ^ this is set to 'Nothing' on EOF , connectionID :: (HostName, PortNumber) -- ^ return a simple tuple of the port and hostname that we're connected to. - , connectionEOF :: MVar Bool } -- | Shared values (certificate store, sessions, ..) between connections diff --git a/connection.cabal b/connection.cabal --- a/connection.cabal +++ b/connection.cabal @@ -1,5 +1,5 @@ Name: connection -Version: 0.1.2 +Version: 0.1.3 Description: Simple network library for all your connection need. .