Merge PR #417

Author: kazu-yamamoto

Network/Socket.hs

@@ -122,6 +122,7 @@ module Network.Socket
     -- ** Closing
     , close
     , close'
+    , gracefulClose
     , shutdown
     , ShutdownCmd(..)
 

Network/Socket/Buffer.hs renamed to Network/Socket/Buffer.hsc

@@ -1,14 +1,21 @@
 {-# LANGUAGE CPP #-}
 
-#include "HsNetDef.h"
+##include "HsNetDef.h"
+#if defined(mingw32_HOST_OS)
+#  include "windows.h"
+#endif
 
 module Network.Socket.Buffer (
     sendBufTo
   , sendBuf
   , recvBufFrom
   , recvBuf
+  , recvBufNoWait
   ) where
 
+#if !defined(mingw32_HOST_OS)
+import Foreign.C.Error (getErrno, eAGAIN, eWOULDBLOCK)
+#endif
 import Foreign.Marshal.Alloc (alloca)
 import GHC.IO.Exception (IOErrorType(InvalidArgument))
 import System.IO.Error (mkIOError, ioeSetErrorString, catchIOError)
@@ -129,6 +136,43 @@ recvBuf s ptr nbytes
 #endif
     return $ fromIntegral len
 
+-- | Receive data from the socket. This function returns immediately
+--   even if data is not available. In other words, IO manager is NOT
+--   involved. The length of data is returned if received.
+--   -1 is returned in the case of EAGAIN or EWOULDBLOCK.
+--   -2 is returned in other error cases.
+recvBufNoWait :: Socket -> Ptr Word8 -> Int -> IO Int
+recvBufNoWait s ptr nbytes = withFdSocket s $ \fd -> do
+#if defined(mingw32_HOST_OS)
+    alloca $ \ptr_bytes -> do
+      res <- c_ioctlsocket fd #{const FIONREAD} ptr_bytes
+      avail <- peek ptr_bytes
+      r <- if res == #{const NO_ERROR} && avail > 0 then
+               c_recv fd (castPtr ptr) (fromIntegral nbytes) 0{-flags-}
+           else if avail == 0 then
+               -- Socket would block, could also mean socket is closed but
+               -- can't distinguish
+               return (-1)
+           else do err <- c_WSAGetLastError
+                   if err == #{const WSAEWOULDBLOCK}
+                       || err == #{const WSAEINPROGRESS} then
+                       return (-1)
+                     else
+                        return (-2)
+      return $ fromIntegral r
+
+#else
+    r <- c_recv fd (castPtr ptr) (fromIntegral nbytes) 0{-flags-}
+    if r >= 0 then
+        return $ fromIntegral r
+      else do
+        err <- getErrno
+        if err == eAGAIN || err == eWOULDBLOCK then
+            return (-1)
+          else
+            return (-2)
+#endif
+
 mkInvalidRecvArgError :: String -> IOError
 mkInvalidRecvArgError loc = ioeSetErrorString (mkIOError
                                     InvalidArgument
@@ -137,9 +181,14 @@ mkInvalidRecvArgError loc = ioeSetErrorString (mkIOError
 #if !defined(mingw32_HOST_OS)
 foreign import ccall unsafe "send"
   c_send :: CInt -> Ptr a -> CSize -> CInt -> IO CInt
+#else
+foreign import CALLCONV SAFE_ON_WIN "ioctlsocket"
+  c_ioctlsocket :: CInt -> CLong -> Ptr CULong -> IO CInt
+foreign import CALLCONV SAFE_ON_WIN "WSAGetLastError"
+  c_WSAGetLastError :: IO CInt
+#endif
 foreign import ccall unsafe "recv"
   c_recv :: CInt -> Ptr CChar -> CSize -> CInt -> IO CInt
-#endif
 foreign import CALLCONV SAFE_ON_WIN "sendto"
   c_sendto :: CInt -> Ptr a -> CSize -> CInt -> Ptr sa -> CInt -> IO CInt
 foreign import CALLCONV SAFE_ON_WIN "recvfrom"

Network/Socket/Shutdown.hs

@@ -6,8 +6,21 @@
 module Network.Socket.Shutdown (
     ShutdownCmd(..)
   , shutdown
+  , gracefulClose
   ) where
 
+import qualified Control.Exception as E
+import Foreign.Marshal.Alloc (mallocBytes, free)
+
+#if defined(mingw32_HOST_OS)
+import Control.Concurrent (threadDelay)
+#else
+import Control.Concurrent (putMVar, takeMVar, newEmptyMVar)
+import qualified GHC.Event as Ev
+import System.Posix.Types (Fd(..))
+#endif
+
+import Network.Socket.Buffer
 import Network.Socket.Imports
 import Network.Socket.Internal
 import Network.Socket.Types
@@ -34,3 +47,75 @@ shutdown s stype = void $ withFdSocket s $ \fd ->
 
 foreign import CALLCONV unsafe "shutdown"
   c_shutdown :: CInt -> CInt -> IO CInt
+
+#if !defined(mingw32_HOST_OS)
+data Wait = MoreData | TimeoutTripped
+#endif
+
+-- | Closing a socket gracefully.
+--   This sends TCP FIN and check if TCP FIN is received from the peer.
+--   The second argument is time out to receive TCP FIN in millisecond.
+--   In both normal cases and error cases, socket is deallocated finally.
+--
+--   Since: 3.1.1.0
+gracefulClose :: Socket -> Int -> IO ()
+gracefulClose s tmout = sendRecvFIN `E.finally` close s
+  where
+    sendRecvFIN = do
+        -- Sending TCP FIN.
+        shutdown s ShutdownSend
+        -- Waiting TCP FIN.
+        recvEOF
+#if defined(mingw32_HOST_OS)
+    -- milliseconds. Taken from BSD fast clock value.
+    clock = 200
+    recvEOF = E.bracket (mallocBytes bufSize) free $ loop 0
+      where
+        loop delay buf = do
+            -- We don't check the (positive) length.
+            -- In normal case, it's 0. That is, only FIN is received.
+            -- In error cases, data is available. But there is no
+            -- application which can read it. So, let's stop receiving
+            -- to prevent attacks.
+            r <- recvBufNoWait s buf bufSize
+            let delay' = delay + clock
+            when (r == -1 && delay' < tmout) $ do
+                threadDelay (clock * 1000)
+                loop delay' buf
+#else
+    recvEOF = do
+        Just evmgr <- Ev.getSystemEventManager
+        tmmgr <- Ev.getSystemTimerManager
+        mvar <- newEmptyMVar
+        E.bracket (register evmgr tmmgr mvar) (unregister evmgr tmmgr) $ \_ -> do
+            wait <- takeMVar mvar
+            case wait of
+              TimeoutTripped -> return ()
+              -- We don't check the (positive) length.
+              -- In normal case, it's 0. That is, only FIN is received.
+              -- In error cases, data is available. But there is no
+              -- application which can read it. So, let's stop receiving
+              -- to prevent attacks.
+              MoreData       -> E.bracket (mallocBytes bufSize)
+                                          free
+                                          (\buf -> void $ recvBufNoWait s buf bufSize)
+    register evmgr tmmgr mvar = do
+        -- millisecond to microsecond
+        key1 <- Ev.registerTimeout tmmgr (tmout * 1000) $
+            putMVar mvar TimeoutTripped
+        key2 <- withFdSocket s $ \fd' -> do
+            let callback _ _ = putMVar mvar MoreData
+                fd = Fd fd'
+#if __GLASGOW_HASKELL__ < 709
+            Ev.registerFd evmgr callback fd Ev.evtRead
+#else
+            Ev.registerFd evmgr callback fd Ev.evtRead Ev.OneShot
+#endif
+        return (key1, key2)
+    unregister evmgr tmmgr (key1,key2) = do
+        Ev.unregisterTimeout tmmgr key1
+        Ev.unregisterFd evmgr key2
+#endif
+    -- Don't use 4092 here. The GHC runtime takes the global lock
+    -- if the length is over 3276 bytes in 32bit or 3272 bytes in 64bit.
+    bufSize = 1024

configure.ac

@@ -1,5 +1,5 @@
 AC_INIT([Haskell network package],
-        [3.1.0.1],
+        [3.1.1.0],
         [[email protected]],
         [network])
 

network.cabal

@@ -1,6 +1,6 @@
 cabal-version:  1.18
 name:           network
-version:        3.1.0.1
+version:        3.1.1.0
 license:        BSD3
 license-file:   LICENSE
 maintainer:     Kazu Yamamoto, Evan Borden

tests/Network/SocketSpec.hs

@@ -196,3 +196,16 @@ spec = do
                     cred1 <- getPeerCredential s
                     cred1 `shouldBe` (Nothing,Nothing,Nothing)
             -}
+
+    describe "gracefulClose" $ do
+        it "does not send TCP RST back" $ do
+            let server sock = do
+                    void $ recv sock 1024 -- receiving "GOAWAY"
+                    gracefulClose sock 3000
+                client sock = do
+                    sendAll sock "GOAWAY"
+                    threadDelay 10000
+                    sendAll sock "PING"
+                    threadDelay 10000
+                    void $ recv sock 1024
+            tcpTest client server