{-# LANGUAGE CPP #-}

#include "HsNetDef.h"

module Network.Socket.Shutdown (
    ShutdownCmd(..)
  , shutdown
  , gracefulClose
  ) where

import qualified Control.Exception as E
import Foreign.Marshal.Alloc (mallocBytes, free)

import Control.Concurrent (threadDelay, yield)

import Network.Socket.Buffer
import Network.Socket.Imports
import Network.Socket.Internal
import Network.Socket.Types

data ShutdownCmd = ShutdownReceive
                 | ShutdownSend
                 | ShutdownBoth

sdownCmdToInt :: ShutdownCmd -> CInt
sdownCmdToInt :: ShutdownCmd -> CInt
sdownCmdToInt ShutdownCmd
ShutdownReceive = CInt
0
sdownCmdToInt ShutdownCmd
ShutdownSend    = CInt
1
sdownCmdToInt ShutdownCmd
ShutdownBoth    = CInt
2

-- | Shut down one or both halves of the connection, depending on the
-- second argument to the function.  If the second argument is
-- 'ShutdownReceive', further receives are disallowed.  If it is
-- 'ShutdownSend', further sends are disallowed.  If it is
-- 'ShutdownBoth', further sends and receives are disallowed.
shutdown :: Socket -> ShutdownCmd -> IO ()
shutdown :: Socket -> ShutdownCmd -> IO ()
shutdown Socket
s ShutdownCmd
stype = IO () -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Socket -> (CInt -> IO ()) -> IO ()
forall r. Socket -> (CInt -> IO r) -> IO r
withFdSocket Socket
s ((CInt -> IO ()) -> IO ()) -> (CInt -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \CInt
fd ->
  String -> IO CInt -> IO ()
forall a. (Eq a, Num a) => String -> IO a -> IO ()
throwSocketErrorIfMinus1Retry_ String
"Network.Socket.shutdown" (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$
    CInt -> CInt -> IO CInt
c_shutdown CInt
fd (CInt -> IO CInt) -> CInt -> IO CInt
forall a b. (a -> b) -> a -> b
$ ShutdownCmd -> CInt
sdownCmdToInt ShutdownCmd
stype

foreign import CALLCONV unsafe "shutdown"
  c_shutdown :: CInt -> CInt -> IO CInt

-- | 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 :: Socket -> Int -> IO ()
gracefulClose Socket
s Int
tmout0 = IO ()
sendRecvFIN IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`E.finally` Socket -> IO ()
close Socket
s
  where
    sendRecvFIN :: IO ()
sendRecvFIN = do
        -- Sending TCP FIN.
        Either SomeException ()
ex <- IO () -> IO (Either SomeException ())
forall e a. Exception e => IO a -> IO (Either e a)
E.try (IO () -> IO (Either SomeException ()))
-> IO () -> IO (Either SomeException ())
forall a b. (a -> b) -> a -> b
$ Socket -> ShutdownCmd -> IO ()
shutdown Socket
s ShutdownCmd
ShutdownSend
        case Either SomeException ()
ex of
          Left (E.SomeException e
_) -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
          Right () -> do
              -- Giving CPU time to other threads hoping that
              -- FIN arrives meanwhile.
              IO ()
yield
              -- Waiting TCP FIN.
              IO (Ptr Word8)
-> (Ptr Word8 -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracket (Int -> IO (Ptr Word8)
forall a. Int -> IO (Ptr a)
mallocBytes Int
bufSize) Ptr Word8 -> IO ()
forall a. Ptr a -> IO ()
free Ptr Word8 -> IO ()
recvEOFloop
    recvEOFloop :: Ptr Word8 -> IO ()
recvEOFloop Ptr Word8
buf = Int -> Int -> IO ()
loop Int
1 Int
0
      where
        loop :: Int -> Int -> IO ()
loop Int
delay Int
tmout = 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.
            Int
r <- Socket -> Ptr Word8 -> Int -> IO Int
recvBufNoWait Socket
s Ptr Word8
buf Int
bufSize
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1 Bool -> Bool -> Bool
&& Int
tmout Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
tmout0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                Int -> IO ()
threadDelay (Int
delay Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000)
                Int -> Int -> IO ()
loop (Int
delay Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2) (Int
tmout Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
delay)
    -- 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 :: Int
bufSize = Int
1024