-- | Internal socket utilities implementing missing
--   features of 'System.Socket' which are yet to be
--   upstreamed.
module Network.Gopher.Util.Socket
  ( gracefulClose
  ) where

import Control.Concurrent.MVar (withMVar)
import Control.Concurrent (threadDelay)
import Control.Concurrent.Async (race)
import Control.Exception.Base (throwIO)
import Control.Monad (void, when)
import Data.Functor ((<&>))
import Foreign.C.Error (Errno (..), getErrno)
import Foreign.C.Types (CInt (..))
import System.Socket (receive, msgNoSignal, SocketException (..), close, Family ())
import System.Socket.Type.Stream (Stream ())
import System.Socket.Protocol.TCP (TCP ())
import System.Socket.Unsafe (Socket (..))

-- Until https://github.com/lpeterse/haskell-socket/pull/67 gets
-- merged, we have to implement shutdown ourselves.
foreign import ccall unsafe "shutdown"
  c_shutdown :: CInt -> CInt -> IO CInt

data ShutdownHow
  -- | Disallow Reading (calls to 'receive' are empty).
  = ShutdownRead
  -- | Disallow Writing (calls to 'send' throw).
  | ShutdownWrite
  -- | Disallow both.
  | ShutdownReadWrite
  deriving (Int -> ShutdownHow -> ShowS
[ShutdownHow] -> ShowS
ShutdownHow -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ShutdownHow] -> ShowS
$cshowList :: [ShutdownHow] -> ShowS
show :: ShutdownHow -> String
$cshow :: ShutdownHow -> String
showsPrec :: Int -> ShutdownHow -> ShowS
$cshowsPrec :: Int -> ShutdownHow -> ShowS
Show, ShutdownHow -> ShutdownHow -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ShutdownHow -> ShutdownHow -> Bool
$c/= :: ShutdownHow -> ShutdownHow -> Bool
== :: ShutdownHow -> ShutdownHow -> Bool
$c== :: ShutdownHow -> ShutdownHow -> Bool
Eq, Eq ShutdownHow
ShutdownHow -> ShutdownHow -> Bool
ShutdownHow -> ShutdownHow -> Ordering
ShutdownHow -> ShutdownHow -> ShutdownHow
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ShutdownHow -> ShutdownHow -> ShutdownHow
$cmin :: ShutdownHow -> ShutdownHow -> ShutdownHow
max :: ShutdownHow -> ShutdownHow -> ShutdownHow
$cmax :: ShutdownHow -> ShutdownHow -> ShutdownHow
>= :: ShutdownHow -> ShutdownHow -> Bool
$c>= :: ShutdownHow -> ShutdownHow -> Bool
> :: ShutdownHow -> ShutdownHow -> Bool
$c> :: ShutdownHow -> ShutdownHow -> Bool
<= :: ShutdownHow -> ShutdownHow -> Bool
$c<= :: ShutdownHow -> ShutdownHow -> Bool
< :: ShutdownHow -> ShutdownHow -> Bool
$c< :: ShutdownHow -> ShutdownHow -> Bool
compare :: ShutdownHow -> ShutdownHow -> Ordering
$ccompare :: ShutdownHow -> ShutdownHow -> Ordering
Ord, Int -> ShutdownHow
ShutdownHow -> Int
ShutdownHow -> [ShutdownHow]
ShutdownHow -> ShutdownHow
ShutdownHow -> ShutdownHow -> [ShutdownHow]
ShutdownHow -> ShutdownHow -> ShutdownHow -> [ShutdownHow]
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
enumFromThenTo :: ShutdownHow -> ShutdownHow -> ShutdownHow -> [ShutdownHow]
$cenumFromThenTo :: ShutdownHow -> ShutdownHow -> ShutdownHow -> [ShutdownHow]
enumFromTo :: ShutdownHow -> ShutdownHow -> [ShutdownHow]
$cenumFromTo :: ShutdownHow -> ShutdownHow -> [ShutdownHow]
enumFromThen :: ShutdownHow -> ShutdownHow -> [ShutdownHow]
$cenumFromThen :: ShutdownHow -> ShutdownHow -> [ShutdownHow]
enumFrom :: ShutdownHow -> [ShutdownHow]
$cenumFrom :: ShutdownHow -> [ShutdownHow]
fromEnum :: ShutdownHow -> Int
$cfromEnum :: ShutdownHow -> Int
toEnum :: Int -> ShutdownHow
$ctoEnum :: Int -> ShutdownHow
pred :: ShutdownHow -> ShutdownHow
$cpred :: ShutdownHow -> ShutdownHow
succ :: ShutdownHow -> ShutdownHow
$csucc :: ShutdownHow -> ShutdownHow
Enum)

-- | Shutdown a stream connection (partially).
--   Will send TCP FIN and prompt a client to
--   close the connection.
--
--   Not exposed to prevent future name clash.
shutdown :: Socket a Stream TCP -> ShutdownHow -> IO ()
shutdown :: forall a. Socket a Stream TCP -> ShutdownHow -> IO ()
shutdown (Socket MVar Fd
mvar) ShutdownHow
how = forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar Fd
mvar forall a b. (a -> b) -> a -> b
$ \Fd
fd -> do
  CInt
res <- CInt -> CInt -> IO CInt
c_shutdown (forall a b. (Integral a, Num b) => a -> b
fromIntegral Fd
fd)
    forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall a. Enum a => a -> Int
fromEnum ShutdownHow
how
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (CInt
res forall a. Eq a => a -> a -> Bool
/= CInt
0) forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => e -> IO a
throwIO forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
    (IO Errno
getErrno forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \(Errno CInt
errno) -> CInt -> SocketException
SocketException CInt
errno)

-- | Shutdown connection and give client a bit
--   of time to clean up on its end before closing
--   the connection to avoid a broken pipe on the
--   other side.
gracefulClose :: Family f => Socket f Stream TCP -> IO ()
gracefulClose :: forall f. Family f => Socket f Stream TCP -> IO ()
gracefulClose Socket f Stream TCP
sock = do
  -- send TCP FIN
  forall a. Socket a Stream TCP -> ShutdownHow -> IO ()
shutdown Socket f Stream TCP
sock ShutdownHow
ShutdownWrite
  -- wait for some kind of read from the
  -- client (either mempty, meaning TCP FIN,
  -- something else which would mean protocol
  -- violation). Give up after 1s.
  Either () ()
_ <- forall a b. IO a -> IO b -> IO (Either a b)
race (forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall f t p. Socket f t p -> Int -> MessageFlags -> IO ByteString
receive Socket f Stream TCP
sock Int
16 MessageFlags
msgNoSignal) (Int -> IO ()
threadDelay Int
1000000)
  forall f t p. Socket f t p -> IO ()
close Socket f Stream TCP
sock