1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE RecordWildCards #-}
3{-# LANGUAGE OverloadedStrings #-}
4
5#include "HsNet.h"
6
7-- |
8-- Module      : Network.Socket.ByteString.IO
9-- Copyright   : (c) Johan Tibell 2007-2010
10-- License     : BSD-style
11--
12-- Maintainer  : johan.tibell@gmail.com
13-- Stability   : stable
14-- Portability : portable
15--
16module Network.Socket.ByteString.IO
17    (
18    -- * Send data to a socket
19      send
20    , sendAll
21    , sendTo
22    , sendAllTo
23
24    -- ** Vectored I/O
25    -- $vectored
26    , sendMany
27    , sendManyTo
28
29    -- * Receive data from a socket
30    , recv
31    , recvFrom
32    , waitWhen0
33
34    -- * Advanced send and recv
35    , sendMsg
36    , recvMsg
37    , MsgFlag(..)
38    , Cmsg(..)
39    ) where
40
41import Control.Concurrent (threadWaitWrite, rtsSupportsBoundThreads)
42import Data.ByteString (ByteString)
43import qualified Data.ByteString as B
44import Data.ByteString.Internal (createAndTrim)
45import Data.ByteString.Unsafe (unsafeUseAsCStringLen)
46import Foreign.Marshal.Alloc (allocaBytes)
47
48import Network.Socket.Buffer
49import Network.Socket.ByteString.Internal
50import Network.Socket.Imports
51import Network.Socket.Types
52
53import Data.ByteString.Internal (create, ByteString(..))
54import Foreign.ForeignPtr (withForeignPtr)
55import Foreign.Marshal.Utils (with)
56import Network.Socket.Internal
57
58import Network.Socket.Flag
59
60#if !defined(mingw32_HOST_OS)
61import Network.Socket.Posix.Cmsg
62import Network.Socket.Posix.IOVec
63import Network.Socket.Posix.MsgHdr (MsgHdr(..))
64#else
65import Foreign.Marshal.Alloc (alloca)
66import Network.Socket.Win32.Cmsg
67import Network.Socket.Win32.WSABuf
68import Network.Socket.Win32.MsgHdr (MsgHdr(..))
69#endif
70
71-- ----------------------------------------------------------------------------
72-- Sending
73
74-- | Send data to the socket.  The socket must be connected to a
75-- remote socket.  Returns the number of bytes sent. Applications are
76-- responsible for ensuring that all data has been sent.
77send :: Socket     -- ^ Connected socket
78     -> ByteString  -- ^ Data to send
79     -> IO Int      -- ^ Number of bytes sent
80send s xs = unsafeUseAsCStringLen xs $ \(str, len) ->
81    sendBuf s (castPtr str) len
82
83waitWhen0 :: Int -> Socket -> IO ()
84waitWhen0 0 s = when rtsSupportsBoundThreads $
85    withFdSocket s $ \fd -> threadWaitWrite $ fromIntegral fd
86waitWhen0 _ _ = return ()
87
88-- | Send data to the socket.  The socket must be connected to a
89-- remote socket.  Unlike 'send', this function continues to send data
90-- until either all data has been sent or an error occurs.  On error,
91-- an exception is raised, and there is no way to determine how much
92-- data, if any, was successfully sent.
93sendAll :: Socket     -- ^ Connected socket
94        -> ByteString  -- ^ Data to send
95        -> IO ()
96sendAll _ "" = return ()
97sendAll s bs = do
98    sent <- send s bs
99    waitWhen0 sent s
100    when (sent >= 0) $ sendAll s $ B.drop sent bs
101
102-- | Send data to the socket.  The recipient can be specified
103-- explicitly, so the socket need not be in a connected state.
104-- Returns the number of bytes sent. Applications are responsible for
105-- ensuring that all data has been sent.
106sendTo :: SocketAddress sa =>
107          Socket     -- ^ Socket
108       -> ByteString  -- ^ Data to send
109       -> sa    -- ^ Recipient address
110       -> IO Int      -- ^ Number of bytes sent
111sendTo s xs sa =
112    unsafeUseAsCStringLen xs $ \(str, len) -> sendBufTo s str len sa
113
114-- | Send data to the socket. The recipient can be specified
115-- explicitly, so the socket need not be in a connected state.  Unlike
116-- 'sendTo', this function continues to send data until either all
117-- data has been sent or an error occurs.  On error, an exception is
118-- raised, and there is no way to determine how much data, if any, was
119-- successfully sent.
120sendAllTo :: SocketAddress sa =>
121             Socket     -- ^ Socket
122          -> ByteString  -- ^ Data to send
123          -> sa    -- ^ Recipient address
124          -> IO ()
125sendAllTo _ "" _  = return ()
126sendAllTo s xs sa = do
127    sent <- sendTo s xs sa
128    waitWhen0 sent s
129    when (sent >= 0) $ sendAllTo s (B.drop sent xs) sa
130
131-- | Send data to the socket.  The socket must be in a connected
132-- state.  The data is sent as if the parts have been concatenated.
133-- This function continues to send data until either all data has been
134-- sent or an error occurs.  On error, an exception is raised, and
135-- there is no way to determine how much data, if any, was
136-- successfully sent.
137sendMany :: Socket       -- ^ Connected socket
138         -> [ByteString]  -- ^ Data to send
139         -> IO ()
140sendMany _ [] = return ()
141sendMany s cs = do
142    sent <- sendManyInner
143    waitWhen0 sent s
144    when (sent >= 0) $ sendMany s $ remainingChunks sent cs
145  where
146    sendManyInner =
147#if !defined(mingw32_HOST_OS)
148      fmap fromIntegral . withIOVecfromBS cs $ \(iovsPtr, iovsLen) ->
149          withFdSocket s $ \fd -> do
150              let len =  fromIntegral $ min iovsLen (#const IOV_MAX)
151              throwSocketErrorWaitWrite s "Network.Socket.ByteString.sendMany" $
152                  c_writev fd iovsPtr len
153#else
154      fmap fromIntegral . withWSABuffromBS cs $ \(wsabsPtr, wsabsLen) ->
155          withFdSocket s $ \fd -> do
156              let len =  fromIntegral wsabsLen
157              alloca $ \send_ptr -> do
158                _ <- throwSocketErrorWaitWrite s "Network.Socket.ByteString.sendMany" $
159                       c_wsasend fd wsabsPtr len send_ptr  0 nullPtr nullPtr
160                peek send_ptr
161#endif
162
163-- | Send data to the socket.  The recipient can be specified
164-- explicitly, so the socket need not be in a connected state.  The
165-- data is sent as if the parts have been concatenated.  This function
166-- continues to send data until either all data has been sent or an
167-- error occurs.  On error, an exception is raised, and there is no
168-- way to determine how much data, if any, was successfully sent.
169sendManyTo :: Socket       -- ^ Socket
170           -> [ByteString]  -- ^ Data to send
171           -> SockAddr      -- ^ Recipient address
172           -> IO ()
173sendManyTo _ [] _    = return ()
174sendManyTo s cs addr = do
175    sent <- fromIntegral <$> sendManyToInner
176    waitWhen0 sent s
177    when (sent >= 0) $ sendManyTo s (remainingChunks sent cs) addr
178  where
179    sendManyToInner =
180      withSockAddr addr $ \addrPtr addrSize ->
181#if !defined(mingw32_HOST_OS)
182        withIOVecfromBS cs $ \(iovsPtr, iovsLen) -> do
183          let msgHdr = MsgHdr {
184                  msgName    = addrPtr
185                , msgNameLen = fromIntegral addrSize
186                , msgIov     = iovsPtr
187                , msgIovLen  = fromIntegral iovsLen
188                , msgCtrl    = nullPtr
189                , msgCtrlLen = 0
190                , msgFlags   = 0
191                }
192          withFdSocket s $ \fd ->
193              with msgHdr $ \msgHdrPtr ->
194                throwSocketErrorWaitWrite s "Network.Socket.ByteString.sendManyTo" $
195                  c_sendmsg fd msgHdrPtr 0
196#else
197        withWSABuffromBS cs $ \(wsabsPtr, wsabsLen) -> do
198          let msgHdr = MsgHdr {
199                  msgName      = addrPtr
200                , msgNameLen   = fromIntegral addrSize
201                , msgBuffer    = wsabsPtr
202                , msgBufferLen = fromIntegral wsabsLen
203                , msgCtrl      = nullPtr
204                , msgCtrlLen   = 0
205                , msgFlags     = 0
206                }
207          withFdSocket s $ \fd ->
208              with msgHdr $ \msgHdrPtr ->
209                alloca $ \send_ptr -> do
210                  _ <- throwSocketErrorWaitWrite s "Network.Socket.ByteString.sendManyTo" $
211                          c_sendmsg fd msgHdrPtr 0 send_ptr nullPtr nullPtr
212                  peek send_ptr
213#endif
214
215-- ----------------------------------------------------------------------------
216-- Receiving
217
218-- | Receive data from the socket.  The socket must be in a connected
219-- state.  This function may return fewer bytes than specified.  If
220-- the message is longer than the specified length, it may be
221-- discarded depending on the type of socket.  This function may block
222-- until a message arrives.
223--
224-- Considering hardware and network realities, the maximum number of bytes to
225-- receive should be a small power of 2, e.g., 4096.
226--
227-- For TCP sockets, a zero length return value means the peer has
228-- closed its half side of the connection.
229recv :: Socket        -- ^ Connected socket
230     -> Int            -- ^ Maximum number of bytes to receive
231     -> IO ByteString  -- ^ Data received
232recv s nbytes
233    | nbytes < 0 = ioError (mkInvalidRecvArgError "Network.Socket.ByteString.recv")
234    | otherwise  = createAndTrim nbytes $ \ptr -> recvBuf s ptr nbytes
235
236-- | Receive data from the socket.  The socket need not be in a
237-- connected state.  Returns @(bytes, address)@ where @bytes@ is a
238-- 'ByteString' representing the data received and @address@ is a
239-- 'SockAddr' representing the address of the sending socket.
240--
241-- If the first return value is zero, it means EOF.
242recvFrom :: SocketAddress sa =>
243            Socket                     -- ^ Socket
244         -> Int                        -- ^ Maximum number of bytes to receive
245         -> IO (ByteString, sa)  -- ^ Data received and sender address
246recvFrom sock nbytes =
247    allocaBytes nbytes $ \ptr -> do
248        (len, sockaddr) <- recvBufFrom sock ptr nbytes
249        str <- B.packCStringLen (ptr, len)
250        return (str, sockaddr)
251
252-- ----------------------------------------------------------------------------
253-- Not exported
254
255
256-- | Suppose we try to transmit a list of chunks @cs@ via a gathering write
257-- operation and find that @n@ bytes were sent. Then @remainingChunks n cs@ is
258-- list of chunks remaining to be sent.
259remainingChunks :: Int -> [ByteString] -> [ByteString]
260remainingChunks _ [] = []
261remainingChunks i (x:xs)
262    | i < len        = B.drop i x : xs
263    | otherwise      = let i' = i - len in i' `seq` remainingChunks i' xs
264  where
265    len = B.length x
266
267#if !defined(mingw32_HOST_OS)
268-- | @withIOVecfromBS cs f@ executes the computation @f@, passing as argument a pair
269-- consisting of a pointer to a temporarily allocated array of pointers to
270-- IOVec made from @cs@ and the number of pointers (@length cs@).
271-- /Unix only/.
272withIOVecfromBS :: [ByteString] -> ((Ptr IOVec, Int) -> IO a) -> IO a
273withIOVecfromBS cs f = withBufSizs cs $ \bufsizs -> withIOVec bufsizs f
274#else
275-- | @withWSABuffromBS cs f@ executes the computation @f@, passing as argument a pair
276-- consisting of a pointer to a temporarily allocated array of pointers to
277-- WSABuf made from @cs@ and the number of pointers (@length cs@).
278-- /Windows only/.
279withWSABuffromBS :: [ByteString] -> ((Ptr WSABuf, Int) -> IO a) -> IO a
280withWSABuffromBS cs f = withBufSizs cs $ \bufsizs -> withWSABuf bufsizs f
281#endif
282
283withBufSizs :: [ByteString] -> ([(Ptr Word8, Int)] -> IO a) -> IO a
284withBufSizs bss0 f = loop bss0 id
285  where
286    loop []                    !build = f $ build []
287    loop (PS fptr off len:bss) !build = withForeignPtr fptr $ \ptr -> do
288        let !ptr' = ptr `plusPtr` off
289        loop bss (build . ((ptr',len) :))
290
291-- | Send data to the socket using sendmsg(2).
292sendMsg :: Socket       -- ^ Socket
293        -> SockAddr     -- ^ Destination address
294        -> [ByteString] -- ^ Data to be sent
295        -> [Cmsg]       -- ^ Control messages
296        -> MsgFlag      -- ^ Message flags
297        -> IO Int       -- ^ The length actually sent
298sendMsg _ _    []  _ _ = return 0
299sendMsg s addr bss cmsgs flags = withBufSizs bss $ \bufsizs ->
300    sendBufMsg s addr bufsizs cmsgs flags
301
302-- | Receive data from the socket using recvmsg(2).
303recvMsg :: Socket  -- ^ Socket
304        -> Int     -- ^ The maximum length of data to be received
305                   --   If the total length is not large enough,
306                   --   'MSG_TRUNC' is returned
307        -> Int     -- ^ The buffer size for control messages.
308                   --   If the length is not large enough,
309                   --   'MSG_CTRUNC' is returned
310        -> MsgFlag -- ^ Message flags
311        -> IO (SockAddr, ByteString, [Cmsg], MsgFlag) -- ^ Source address, received data, control messages and message flags
312recvMsg s siz clen flags = do
313    bs@(PS fptr _ _) <- create siz $ \ptr -> zeroMemory ptr (fromIntegral siz)
314    withForeignPtr fptr $ \ptr -> do
315        (addr,len,cmsgs,flags') <- recvBufMsg s [(ptr,siz)] clen flags
316        let bs' | len < siz = PS fptr 0 len
317                | otherwise = bs
318        return (addr, bs', cmsgs, flags')
319