{-# LINE 1 "Network/Socket/ByteString/IO.hsc" #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE OverloadedStrings #-}



-- |
-- Module      : Network.Socket.ByteString
-- Copyright   : (c) Johan Tibell 2007-2010
-- License     : BSD-style
--
-- Maintainer  : johan.tibell@gmail.com
-- Stability   : stable
-- Portability : portable
--
-- This module provides access to the BSD /socket/ interface.  This
-- module is generally more efficient than the 'String' based network
-- functions in 'Network.Socket'.  For detailed documentation, consult
-- your favorite POSIX socket reference. All functions communicate
-- failures by converting the error number to 'System.IO.IOError'.
--
-- This module is made to be imported with 'Network.Socket' like so:
--
-- > import Network.Socket hiding (send, sendTo, recv, recvFrom)
-- > import Network.Socket.ByteString
--
module Network.Socket.ByteString.IO
    (
    -- * Send data to a socket
      send
    , sendAll
    , sendTo
    , sendAllTo

    -- ** Vectored I/O
    -- $vectored
    , sendMany
    , sendManyTo

    -- * Receive data from a socket
    , recv
    , recvFrom
    , waitWhen0
    ) where

import Control.Concurrent (threadWaitWrite, rtsSupportsBoundThreads)
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import Data.ByteString.Internal (createAndTrim)
import Data.ByteString.Unsafe (unsafeUseAsCStringLen)
import Foreign.Marshal.Alloc (allocaBytes)
import Foreign.Ptr (castPtr)

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


{-# LINE 59 "Network/Socket/ByteString/IO.hsc" #-}
import Control.Monad (zipWithM_)
import Foreign.Marshal.Array (allocaArray)
import Foreign.Marshal.Utils (with)
import Foreign.Ptr (Ptr, plusPtr)
import Foreign.Storable (Storable(..))
import Network.Socket.Internal

import Network.Socket.ByteString.IOVec (IOVec(..))
import Network.Socket.ByteString.MsgHdr (MsgHdr(..))

{-# LINE 69 "Network/Socket/ByteString/IO.hsc" #-}

-- ----------------------------------------------------------------------------
-- Sending

-- | Send data to the socket.  The socket must be connected to a
-- remote socket.  Returns the number of bytes sent. Applications are
-- responsible for ensuring that all data has been sent.
send :: Socket     -- ^ Connected socket
     -> ByteString  -- ^ Data to send
     -> IO Int      -- ^ Number of bytes sent
send :: Socket -> ByteString -> IO Int
send Socket
s ByteString
xs = ByteString -> (CStringLen -> IO Int) -> IO Int
forall a. ByteString -> (CStringLen -> IO a) -> IO a
unsafeUseAsCStringLen ByteString
xs ((CStringLen -> IO Int) -> IO Int)
-> (CStringLen -> IO Int) -> IO Int
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
str, Int
len) ->
    Socket -> Ptr Word8 -> Int -> IO Int
sendBuf Socket
s (Ptr CChar -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr CChar
str) Int
len

waitWhen0 :: Int -> Socket -> IO ()
waitWhen0 :: Int -> Socket -> IO ()
waitWhen0 Int
0 Socket
s = Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
rtsSupportsBoundThreads (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 -> Fd -> IO ()
threadWaitWrite (Fd -> IO ()) -> Fd -> IO ()
forall a b. (a -> b) -> a -> b
$ CInt -> Fd
forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
fd
waitWhen0 Int
_ Socket
_ = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Send data to the socket.  The socket must be connected to a
-- remote socket.  Unlike 'send', this function continues to send data
-- until either all data has been sent or an error occurs.  On error,
-- an exception is raised, and there is no way to determine how much
-- data, if any, was successfully sent.
sendAll :: Socket     -- ^ Connected socket
        -> ByteString  -- ^ Data to send
        -> IO ()
sendAll :: Socket -> ByteString -> IO ()
sendAll Socket
_ ByteString
"" = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
sendAll Socket
s ByteString
bs = do
    Int
sent <- Socket -> ByteString -> IO Int
send Socket
s ByteString
bs
    Int -> Socket -> IO ()
waitWhen0 Int
sent Socket
s
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
sent Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Socket -> ByteString -> IO ()
sendAll Socket
s (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
B.drop Int
sent ByteString
bs

-- | Send data to the socket.  The recipient can be specified
-- explicitly, so the socket need not be in a connected state.
-- Returns the number of bytes sent. Applications are responsible for
-- ensuring that all data has been sent.
sendTo :: SocketAddress sa =>
          Socket     -- ^ Socket
       -> ByteString  -- ^ Data to send
       -> sa    -- ^ Recipient address
       -> IO Int      -- ^ Number of bytes sent
sendTo :: Socket -> ByteString -> sa -> IO Int
sendTo Socket
s ByteString
xs sa
sa =
    ByteString -> (CStringLen -> IO Int) -> IO Int
forall a. ByteString -> (CStringLen -> IO a) -> IO a
unsafeUseAsCStringLen ByteString
xs ((CStringLen -> IO Int) -> IO Int)
-> (CStringLen -> IO Int) -> IO Int
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
str, Int
len) -> Socket -> Ptr CChar -> Int -> sa -> IO Int
forall sa a.
SocketAddress sa =>
Socket -> Ptr a -> Int -> sa -> IO Int
sendBufTo Socket
s Ptr CChar
str Int
len sa
sa

-- | Send data to the socket. The recipient can be specified
-- explicitly, so the socket need not be in a connected state.  Unlike
-- 'sendTo', this function continues to send data until either all
-- data has been sent or an error occurs.  On error, an exception is
-- raised, and there is no way to determine how much data, if any, was
-- successfully sent.
sendAllTo :: SocketAddress sa =>
             Socket     -- ^ Socket
          -> ByteString  -- ^ Data to send
          -> sa    -- ^ Recipient address
          -> IO ()
sendAllTo :: Socket -> ByteString -> sa -> IO ()
sendAllTo Socket
_ ByteString
"" sa
_  = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
sendAllTo Socket
s ByteString
xs sa
sa = do
    Int
sent <- Socket -> ByteString -> sa -> IO Int
forall sa. SocketAddress sa => Socket -> ByteString -> sa -> IO Int
sendTo Socket
s ByteString
xs sa
sa
    Int -> Socket -> IO ()
waitWhen0 Int
sent Socket
s
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
sent Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Socket -> ByteString -> sa -> IO ()
forall sa. SocketAddress sa => Socket -> ByteString -> sa -> IO ()
sendAllTo Socket
s (Int -> ByteString -> ByteString
B.drop Int
sent ByteString
xs) sa
sa

-- | Send data to the socket.  The socket must be in a connected
-- state.  The data is sent as if the parts have been concatenated.
-- This function continues to send data until either all data has been
-- sent or an error occurs.  On error, an exception is raised, and
-- there is no way to determine how much data, if any, was
-- successfully sent.
sendMany :: Socket       -- ^ Connected socket
         -> [ByteString]  -- ^ Data to send
         -> IO ()

sendMany :: Socket -> [ByteString] -> IO ()
{-# LINE 140 "Network/Socket/ByteString/IO.hsc" #-}
sendMany _ [] = return ()
sendMany s cs = do
    sent <- sendManyInner
    waitWhen0 sent s
    when (sent >= 0) $ sendMany s $ remainingChunks sent cs
  where
    sendManyInner =
      fmap fromIntegral . withIOVec cs $ \(iovsPtr, iovsLen) ->
          withFdSocket s $ \fd -> do
              let len =  fromIntegral $ min iovsLen (1024)
{-# LINE 150 "Network/Socket/ByteString/IO.hsc" #-}
              throwSocketErrorWaitWrite s "Network.Socket.ByteString.sendMany" $
                  c_writev fd iovsPtr len

{-# LINE 155 "Network/Socket/ByteString/IO.hsc" #-}

-- | Send data to the socket.  The recipient can be specified
-- explicitly, so the socket need not be in a connected state.  The
-- data is sent as if the parts have been concatenated.  This function
-- continues to send data until either all data has been sent or an
-- error occurs.  On error, an exception is raised, and there is no
-- way to determine how much data, if any, was successfully sent.
sendManyTo :: Socket       -- ^ Socket
           -> [ByteString]  -- ^ Data to send
           -> SockAddr      -- ^ Recipient address
           -> IO ()

{-# LINE 167 "Network/Socket/ByteString/IO.hsc" #-}
sendManyTo _ [] _    = return ()
sendManyTo s cs addr = do
    sent <- fromIntegral <$> sendManyToInner
    waitWhen0 sent s
    when (sent >= 0) $ sendManyTo s (remainingChunks sent cs) addr
  where
    sendManyToInner =
      withSockAddr addr $ \addrPtr addrSize ->
        withIOVec cs $ \(iovsPtr, iovsLen) -> do
          let msgHdr = MsgHdr
                addrPtr (fromIntegral addrSize)
                iovsPtr (fromIntegral iovsLen)
          withFdSocket s $ \fd ->
              with msgHdr $ \msgHdrPtr ->
                throwSocketErrorWaitWrite s "Network.Socket.ByteString.sendManyTo" $
                  c_sendmsg fd msgHdrPtr 0

{-# LINE 186 "Network/Socket/ByteString/IO.hsc" #-}

-- ----------------------------------------------------------------------------
-- Receiving

-- | Receive data from the socket.  The socket must be in a connected
-- state.  This function may return fewer bytes than specified.  If
-- the message is longer than the specified length, it may be
-- discarded depending on the type of socket.  This function may block
-- until a message arrives.
--
-- Considering hardware and network realities, the maximum number of bytes to
-- receive should be a small power of 2, e.g., 4096.
--
-- For TCP sockets, a zero length return value means the peer has
-- closed its half side of the connection.
recv :: Socket        -- ^ Connected socket
     -> Int            -- ^ Maximum number of bytes to receive
     -> IO ByteString  -- ^ Data received
recv :: Socket -> Int -> IO ByteString
recv Socket
s Int
nbytes
    | Int
nbytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = IOError -> IO ByteString
forall a. IOError -> IO a
ioError (String -> IOError
mkInvalidRecvArgError String
"Network.Socket.ByteString.recv")
    | Bool
otherwise  = Int -> (Ptr Word8 -> IO Int) -> IO ByteString
createAndTrim Int
nbytes ((Ptr Word8 -> IO Int) -> IO ByteString)
-> (Ptr Word8 -> IO Int) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> Socket -> Ptr Word8 -> Int -> IO Int
recvBuf Socket
s Ptr Word8
ptr Int
nbytes

-- | Receive data from the socket.  The socket need not be in a
-- connected state.  Returns @(bytes, address)@ where @bytes@ is a
-- 'ByteString' representing the data received and @address@ is a
-- 'SockAddr' representing the address of the sending socket.
--
-- If the first return value is zero, it means EOF.
recvFrom :: SocketAddress sa =>
            Socket                     -- ^ Socket
         -> Int                        -- ^ Maximum number of bytes to receive
         -> IO (ByteString, sa)  -- ^ Data received and sender address
recvFrom :: Socket -> Int -> IO (ByteString, sa)
recvFrom Socket
sock Int
nbytes =
    Int -> (Ptr CChar -> IO (ByteString, sa)) -> IO (ByteString, sa)
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
nbytes ((Ptr CChar -> IO (ByteString, sa)) -> IO (ByteString, sa))
-> (Ptr CChar -> IO (ByteString, sa)) -> IO (ByteString, sa)
forall a b. (a -> b) -> a -> b
$ \Ptr CChar
ptr -> do
        (Int
len, sa
sockaddr) <- Socket -> Ptr CChar -> Int -> IO (Int, sa)
forall sa a.
SocketAddress sa =>
Socket -> Ptr a -> Int -> IO (Int, sa)
recvBufFrom Socket
sock Ptr CChar
ptr Int
nbytes
        ByteString
str <- CStringLen -> IO ByteString
B.packCStringLen (Ptr CChar
ptr, Int
len)
        (ByteString, sa) -> IO (ByteString, sa)
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
str, sa
sockaddr)

-- ----------------------------------------------------------------------------
-- Not exported


{-# LINE 228 "Network/Socket/ByteString/IO.hsc" #-}
-- | Suppose we try to transmit a list of chunks @cs@ via a gathering write
-- operation and find that @n@ bytes were sent. Then @remainingChunks n cs@ is
-- list of chunks remaining to be sent.
remainingChunks :: Int -> [ByteString] -> [ByteString]
remainingChunks :: Int -> [ByteString] -> [ByteString]
remainingChunks Int
_ [] = []
remainingChunks Int
i (ByteString
x:[ByteString]
xs)
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len        = Int -> ByteString -> ByteString
B.drop Int
i ByteString
x ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
xs
    | Bool
otherwise      = let i' :: Int
i' = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len in Int
i' Int -> [ByteString] -> [ByteString]
`seq` Int -> [ByteString] -> [ByteString]
remainingChunks Int
i' [ByteString]
xs
  where
    len :: Int
len = ByteString -> Int
B.length ByteString
x

-- | @withIOVec cs f@ executes the computation @f@, passing as argument a pair
-- consisting of a pointer to a temporarily allocated array of pointers to
-- IOVec made from @cs@ and the number of pointers (@length cs@).
-- /Unix only/.
withIOVec :: [ByteString] -> ((Ptr IOVec, Int) -> IO a) -> IO a
withIOVec :: [ByteString] -> ((Ptr IOVec, Int) -> IO a) -> IO a
withIOVec [ByteString]
cs (Ptr IOVec, Int) -> IO a
f =
    Int -> (Ptr IOVec -> IO a) -> IO a
forall a b. Storable a => Int -> (Ptr a -> IO b) -> IO b
allocaArray Int
csLen ((Ptr IOVec -> IO a) -> IO a) -> (Ptr IOVec -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr IOVec
aPtr -> do
        (Ptr IOVec -> ByteString -> IO ())
-> [Ptr IOVec] -> [ByteString] -> IO ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ Ptr IOVec -> ByteString -> IO ()
pokeIov (Ptr IOVec -> [Ptr IOVec]
forall a. Ptr a -> [Ptr a]
ptrs Ptr IOVec
aPtr) [ByteString]
cs
        (Ptr IOVec, Int) -> IO a
f (Ptr IOVec
aPtr, Int
csLen)
  where
    csLen :: Int
csLen = [ByteString] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
cs
    ptrs :: Ptr a -> [Ptr a]
ptrs = (Ptr a -> Ptr a) -> Ptr a -> [Ptr a]
forall a. (a -> a) -> a -> [a]
iterate (Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` IOVec -> Int
forall a. Storable a => a -> Int
sizeOf (IOVec
forall a. HasCallStack => a
undefined :: IOVec))
    pokeIov :: Ptr IOVec -> ByteString -> IO ()
pokeIov Ptr IOVec
ptr ByteString
s =
        ByteString -> (CStringLen -> IO ()) -> IO ()
forall a. ByteString -> (CStringLen -> IO a) -> IO a
unsafeUseAsCStringLen ByteString
s ((CStringLen -> IO ()) -> IO ()) -> (CStringLen -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
sPtr, Int
sLen) ->
        Ptr IOVec -> IOVec -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr IOVec
ptr (IOVec -> IO ()) -> IOVec -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr CChar -> CSize -> IOVec
IOVec Ptr CChar
sPtr (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
sLen)

{-# LINE 255 "Network/Socket/ByteString/IO.hsc" #-}