{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
module Network.Unexceptional
( accept_
, socket
, connect
, connectInterruptible
) where
import Control.Applicative ((<|>))
import Control.Concurrent.STM (STM, TVar)
import Control.Exception (mask_)
import Control.Monad ((<=<))
import Data.Functor (($>))
import Foreign.C.Error (Errno (Errno))
import Foreign.C.Error.Pattern (pattern EAGAIN, pattern EINPROGRESS, pattern EINTR, pattern EWOULDBLOCK)
import Foreign.Marshal.Alloc (allocaBytes)
import Foreign.Ptr (castPtr, nullPtr)
import GHC.Conc (threadWaitRead, threadWaitWrite, threadWaitWriteSTM)
import GHC.Exts (Ptr)
import Network.Socket (SockAddr, Socket, SocketOption (SoError), getSocketOption, mkSocket, withFdSocket)
import Network.Socket.Address (SocketAddress, pokeSocketAddress, sizeOfSocketAddress)
import System.Posix.Types (Fd (Fd))
import qualified Control.Concurrent.STM as STM
import qualified Linux.Socket as X
import qualified Network.Socket as N
import qualified Posix.Socket as X
{- | Accept a connection. See the documentation in @network@ for @accept@.
Note: This may leak a file descriptor if an asynchronous exception is
received while this function is running.
-}
accept_ :: Socket -> IO (Either Errno Socket)
accept_ listing_sock = withFdSocket listing_sock $ \listing_fd -> do
let acceptLoop = do
threadWaitRead (Fd listing_fd)
X.uninterruptibleAccept4_ (Fd listing_fd) (X.nonblocking <> X.closeOnExec) >>= \case
Left e ->
if e == EAGAIN || e == EWOULDBLOCK
then acceptLoop
else pure (Left e)
Right (Fd fd) -> fmap Right (mkSocket fd)
acceptLoop
{- | Connect to a socket address. See the documentation in @network@
for @connect@.
-}
connect :: Socket -> SockAddr -> IO (Either Errno ())
connect s sa = withSocketAddress sa $ \p_sa sz -> withFdSocket s $ \fd ->
let loop = do
r <- X.uninterruptibleConnectPtr (Fd fd) p_sa sz
case r of
Right _ -> pure (Right ())
Left err -> case err of
EINTR -> loop
EINPROGRESS -> do
threadWaitWrite (Fd fd)
errB <- getSocketOption s SoError
case errB of
0 -> pure (Right ())
_ -> pure (Left (Errno (fromIntegral errB)))
_ -> pure (Left err)
in loop
{- | Variant of 'connect' that can be interrupted by setting the interrupt
variable to @True@. If interrupted in this way, this function returns
@EAGAIN@. For example, to attempt to connect for no more than 1 second:
> interrupt <- Control.Concurrent.STM.registerDelay 1_000_000
> connectInterruptible interrupt sock sockAddr
-}
connectInterruptible :: TVar Bool -> Socket -> SockAddr -> IO (Either Errno ())
connectInterruptible !interrupt s sa = withSocketAddress sa $ \p_sa sz -> withFdSocket s $ \fd ->
let loop = do
r <- X.uninterruptibleConnectPtr (Fd fd) p_sa sz
case r of
Right _ -> pure (Right ())
Left err -> case err of
EINTR -> loop
EINPROGRESS ->
waitUntilWriteable interrupt (Fd fd) >>= \case
Interrupted -> pure (Left EAGAIN)
Ready -> do
errB <- getSocketOption s SoError
case errB of
0 -> pure (Right ())
_ -> pure (Left (Errno (fromIntegral errB)))
_ -> pure (Left err)
in loop
-- Copied this from the network library. TODO: See if network can
-- just export this.
withSocketAddress :: (SocketAddress sa) => sa -> (Ptr sa -> Int -> IO a) -> IO a
withSocketAddress addr f = do
let sz = sizeOfSocketAddress addr
if sz == 0
then f nullPtr 0
else allocaBytes sz $ \p -> pokeSocketAddress p addr >> f (castPtr p) sz
data Outcome = Ready | Interrupted
checkFinished :: TVar Bool -> STM ()
checkFinished = STM.check <=< STM.readTVar
waitUntilWriteable :: TVar Bool -> Fd -> IO Outcome
waitUntilWriteable !interrupt !fd = do
(isReadyAction, deregister) <- threadWaitWriteSTM fd
outcome <- STM.atomically $ (isReadyAction $> Ready) <|> (checkFinished interrupt $> Interrupted)
deregister
pure outcome
{- | Create a socket. See the documentation in @network@ for @socket@.
There is no interruptible variant of this function because it cannot
block. (It does not actually perform network activity.)
-}
socket ::
N.Family -> -- Family Name (usually AF_INET)
N.SocketType -> -- Socket Type (usually Stream)
N.ProtocolNumber -> -- Protocol Number (getProtocolByName to find value)
IO (Either Errno Socket) -- Unconnected Socket
socket !fam !stype !protocol = case stype of
N.Stream -> finish X.stream
N.Datagram -> finish X.datagram
_ -> fail "Network.Unexceptional.socket: Currently only supports stream and datagram types"
where
finish !sockTy = mask_ $ do
X.uninterruptibleSocket (X.Family (N.packFamily fam)) (X.applySocketFlags (X.closeOnExec <> X.nonblocking) sockTy) (X.Protocol protocol) >>= \case
Left err -> pure (Left err)
Right (Fd fd) -> do
s <- mkSocket fd
pure (Right s)