{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables #-}

#include "HsNetDef.h"

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

import Control.Concurrent (rtsSupportsBoundThreads, yield)
import Control.Concurrent.STM
import qualified Control.Exception as E
import Foreign.Marshal.Alloc (mallocBytes, free)
import System.Timeout

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

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 IOException ()
ex <- IO () -> IO (Either IOException ())
forall e a. Exception e => IO a -> IO (Either e a)
E.try (IO () -> IO (Either IOException ()))
-> IO () -> IO (Either IOException ())
forall a b. (a -> b) -> a -> b
$ Socket -> ShutdownCmd -> IO ()
shutdown Socket
s ShutdownCmd
ShutdownSend
        case Either IOException ()
ex of
          -- Don't catch asynchronous exceptions
          Left (IOException
_ :: E.IOException) -> () -> 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 (Socket -> Int -> Ptr Word8 -> IO ()
recvEOF Socket
s Int
tmout0)

-- Waiting for the peer's FIN, draining what arrives before it.  The
-- receive is given up on rather than interrupted: an asynchronous
-- exception does not reach a thread blocked in one on Windows.
recvEOF :: Socket -> Int -> Ptr Word8 -> IO ()
recvEOF :: Socket -> Int -> Ptr Word8 -> IO ()
recvEOF Socket
s Int
tmout0 Ptr Word8
buf
    | Bool -> Bool
not Bool
rtsSupportsBoundThreads = Socket -> Int -> Ptr Word8 -> IO ()
recvEOFtimeout Socket
s Int
tmout0 Ptr Word8
buf
    | Bool
otherwise = do
        -- millisecond to microsecond
        TVar Bool
late <- Int -> IO (TVar Bool)
registerDelay (Int -> IO (TVar Bool)) -> Int -> IO (TVar Bool)
forall a b. (a -> b) -> a -> b
$ Int
tmout0 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000
        TVar Bool -> Int -> IO ()
loop TVar Bool
late Int
0
  where
    loop :: TVar Bool -> Int -> IO ()
loop TVar Bool
late Int
n0 = do
        Either IOException (Either () Int)
r <- IO (Either () Int) -> IO (Either IOException (Either () Int))
forall e a. Exception e => IO a -> IO (Either e a)
E.try (IO (Either () Int) -> IO (Either IOException (Either () Int)))
-> IO (Either () Int) -> IO (Either IOException (Either () Int))
forall a b. (a -> b) -> a -> b
$ Socket -> Ptr Word8 -> Int -> STM () -> IO (Either () Int)
forall a. Socket -> Ptr Word8 -> Int -> STM a -> IO (Either a Int)
recvBufWith Socket
s Ptr Word8
buf Int
bufSize (STM () -> IO (Either () Int)) -> STM () -> IO (Either () Int)
forall a b. (a -> b) -> a -> b
$ Bool -> STM ()
check (Bool -> STM ()) -> STM Bool -> STM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
late
        case Either IOException (Either () Int)
r of
            Left (IOException
_ :: E.IOException) -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            -- The deadline.
            Right (Left ()) -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            Right (Right Int
n1) -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                let n :: Int
n = Int
n0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n1
                Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
drainLimit) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Int -> IO ()
loop TVar Bool
late Int
n

-- registerDelay needs -threaded, and so does the give-up on Windows --
-- where this cannot cut a receive short either, but nothing else works
-- there without it.
recvEOFtimeout :: Socket -> Int -> Ptr Word8 -> IO ()
-- millisecond to microsecond
recvEOFtimeout :: Socket -> Int -> Ptr Word8 -> IO ()
recvEOFtimeout Socket
s Int
tmout0 Ptr Word8
buf = IO (Maybe Int) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe Int) -> IO ()) -> IO (Maybe Int) -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> IO Int -> IO (Maybe Int)
forall a. Int -> IO a -> IO (Maybe a)
timeout (Int
tmout0 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000) (IO Int -> IO (Maybe Int)) -> IO Int -> IO (Maybe Int)
forall a b. (a -> b) -> a -> b
$ Socket -> Ptr Word8 -> Int -> IO Int
recvBuf Socket
s Ptr Word8
buf Int
bufSize

-- 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
bufSize = Int
1024

-- Data left unread makes close send RST rather than FIN, so what has
-- arrived is drained -- but only this far: nothing is going to read it.
drainLimit :: Int
drainLimit :: Int
drainLimit = Int
128 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024