{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- | The Unix domain socket server backing the POSIX semaphore implementation.
--
-- The server manages a shared token pool and accepts multiple client
-- connections, each served on its own thread.
--
-- Protocol (SOCK_STREAM, one byte per command):
--
--   "-"  Wait (blocking acquire).  Decrements semaphore; replies ".".
--   "?"  Try-wait.  Decrements if positive and replies "."; otherwise replies "!".
--   "+"  Release.  Increments pool; replies ".".  Rejected with "!" if
--        myCount <= 0 (client has not acquired any tokens on this connection).
--
-- Unrecognised bytes are rejected with 'RspFail'.
--
-- Per-connection token tracking: the server counts tokens held by each
-- connection (myCount).  On disconnect (EOF / ResourceVanished) held
-- tokens are returned to the pool, so a crashing client cannot leak
-- tokens.
--
-- Connections are tied to a token, so each token (obtained by
-- @waitOnSemaphore@) has its own connection.
module System.Semaphore.Internal.Posix.Server
  ( serverLoop
  , pattern CmdWait, pattern CmdTryWait, pattern CmdRelease
  , pattern RspOk, pattern RspFail
  ) where

-- base
import Control.Concurrent
  ( ThreadId, forkIOWithUnmask, killThread )
import Control.Concurrent.MVar
  ( MVar, newEmptyMVar, newMVar, putMVar
  , readMVar, takeMVar, tryTakeMVar )
import Control.Exception
  ( IOException, SomeException
  , catch, finally, mask, mask_
  , onException, throw, try, uninterruptibleMask_
  )
import Control.Monad
  ( forM_, forever, void, when )
import Data.Word ( Word8 )
import Foreign.C.Error ( Errno(Errno), eCONNABORTED )
import GHC.IO.Exception ( ioe_errno )
import System.IO.Error ( isFullError )

-- stm
import Control.Concurrent.STM
  ( TVar, atomically, newTVarIO, readTVar, readTVarIO
  , modifyTVar', writeTVar, retry )

-- unix
import System.Posix.IO ( closeFd )
import System.Posix.Types ( Fd )

import System.Semaphore.Internal.DomainSocket
  ( pollAcceptSocket, AcceptResult(..)
  , fdReadByte, fdWriteByte
  , fdShutdown )

---------------------------------------
-- Protocol byte constants

pattern CmdWait, CmdTryWait, CmdRelease :: Word8
pattern $mCmdWait :: forall {r}. Word8 -> ((# #) -> r) -> ((# #) -> r) -> r
$bCmdWait :: Word8
CmdWait    = 0x2D -- '-'
pattern $mCmdTryWait :: forall {r}. Word8 -> ((# #) -> r) -> ((# #) -> r) -> r
$bCmdTryWait :: Word8
CmdTryWait = 0x3F -- '?'
pattern $mCmdRelease :: forall {r}. Word8 -> ((# #) -> r) -> ((# #) -> r) -> r
$bCmdRelease :: Word8
CmdRelease = 0x2B -- '+'

pattern RspOk, RspFail :: Word8
pattern $mRspOk :: forall {r}. Word8 -> ((# #) -> r) -> ((# #) -> r) -> r
$bRspOk :: Word8
RspOk   = 0x2E -- '.'
pattern $mRspFail :: forall {r}. Word8 -> ((# #) -> r) -> ((# #) -> r) -> r
$bRspFail :: Word8
RspFail = 0x21 -- '!'

-- | Children of the server, to be cleaned up
-- childFdLock is full when the 'Fd' is still known to be valid,
-- and subject to being shutdown or closed
data Child = Child
  { Child -> MVar ThreadId
childThread :: !(MVar ThreadId)
  , Child -> MVar Fd
childFdLock :: !(MVar Fd)
  }

-- | Run the server accept loop on a listening socket until cancellation is
-- signalled via the cancel pipe.  All accepted connections are served on
-- child threads, which are cleaned up before this returns.
serverLoop :: TVar Int -- ^ shared token pool
           -> Fd        -- ^ listening socket
           -> Fd        -- ^ read end of the cancel pipe
           -> IO ()
serverLoop :: TVar Int -> Fd -> Fd -> IO ()
serverLoop TVar Int
pool Fd
listenFd Fd
cancelFd = do
    children <- [Child] -> IO (TVar [Child])
forall a. a -> IO (TVar a)
newTVarIO ([] :: [Child])
    loop children `finally` killChildren children
  where
    loop :: TVar [Child] -> IO ()
loop TVar [Child]
children = do
      continueLoop <- IO Bool -> IO Bool
forall a. IO a -> IO a
mask_ (IO Bool -> IO Bool) -> IO Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
        r <- IO AcceptResult
acceptWithRetry
        case r of
          AcceptResult
AcceptCancelled     -> Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
          AcceptedFd Fd
clientFd -> do
            TVar [Child] -> Fd -> IO ()
forkServeChild TVar [Child]
children Fd
clientFd
            Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
      when continueLoop $ loop children

    -- pollAcceptSocket only returns if the accept succeeded, or cancellation
    -- was signalled via the cancel pipe.
    acceptWithRetry :: IO AcceptResult
    acceptWithRetry :: IO AcceptResult
acceptWithRetry = Fd -> Fd -> IO AcceptResult
pollAcceptSocket Fd
listenFd Fd
cancelFd IO AcceptResult
-> (IOException -> IO AcceptResult) -> IO AcceptResult
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` IOException -> IO AcceptResult
handleIOError

    -- Retry accept on transient errors.
    handleIOError :: IOException -> IO AcceptResult
    handleIOError :: IOException -> IO AcceptResult
handleIOError IOException
e
      -- 'isFullError' catches ResourceExhausted (EMFILE/ENFILE/ENOBUFS/ENOMEM and more that accept doesn't produce but are harmless to retry).
      | IOException -> Bool
isFullError IOException
e                             = IO AcceptResult
acceptWithRetry
      -- ECONNABORTED is also transient but categorised as OtherError, we additionaly match that.
      | Just CInt
err <- IOException -> Maybe CInt
ioe_errno IOException
e
      , CInt -> Errno
Errno CInt
err Errno -> Errno -> Bool
forall a. Eq a => a -> a -> Bool
== Errno
eCONNABORTED                 = IO AcceptResult
acceptWithRetry
      -- EINTR is absorbed in hs_poll_accept.
      -- everything else (EBADF, EINVAL, ENOTSOCK, ...) is rethrown.
      | Bool
otherwise                                 = IOException -> IO AcceptResult
forall a e. (HasCallStack, Exception e) => e -> a
throw IOException
e

    forkServeChild :: TVar [Child] -> Fd -> IO ()
forkServeChild TVar [Child]
children Fd
clientFd = do
      fdLock <- Fd -> IO (MVar Fd)
forall a. a -> IO (MVar a)
newMVar Fd
clientFd
      tidVar <- newEmptyMVar
      let child = MVar ThreadId -> MVar Fd -> Child
Child MVar ThreadId
tidVar MVar Fd
fdLock
      atomically $ modifyTVar' children (child :)
      childTid <- forkIOWithUnmask $ \forall a. IO a -> IO a
unmask ->
        (forall a. IO a -> IO a)
-> TVar Int -> TVar [Child] -> Fd -> Child -> IO ()
serve IO a -> IO a
forall a. IO a -> IO a
unmask TVar Int
pool TVar [Child]
children Fd
clientFd Child
child
      putMVar tidVar childTid

    -- Interrupt all children blocked on a read from the FD, then kill them.
    killChildren :: TVar [Child] -> IO ()
    killChildren :: TVar [Child] -> IO ()
killChildren TVar [Child]
children = do
      kids <- TVar [Child] -> IO [Child]
forall a. TVar a -> IO a
readTVarIO TVar [Child]
children
      forM_ kids $ \Child
child -> do
        -- If the child is in a read(), interrupt it
        -- by calling 'fdShutdown'.
        mb <- MVar Fd -> IO (Maybe Fd)
forall a. MVar a -> IO (Maybe a)
tryTakeMVar (Child -> MVar Fd
childFdLock Child
child)
        case mb of
          Just Fd
cfd -> do
            IO (Either IOException ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either IOException ()) -> IO ())
-> IO (Either IOException ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try @IOException (IO () -> IO (Either IOException ()))
-> IO () -> IO (Either IOException ())
forall a b. (a -> b) -> a -> b
$ Fd -> IO ()
fdShutdown Fd
cfd
            MVar Fd -> Fd -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (Child -> MVar Fd
childFdLock Child
child) Fd
cfd
          Maybe Fd
Nothing  -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()  -- serve thread already closing
      -- No children blocked on read: terminate them all.
      -- childThread was filled by the parent before mask exit; readMVar of
      -- a definitely-full MVar is non-interruptible.
      forM_ kids $ \Child
child -> do
        tid <- MVar ThreadId -> IO ThreadId
forall a. MVar a -> IO a
readMVar (Child -> MVar ThreadId
childThread Child
child)
        killThread tid

-- | Per-connection server loop.
serve :: (forall a. IO a -> IO a)
      -> TVar Int -> TVar [Child] -> Fd -> Child
      -> IO ()
serve :: (forall a. IO a -> IO a)
-> TVar Int -> TVar [Child] -> Fd -> Child -> IO ()
serve forall a. IO a -> IO a
restore TVar Int
pool TVar [Child]
children Fd
fd (Child MVar ThreadId
_ MVar Fd
fdLock) = do
    myCount <- Int -> IO (TVar Int)
forall a. a -> IO (TVar a)
newTVarIO (Int
0 :: Int)
    let loop = IO () -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO b) -> IO () -> IO b
forall a b. (a -> b) -> a -> b
$ ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO ()) -> IO ())
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restoreInner -> do
            -- fdReadByte is a safe-FFI read(2), interrupted by
            -- fdShutdown from killChildren, not by throwTo.
            msg <- HasCallStack => Fd -> IO Word8
Fd -> IO Word8
fdReadByte Fd
fd
            case msg of
              Word8
CmdWait -> do
                -- Block until a token is available.
                -- restoreInner keeps retry interruptible under mask.
                IO () -> IO ()
forall a. IO a -> IO a
restoreInner (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                    n <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
pool
                    when (n <= 0) retry
                    writeTVar pool (n - 1)
                    modifyTVar' myCount (+ 1)
                HasCallStack => Fd -> Word8 -> IO ()
Fd -> Word8 -> IO ()
fdWriteByte Fd
fd Word8
RspOk

              Word8
CmdRelease -> do
                ok <- STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
                    mc <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
myCount
                    if mc > 0
                      then do
                        modifyTVar' pool (+ 1)
                        modifyTVar' myCount (subtract 1)
                        return True
                      else return False
                fdWriteByte fd (if ok then RspOk else RspFail)

              Word8
CmdTryWait -> do
                acquired <- STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
                    n <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
pool
                    if n > 0
                      then do
                        writeTVar pool (n - 1)
                        modifyTVar' myCount (+ 1)
                        return True
                      else return False
                if acquired
                  then fdWriteByte fd RspOk
                  else fdWriteByte fd RspFail

              -- Unknown command: reply so the client doesn't hang in read().
              Word8
_ -> HasCallStack => Fd -> Word8 -> IO ()
Fd -> Word8 -> IO ()
fdWriteByte Fd
fd Word8
RspFail

        cleanup = do
          -- Return tokens to the pool and remove from children list.
          STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            n <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
myCount
            when (n > 0) $ modifyTVar' pool (+ n)
            modifyTVar' children (filter (\Child
c -> Child -> MVar Fd
childFdLock Child
c MVar Fd -> MVar Fd -> Bool
forall a. Eq a => a -> a -> Bool
/= MVar Fd
fdLock))
          -- Take fd ownership and close.
          -- prevents killChildren from double closing fd
          IO Fd -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Fd -> IO ()) -> IO Fd -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar Fd -> IO Fd
forall a. MVar a -> IO a
takeMVar MVar Fd
fdLock
          IO (Either IOException ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either IOException ()) -> IO ())
-> IO (Either IOException ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try @IOException (IO () -> IO (Either IOException ()))
-> IO () -> IO (Either IOException ())
forall a b. (a -> b) -> a -> b
$ Fd -> IO ()
closeFd Fd
fd

    -- restore so thread can be killed in between loop iterations
    -- Catch IOException (EOF/disconnect) silently.
    (restore loop `catch` \(IOException
_ :: IOException) -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
      `finally` cleanup