Network.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE CApiFFI #-}
{- | This module exports everything you need to engage with network-light. The
API is, in many places, similar to that of the standard Network package. However, this
package is smaller and works with both MicroHs and GHC. It is not intended to replace
or improve upon network, but is rather deliberately kept simple.

It exports a basic API to create and use sockets, and some helper functions that
send and receive ByteString.

The data types modeling the addressing families, socket options etc, are not complete.
They mirror what we've needed for our purposes, but please make a fork and open a PR to
add what you need, and I will gladly merge it.

At some point the hope is that MicroHs will be able to compile all of network, but until
then, using this library is the quickest work-around.
-}
module System.Network
    ( -- * Data types
      {- | @Socket@s are created by @socket@ or @accept@. When you get them, they are already configured to
      be non-blocking.-}
      Socket
    , Domain(..)
    , StreamType(..)
    , SockOpt(..)
    , SockAddr
    , mkSockAddr

      -- * Basic operations
    , socket
    , setsocketopt
    , close
    , connect
    , connect'
    , bind
    , accept
    , listen
    -- * Sending data
    , sendBuf
    , sendBufFull
    , sendString
    , sendByteString
    -- * Receiving data
    , recvBuf
    , recvString
    , recvByteString
    , recvByteStringFull
    ) where

import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BS
import Data.Word

#ifdef __MHS__
import System.IO.FD (waitForReadFD, waitForWriteFD)
#endif

import Foreign.C.Error
import Foreign.C.String
import Foreign.C.Types
import Foreign.Marshal.Alloc
import Foreign.Marshal.Utils
import Foreign.Ptr
import Foreign.Storable

import System.Network.Types

-- ---------------------------------------------------------------------------
-- FFI
-- ---------------------------------------------------------------------------

type FD = CInt

foreign import ccall      "sys/socket.h socket"     c_socket     :: CInt -> CInt -> CInt -> IO FD
-- connect/accept/send/recv can block for an unbounded time (waiting on
-- a peer), so they're "safe": under GHC with -threaded, a safe call
-- runs on its own OS thread instead of blocking the whole capability,
-- so other Haskell threads (e.g. forkIO'd connection handlers) keep
-- running. A single capability is enough -- no +RTS -N2 needed. mhs
-- accepts the same "safe" keyword but doesn't need it: its own
-- concurrency comes from non-blocking sockets plus waitForReadFD /
-- waitForWriteFD, independent of this annotation.
foreign import ccall safe "sys/socket.h connect"    c_connect    :: FD   -> Ptr SockAddr -> CInt -> IO CInt
foreign import ccall      "sys/socket.h bind"       c_bind       :: FD   -> Ptr SockAddr -> CInt -> IO CInt
foreign import ccall safe "sys/socket.h accept"     c_accept     :: FD   -> Ptr SockAddr -> Ptr CInt -> IO FD
foreign import ccall      "sys/socket.h listen"     c_listen     :: FD   -> CInt -> IO CInt
foreign import ccall safe "sys/socket.h send"       c_send       :: FD   -> Ptr Word8 -> CSize -> CInt -> IO CInt
foreign import ccall safe "sys/socket.h recv"       c_recv       :: FD   -> Ptr Word8 -> CSize -> CInt -> IO CInt
foreign import ccall      "sys/socket.h setsockopt" c_setsockopt :: FD   -> CInt -> CInt -> Ptr CInt -> CInt -> IO CInt
foreign import ccall      "unistd.h close"          c_close      :: FD   -> IO CInt

#ifdef __MHS__
foreign import ccall "fcntl.h fcntl" c_fcntl :: CInt -> CInt -> CInt -> IO CInt
foreign import ccall "sys/socket.h getsockopt" c_getsockopt :: CInt -> CInt -> CInt -> Ptr CInt -> Ptr CInt -> IO CInt

foreign import capi "fcntl.h value F_SETFL"    f_SETFL :: CInt
foreign import capi "fcntl.h value O_NONBLOCK" o_NONBLOCK :: CInt

fdInt :: CInt -> Int
fdInt = fromIntegral

-- | Read SO_ERROR after a non-blocking connect completes.
peekSockError :: CInt -> IO CInt
peekSockError fd =
    alloca $ \errPtr -> alloca $ \lenPtr -> do
        poke lenPtr (cSizeOf (0 :: CInt))
        _ <- c_getsockopt fd sOL_SOCKET sO_ERROR errPtr lenPtr
        peek errPtr
#endif

foreign import capi "sys/socket.h value SOL_SOCKET"   sOL_SOCKET   :: CInt
foreign import capi "sys/socket.h value SO_REUSEADDR" sO_REUSEADDR :: CInt
foreign import capi "sys/socket.h value SO_DEBUG"     sO_DEBUG     :: CInt
foreign import capi "sys/socket.h value SO_TYPE"      sO_TYPE      :: CInt
foreign import capi "sys/socket.h value SO_ERROR"     sO_ERROR     :: CInt


-- ---------------------------------------------------------------------------
-- Operations
-- ---------------------------------------------------------------------------

-- | Create a new socket.
socket :: Domain -> StreamType -> IO Socket
socket d st = do
    fd <- throwErrnoIfMinus1 "socket" $
              c_socket (cFromEnum d) (cFromEnum st) 0
#ifdef __MHS__
    let sock = Socket fd
    setsocketopt (Socket fd) O_NONBLOCK (error "do not evaluate")
#endif
    return (Socket fd)

-- | Set a socket option. There should probably be a @getsocketopt@ as well, but there isn't for now.
setsocketopt :: Socket -> SockOpt -> Int -> IO ()
#ifdef __MHS__
setsocketopt (Socket fd) O_NONBLOCK _ =
    throwErrnoIfMinus1_ "setsocketopt/O_NONBLOCK" $
        c_fcntl fd f_SETFL o_NONBLOCK
#else
setsocketopt _ O_NONBLOCK _ = return () -- sockets created via GHC already have this setting from the IO manager?
#endif
setsocketopt (Socket fd) so value =
    with (fromIntegral value :: CInt) $ \ opt ->
        throwErrnoIfMinus1_ "setsocketopt" $ do
            let (option, level) = case so of
                  SO_REUSEADDR -> (sO_REUSEADDR, sOL_SOCKET)
                  SO_DEBUG     -> (sO_DEBUG,     sOL_SOCKET)
                  SO_TYPE      -> (sO_TYPE,      sOL_SOCKET)
            c_setsockopt fd level option opt (cSizeOf (0 :: CInt))

-- | Close a socket.
close :: Socket -> IO ()
close (Socket fd@(CInt n)) =
    throwErrnoIfMinus1_ ("close socket " ++ show n) $
        c_close fd

-- | Connect.  Throws on failure.
connect :: Socket -> SockAddr -> IO ()
connect (Socket fd) sockaddr =
#ifdef __MHS__
    with sockaddr $ \p -> do
        r <- c_connect fd p (cSizeOf sockaddr)
        if r /= -1 then return () else do
            errno <- getErrno
            if errno == eINPROGRESS
              then do waitForWriteFD (fdInt fd)
                      err <- peekSockError fd
                      if err /= 0 then do
                        setErrno (Errno err)
                        throwErrno "connect"
                      else
                        return ()
              else throwErrno "connect"
#else
    with sockaddr $ \p ->
        throwErrnoIfMinus1_ "connect" $
            c_connect fd p (cSizeOf sockaddr)
#endif

-- | Same as 'connect', but returns @False@ rather than throwing on error.
connect' :: Socket -> SockAddr -> IO Bool
connect' (Socket fd) sockaddr =
#ifdef __MHS__
    with sockaddr $ \p -> do
        r <- c_connect fd p (cSizeOf sockaddr)
        if r /= -1 then return True else do
            errno <- getErrno
            if errno == eINPROGRESS
              then do waitForWriteFD (fdInt fd)
                      err <- peekSockError fd
                      return (err == 0)
              else return False
#else
    with sockaddr $ \p -> do
        CInt e <- c_connect fd p (cSizeOf sockaddr)
        return (e >= 0)
#endif

-- | Bind a socket to an address.
bind :: Socket -> SockAddr -> IO ()
bind (Socket fd) sockaddr =
    with sockaddr $ \p ->
        throwErrnoIfMinus1_ "bind" $
            c_bind fd p (cSizeOf sockaddr)

-- | Accept an incoming connection. The returned 'Socket' is already in non-blocking mode.
accept :: Socket -> IO (Socket, SockAddr)
accept (Socket serverFd) =
#ifdef __MHS__
    allocaBytes (sizeOf (undefined :: SockAddr)) $ \p ->
        with (cSizeOf (undefined :: SockAddr)) $ \pSize ->
            go p pSize
  where
    go p pSize = do
        r <- c_accept serverFd p pSize
        if r /= -1
          then do addr <- peek p
                  throwErrnoIfMinus1_ "accept/setnonblock" $
                      c_fcntl r f_SETFL o_NONBLOCK
                  return (Socket r, addr)
          else do errno <- getErrno
                  if errno == eAGAIN || errno == eWOULDBLOCK
                    then waitForReadFD (fdInt serverFd) >> go p pSize
                    else throwErrno "accept"
#else
    allocaBytes (sizeOf (undefined :: SockAddr)) $ \p ->
        with (cSizeOf (undefined :: SockAddr)) $ \pSize -> do
            clientFd <- throwErrnoIfMinus1 "accept" $
                            c_accept serverFd p pSize
            addr <- peek p
            return (Socket clientFd, addr)
#endif

-- | Set the socket to listening mode.
listen :: Socket -> Int -> IO ()
listen (Socket fd) n =
    throwErrnoIfMinus1_ "listen" $
        c_listen fd (fromIntegral n)

-- | Send raw bytes. Returns the number of bytes that were actually sent.
sendBuf :: Socket -> Ptr Word8 -> Int -> IO Int
#ifdef __MHS__
sendBuf (Socket fd) buf len = go
  where
    go = do
        CInt n <- c_send fd buf (CSize (fromIntegral len)) (CInt 0)
        if n /= -1 then return (fromIntegral n) else do
            errno <- getErrno
            if errno == eAGAIN || errno == eWOULDBLOCK
              then waitForWriteFD (fdInt fd) >> go
              else throwErrno "sendBuf"
#else
sendBuf (Socket fd) buf len =
    throwErrnoIfMinus1 "sendBuf" $ do
        CInt n <- c_send fd buf (CSize (fromIntegral len)) (CInt 0)
        return (fromIntegral n)
#endif

-- | Send raw bytes. Sends the total number of bytes.
sendBufFull :: Socket -> Ptr Word8 -> Int -> IO ()
sendBufFull sock ptr len | len == 0 = return ()
                         | otherwise = do
    n <- sendBuf sock ptr len
    sendBufFull sock (plusPtr ptr n) (len - n)

-- | A helper function that behaves like 'sendBufFull', but which takes a 'String' rather than a pointer to a buffer.
sendString :: Socket -> String -> IO Int
sendString sock str =
    withCAStringLen str $ \(ptr, len) ->
        sendBuf sock (castPtr ptr) len

-- | A helper function that behaves like 'sendBufFull', but which takes a 'ByteString' rather than a pointer to a buffer.
sendByteString :: Socket -> BS.ByteString -> IO ()
sendByteString sock bs =
    BS.unsafeUseAsCStringLen bs $ \ (ptr, len) ->
        sendBufFull sock (castPtr ptr) len

-- | @recvBuf socket buf len@ -- read at most @len@ bytes from @socket@ into @buf@. Returns the number of bytes that was read.
recvBuf :: Socket -> Ptr Word8 -> Int -> IO Int
#ifdef __MHS__
recvBuf (Socket fd) buf len = go
  where
    go = do
        CInt n <- c_recv fd buf (CSize (fromIntegral len)) (CInt 0)
        if n /= -1 then return (fromIntegral n) else do
            errno <- getErrno
            if errno == eAGAIN || errno == eWOULDBLOCK
              then waitForReadFD (fdInt fd) >> go
              else throwErrno "recvBuf"
#else
recvBuf (Socket fd) buf len =
    throwErrnoIfMinus1 "recvBuf" $ do
        CInt n <- c_recv fd buf (CSize (fromIntegral len)) (CInt 0)
        return (fromIntegral n)
#endif

-- | @recvBufFull socket buf len@ -- read exactly len bytes from @socket@ into @buf@.
recvBufFull :: Socket -> Ptr Word8 -> Int -> IO ()
recvBufFull sock ptr len | len == 0 = return ()
                         | otherwise = do
    n <- recvBuf sock ptr len
    recvBufFull sock (plusPtr ptr n) (len - n)

-- | @recvString socket len@ -- Receive up to @len@ bytes and decode as a 'String'.
recvString :: Socket -> Int -> IO String
recvString sock maxLen =
    allocaBytes maxLen $ \buf -> do
        n <- recvBuf sock buf maxLen
        peekCAStringLen (castPtr buf, n)

-- | @recvByteString socket len@ -- read at most @len@ bytes and decode as a 'ByteString'.
recvByteString :: Socket -> Int -> IO BS.ByteString
recvByteString sock maxLen =
    allocaBytes maxLen $ \buf -> do
        n <- recvBuf sock buf maxLen
        BS.packCStringLen (castPtr buf, n)

-- | @recvByteStringFull socket len@ -- read exactly @len@ bytes and decode as a 'ByteString'.
recvByteStringFull :: Socket -> Int -> IO BS.ByteString
recvByteStringFull sock len = do
    buf <- mallocBytes len
    recvBufFull sock buf len
#ifdef __MHS__
    BS.unsafePackMallocCStringLen (castPtr buf) len
#else
    BS.unsafePackMallocCStringLen (castPtr buf, len)
#endif