Add safe modules

Network.hs

@@ -317,8 +317,8 @@ accept (MkSocket _ family _ _ _) =
   error $ "Sorry, address family " ++ (show family) ++ " is not supported!"
 
 
--- | Close the socket. All future operations on the socket object will fail.
---   The remote end will receive no more data (after queued data is flushed).
+-- | Close the socket. Sending data to or receiving data from a closed socket
+--   causes undefined behaviour.
 sClose :: Socket -> IO ()
 sClose = close -- Explicit redefinition because Network.sClose is deperecated,
                -- hence the re-export would also be marked as such.

Network/Socket.hsc

@@ -561,6 +561,10 @@ foreign import ccall unsafe "free"
 --
 -- NOTE: blocking on Windows unless you compile with -threaded (see
 -- GHC ticket #1129)
+--
+-- Sending data to a closed socket causes undefined behaviour. To always get an
+-- exception, use the version in "Network.Socket.Safe".
+
 {-# WARNING sendTo "Use sendTo defined in \"Network.Socket.ByteString\"" #-}
 sendTo :: Socket        -- (possibly) bound/connected Socket
        -> String        -- Data to send
@@ -640,6 +644,10 @@ recvBufFrom sock@(MkSocket s family _stype _protocol _status) ptr nbytes
 -- | Send data to the socket. The socket must be connected to a remote
 -- socket. Returns the number of bytes sent.  Applications are
 -- responsible for ensuring that all data has been sent.
+--
+-- Sending data to a closed socket causes undefined behaviour. To always get an
+-- exception, use the version in "Network.Socket.Safe".
+
 {-# WARNING send "Use send defined in \"Network.Socket.ByteString\"" #-}
 send :: Socket  -- Bound/Connected Socket
      -> String  -- Data to send
@@ -650,6 +658,10 @@ send sock xs = withCStringLen xs $ \(str, len) ->
 -- | Send data to the socket. The socket must be connected to a remote
 -- socket. Returns the number of bytes sent.  Applications are
 -- responsible for ensuring that all data has been sent.
+--
+-- Sending data to a closed socket causes undefined behaviour. To always get an
+-- exception, use the version in "Network.Socket.Safe".
+
 sendBuf :: Socket     -- Bound/Connected Socket
         -> Ptr Word8  -- Pointer to the data to send
         -> Int        -- Length of the buffer
@@ -684,6 +696,10 @@ sendBuf sock@(MkSocket s _family _stype _protocol _status) str len = do
 --
 -- For TCP sockets, a zero length return value means the peer has
 -- closed its half side of the connection.
+--
+-- Receiving data from a closed socket causes undefined behaviour. To always get
+-- an exception, use the version in "Network.Socket.Safe".
+
 {-# WARNING recv "Use recv defined in \"Network.Socket.ByteString\"" #-}
 recv :: Socket -> Int -> IO String
 recv sock l = fst <$> recvLen sock l
@@ -707,6 +723,9 @@ recvLen sock nbytes =
 --
 -- For TCP sockets, a zero length return value means the peer has
 -- closed its half side of the connection.
+--
+-- Receiving data from a closed socket causes undefined behaviour. To always get
+-- an exception, use "Network.Socket.Safe".
 recvBuf :: Socket -> Ptr Word8 -> Int -> IO Int
 recvBuf sock@(MkSocket s _family _stype _protocol _status) ptr nbytes
  | nbytes <= 0 = ioError (mkInvalidRecvArgError "Network.Socket.recvBuf")
@@ -1072,9 +1091,8 @@ shutdown (MkSocket s _ _ _ _) stype = do
 
 -- -----------------------------------------------------------------------------
 
--- | Close the socket.  All future operations on the socket object
--- will fail.  The remote end will receive no more data (after queued
--- data is flushed).
+-- | Close the socket. Sending data to or receiving data from a closed socket
+-- causes undefined behaviour.
 close :: Socket -> IO ()
 close (MkSocket s _ _ _ socketStatus) = do
  modifyMVar_ socketStatus $ \ status ->

Network/Socket/ByteString.hsc

@@ -75,6 +75,10 @@ import Network.Socket.ByteString.MsgHdr (MsgHdr(..))
 -- | Send data to the socket.  The socket must be connected to a
 -- remote socket.  Returns the number of bytes sent. Applications are
 -- responsible for ensuring that all data has been sent.
+--
+-- Sending data to a closed socket causes undefined behaviour. To always get
+-- an exception, use the verion in "Network.Socket.ByteString.Safe".
+
 send :: Socket      -- ^ Connected socket
      -> ByteString  -- ^ Data to send
      -> IO Int      -- ^ Number of bytes sent
@@ -86,6 +90,10 @@ send sock xs = unsafeUseAsCStringLen xs $ \(str, len) ->
 -- until either all data has been sent or an error occurs.  On error,
 -- an exception is raised, and there is no way to determine how much
 -- data, if any, was successfully sent.
+--
+-- Sending data to a closed socket causes undefined behaviour. To always get an
+-- exception, use the verion in "Network.Socket.ByteString.Safe".
+
 sendAll :: Socket      -- ^ Connected socket
         -> ByteString  -- ^ Data to send
         -> IO ()
@@ -97,6 +105,10 @@ sendAll sock bs = do
 -- explicitly, so the socket need not be in a connected state.
 -- Returns the number of bytes sent. Applications are responsible for
 -- ensuring that all data has been sent.
+--
+-- Sending data to a closed socket causes undefined behaviour. To always get an
+-- exception, use the verion in "Network.Socket.ByteString.Safe".
+
 sendTo :: Socket      -- ^ Socket
        -> ByteString  -- ^ Data to send
        -> SockAddr    -- ^ Recipient address
@@ -110,6 +122,10 @@ sendTo sock xs addr =
 -- data has been sent or an error occurs.  On error, an exception is
 -- raised, and there is no way to determine how much data, if any, was
 -- successfully sent.
+--
+-- Sending data to a closed socket causes undefined behaviour. To always get an
+-- exception, use the verion in "Network.Socket.ByteString.Safe".
+
 sendAllTo :: Socket      -- ^ Socket
           -> ByteString  -- ^ Data to send
           -> SockAddr    -- ^ Recipient address
@@ -146,6 +162,10 @@ sendAllTo sock xs addr = do
 -- sent or an error occurs.  On error, an exception is raised, and
 -- there is no way to determine how much data, if any, was
 -- successfully sent.
+--
+-- Sending data to a closed socket causes undefined behaviour. To always get an
+-- exception, use the verion in "Network.Socket.ByteString.Safe".
+
 sendMany :: Socket        -- ^ Connected socket
          -> [ByteString]  -- ^ Data to send
          -> IO ()
@@ -169,6 +189,10 @@ sendMany sock = sendAll sock . B.concat
 -- continues to send data until either all data has been sent or an
 -- error occurs.  On error, an exception is raised, and there is no
 -- way to determine how much data, if any, was successfully sent.
+--
+-- Sending data to a closed socket causes undefined behaviour. To always get an
+-- exception, use the verion in "Network.Socket.ByteString.Safe".
+
 sendManyTo :: Socket        -- ^ Socket
            -> [ByteString]  -- ^ Data to send
            -> SockAddr      -- ^ Recipient address
@@ -205,6 +229,10 @@ sendManyTo sock cs = sendAllTo sock (B.concat cs)
 --
 -- For TCP sockets, a zero length return value means the peer has
 -- closed its half side of the connection.
+--
+-- Receiving data from a closed socket causes undefined behaviour. To
+-- always get an exception, use the verion in "Network.Socket.ByteString.Safe".
+
 recv :: Socket         -- ^ Connected socket
      -> Int            -- ^ Maximum number of bytes to receive
      -> IO ByteString  -- ^ Data received
@@ -217,6 +245,10 @@ recv sock nbytes
 -- connected state.  Returns @(bytes, address)@ where @bytes@ is a
 -- 'ByteString' representing the data received and @address@ is a
 -- 'SockAddr' representing the address of the sending socket.
+--
+-- Receiving data from a closed socket causes undefined behaviour. To always get
+-- an exception, use the verion in "Network.Socket.ByteString.Safe".
+
 recvFrom :: Socket                     -- ^ Socket
          -> Int                        -- ^ Maximum number of bytes to receive
          -> IO (ByteString, SockAddr)  -- ^ Data received and sender address

Network/Socket/ByteString/Lazy.hs

@@ -61,6 +61,9 @@ import Network.Socket.ByteString.Lazy.Posix (send, sendAll)
 -- more data to be received, the receiving side of the socket is shut
 -- down.  If there is an error and an exception is thrown, the socket
 -- is not shut down.
+--
+-- Receiving data from a closed socket causes undefined behaviour. To always get
+-- an exception, use the verion in "Network.Socket.ByteString.Lazy.Safe".
 getContents :: Socket         -- ^ Connected socket
             -> IO ByteString  -- ^ Data received
 getContents sock = loop where
@@ -77,6 +80,9 @@ getContents sock = loop where
 -- until a message arrives.
 --
 -- If there is no more data to be received, returns an empty 'ByteString'.
+--
+-- Receiving data from a closed socket causes undefined behaviour. To always get
+-- an exception, use the verion in "Network.Socket.ByteString.Lazy.Safe".
 recv :: Socket         -- ^ Connected socket
      -> Int64          -- ^ Maximum number of bytes to receive
      -> IO ByteString  -- ^ Data received

Network/Socket/ByteString/Lazy/Safe.hs

@@ -0,0 +1,42 @@
+-----------------------------------------------------------------------------
+-- |
+-- Module      :  Network.Socket.ByteString.Lazy.Safe
+-- Copyright   :  Echo Nolan 2016
+-- License     :  BSD-style (see the file libraries/network/LICENSE)
+--
+-- Maintainer  :  [email protected]
+-- Stability   :  provisional
+-- Portability :  portable
+--
+-- A drop in replacement for "Network.Socket.ByteString.Lazy" that sacrifices
+-- some performance for correctness. See "Network.Socket.Safe" for what exactly
+-- that means. See "Network.Socket.ByteString.Lazy" for API documentation.
+-----------------------------------------------------------------------------
+
+module Network.Socket.ByteString.Lazy.Safe
+  (
+    send
+  , sendAll
+  , getContents
+  , recv
+  ) where
+
+import qualified Network.Socket.ByteString.Lazy as Unsafe
+import Network.Socket.Internal
+import Network.Socket.Types
+
+import Prelude hiding (getContents)
+import Data.ByteString.Lazy (ByteString)
+import Data.Int (Int64)
+
+send :: Socket -> ByteString -> IO Int64
+send = wrapCheckStatus2 Unsafe.send "Network.Socket.ByteString.Lazy.Safe.send"
+
+sendAll :: Socket -> ByteString -> IO ()
+sendAll = wrapCheckStatus2 Unsafe.sendAll "Network.Socket.ByteString.Lazy.Safe.sendAll"
+
+getContents :: Socket -> IO ByteString
+getContents = wrapCheckStatus Unsafe.getContents "Network.Socket.ByteString.Lazy.Safe.getContents"
+
+recv :: Socket -> Int64 -> IO ByteString
+recv = wrapCheckStatus2 Unsafe.recv "Network.Socket.ByteString.Lazy.Safe.recv"

Network/Socket/ByteString/Safe.hs

@@ -0,0 +1,56 @@
+-----------------------------------------------------------------------------
+-- |
+-- Module      :  Network.Socket.ByteString.Safe
+-- Copyright   :  Echo Nolan 2016
+-- License     :  BSD-style (see the file libraries/network/LICENSE)
+--
+-- Maintainer  :  [email protected]
+-- Stability   :  provisional
+-- Portability :  portable
+--
+-- A drop in replacement for "Network.Socket.ByteString" that sacrifices some
+-- performance for correctness. See "Network.Socket.Safe" for what exactly that
+-- means. See "Network.Socket.ByteString" for API documentation.
+-----------------------------------------------------------------------------
+
+module Network.Socket.ByteString.Safe
+  (
+    send
+  , sendAll
+  , sendTo
+  , sendAllTo
+  , sendMany
+  , sendManyTo
+  , recv
+  , recvFrom
+  ) where
+
+import qualified Network.Socket.ByteString as Unsafe
+import Network.Socket.Internal
+import Network.Socket.Types
+
+import Data.ByteString (ByteString)
+
+send :: Socket -> ByteString -> IO Int
+send = wrapCheckStatus2 Unsafe.send "Network.Socket.ByteString.Safe.send"
+
+sendAll :: Socket -> ByteString -> IO ()
+sendAll = wrapCheckStatus2 Unsafe.sendAll "Network.Socket.ByteString.Safe.sendAll"
+
+sendTo :: Socket -> ByteString -> SockAddr -> IO Int
+sendTo = wrapCheckStatus3 Unsafe.sendTo "Network.Socket.ByteString.Safe.sendTo"
+
+sendAllTo :: Socket -> ByteString -> SockAddr -> IO ()
+sendAllTo = wrapCheckStatus3 Unsafe.sendAllTo "Network.Socket.ByteString.Safe.sendAllTo"
+
+sendMany :: Socket -> [ByteString] -> IO ()
+sendMany = wrapCheckStatus2 Unsafe.sendMany "Network.Socket.ByteString.Safe.sendMany"
+
+sendManyTo :: Socket -> [ByteString] -> SockAddr -> IO ()
+sendManyTo = wrapCheckStatus3 Unsafe.sendManyTo "Network.Socket.ByteString.Safe.sendManyTo"
+
+recv :: Socket -> Int -> IO ByteString
+recv = wrapCheckStatus2 Unsafe.recv "Network.Socket.ByteString.Safe.recv"
+
+recvFrom :: Socket -> Int -> IO (ByteString,SockAddr)
+recvFrom = wrapCheckStatus2 Unsafe.recvFrom "Network.Socket.ByteString.Safe.recvFrom"

Network/Socket/Internal.hsc

@@ -67,6 +67,12 @@ module Network.Socket.Internal
 
     -- * Low-level helpers
     , zeroMemory
+
+    -- * Helpers for Network.Socket.Safe
+    , wrapCheckStatus
+    , wrapCheckStatus2
+    , wrapCheckStatus3
+    , wrapCheckStatus4
     ) where
 
 import Foreign.C.Error (throwErrno, throwErrnoIfMinus1Retry,
@@ -92,6 +98,9 @@ import Foreign.C.Types ( CChar )
 import System.IO.Error ( ioeSetErrorString, mkIOError )
 #endif
 
+import Control.Concurrent (withMVar)
+import Control.Exception (throwIO)
+
 import Network.Socket.Types
 
 -- ---------------------------------------------------------------------
@@ -271,3 +280,26 @@ withSocketsInit = unsafePerformIO $ do
 foreign import ccall unsafe "initWinSock" initWinSock :: IO Int
 
 #endif
+
+wrapCheckStatus :: (Socket -> IO a) -> String -> Socket -> IO a
+wrapCheckStatus act fnName sock@(MkSocket _ _ _ _ statusVar) =
+  withMVar statusVar $ \status ->
+    case status of
+      Closed -> throwIO $ userError $
+        fnName ++ ": attempted to use a closed socket"
+      _ -> act sock
+
+wrapCheckStatus2 :: (Socket -> a -> IO b) ->
+                    String -> Socket -> a -> IO b
+wrapCheckStatus2 act fnName sock a =
+  wrapCheckStatus (\s -> act s a) fnName sock
+
+wrapCheckStatus3 :: (Socket -> a -> b -> IO c) ->
+                    String -> Socket -> a -> b -> IO c
+wrapCheckStatus3 act fnName sock a b =
+  wrapCheckStatus (\s -> act s a b) fnName sock
+
+wrapCheckStatus4 :: (Socket -> a -> b -> c -> IO d) ->
+                    String -> Socket -> a -> b -> c -> IO d
+wrapCheckStatus4 act fnName sock a b c =
+  wrapCheckStatus (\s -> act s a b c) fnName sock