git.iko.soy/streamly
on d003fe5 help sign in
streamly/src/Streamly/Internal/Network/Inet/TCP.hs
458 lines 15 kB history raw
1{-# LANGUAGE CPP #-}
2{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE FlexibleContexts #-}
4{-# LANGUAGE MagicHash #-}
5{-# LANGUAGE RecordWildCards #-}
6{-# LANGUAGE UnboxedTuples #-}
8#include "inline.hs"
10-- |
11-- Module : Streamly.Internal.Network.Inet.TCP
12-- Copyright : (c) 2019 Composewell Technologies
13--
14-- License : BSD3
15-- Maintainer : streamly@composewell.com
16-- Stability : experimental
17-- Portability : GHC
18--
19-- Combinators to build Inet/TCP clients and servers.
21module Streamly.Internal.Network.Inet.TCP
22 (
23 -- * TCP Servers
24 -- ** Unfolds
25 acceptOnAddr
26 , acceptOnAddrWith
27 , acceptOnPort
28 , acceptOnPortWith
29 , acceptOnPortLocal
31 -- ** Streams
32 , connectionsOnAddr
33 , connectionsOnAddrWith
34 , connectionsOnPort
35 , connectionsOnLocalHost
37 -- * TCP clients
38 -- | IP Address based operations.
39 , connect
40 , withConnectionM
42 -- ** Unfolds
43 , usingConnection
44 , read
46 -- ** Streams
47 , withConnection
48 -- *** Source
49 , toBytes
50 -- , readUtf8
51 -- , readLines
52 -- , readFrames
53 -- , readByChunks
55 -- -- * Array Read
56 -- , readArrayUpto
57 -- , readArrayOf
59 -- , readChunksUpto
60 -- , readChunksOf
61 -- , readChunks
63 -- *** Sink
64 , write
65 -- , writeUtf8
66 -- , writeUtf8ByLines
67 -- , writeByFrames
68 , writeWithBufferOf
69 , fromBytes
70 , fromBytesWithBufferOf
72 -- -- * Array Write
73 -- , writeArray
74 , writeChunks
75 , fromChunks
77 -- ** Transformation
78 , transformBytesWith
79 {-
80 -- ** Sink Servers
82 -- These abstractions can be applied to any setting where we need to do a
83 -- sink processing of multiple streams e.g. output from multiple processes
84 -- or data coming from multiple files.
86 -- handle connections concurrently using a specified fold
87 -- , handleConnections
89 -- handle frames concurrently using a specified fold
90 , handleFrames
92 -- merge frames from all connection into a single stream. Frames can be
93 -- created by a specified fold.
94 , mergeFrames
96 -- * UDP Servers
97 , datagrams
98 , datagramsOn
99 -}
100 )
101where
103import Control.Monad.Catch (MonadCatch, MonadMask, bracket)
104import Control.Monad.IO.Class (MonadIO(..))
105import Data.Word (Word8)
106import Network.Socket
107 (Socket, PortNumber, SocketOption(..), Family(..), SockAddr(..),
108 SocketType(..), defaultProtocol, maxListenQueue, tupleToHostAddress,
109 socket)
110import Prelude hiding (read)
112import Streamly (MonadAsync)
113import Streamly.Internal.Data.Fold.Types (Fold(..))
114import Streamly.Internal.Data.SVar (fork)
115import Streamly.Internal.Data.Unfold.Types (Unfold(..))
116import Streamly.Internal.Network.Socket (SockSpec(..), accept, connections)
117import Streamly.Internal.Data.Stream.Serial (SerialT)
118import Streamly.Internal.Memory.Array.Types (Array(..), defaultChunkSize, writeNUnsafe)
119import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
121import qualified Control.Monad.Catch as MC
122import qualified Network.Socket as Net
124import qualified Streamly.Internal.Data.Unfold as UF
125import qualified Streamly.Internal.Memory.Array as A
126import qualified Streamly.Internal.Memory.ArrayStream as AS
127import qualified Streamly.Internal.Data.Fold.Types as FL
128import qualified Streamly.Prelude as S
129import qualified Streamly.Network.Socket as SK
130import qualified Streamly.Internal.Network.Socket as ISK
132-------------------------------------------------------------------------------
133-- Accept (unfolds)
134-------------------------------------------------------------------------------
136{-# INLINE acceptOnAddrWith #-}
137acceptOnAddrWith
138 :: MonadIO m
139 => [(SocketOption, Int)]
140 -> Unfold m ((Word8, Word8, Word8, Word8), PortNumber) Socket
141acceptOnAddrWith opts = UF.lmap f accept
142 where
143 f (addr, port) =
144 (maxListenQueue
145 , SockSpec
146 { sockFamily = AF_INET
147 , sockType = Stream
148 , sockProto = defaultProtocol -- TCP
149 , sockOpts = opts
150 }
151 , SockAddrInet port (tupleToHostAddress addr)
152 )
154-- | Unfold a tuple @(ipAddr, port)@ into a stream of connected TCP sockets.
155-- @ipAddr@ is the local IP address and @port@ is the local port on which
156-- connections are accepted.
157--
158-- @since 0.7.0
159{-# INLINE acceptOnAddr #-}
160acceptOnAddr
161 :: MonadIO m
162 => Unfold m ((Word8, Word8, Word8, Word8), PortNumber) Socket
163acceptOnAddr = acceptOnAddrWith []
165{-# INLINE acceptOnPortWith #-}
166acceptOnPortWith :: MonadIO m
167 => [(SocketOption, Int)]
168 -> Unfold m PortNumber Socket
169acceptOnPortWith opts = UF.supplyFirst (acceptOnAddrWith opts) (0,0,0,0)
171-- | Like 'acceptOnAddr' but binds on the IPv4 address @0.0.0.0@ i.e. on all
172-- IPv4 addresses/interfaces of the machine and listens for TCP connections on
173-- the specified port.
174--
175-- > acceptOnPort = UF.supplyFirst acceptOnAddr (0,0,0,0)
176--
177-- @since 0.7.0
178{-# INLINE acceptOnPort #-}
179acceptOnPort :: MonadIO m => Unfold m PortNumber Socket
180acceptOnPort = UF.supplyFirst acceptOnAddr (0,0,0,0)
182-- | Like 'acceptOnAddr' but binds on the localhost IPv4 address @127.0.0.1@.
183-- The server can only be accessed from the local host, it cannot be accessed
184-- from other hosts on the network.
185--
186-- > acceptOnPortLocal = UF.supplyFirst acceptOnAddr (127,0,0,1)
187--
188-- @since 0.7.0
189{-# INLINE acceptOnPortLocal #-}
190acceptOnPortLocal :: MonadIO m => Unfold m PortNumber Socket
191acceptOnPortLocal = UF.supplyFirst acceptOnAddr (127,0,0,1)
193-------------------------------------------------------------------------------
194-- Accept (streams)
195-------------------------------------------------------------------------------
197{-# INLINE connectionsOnAddrWith #-}
198connectionsOnAddrWith
199 :: MonadAsync m
200 => [(SocketOption, Int)]
201 -> (Word8, Word8, Word8, Word8)
202 -> PortNumber
203 -> SerialT m Socket
204connectionsOnAddrWith opts addr port =
205 connections maxListenQueue SockSpec
206 { sockFamily = AF_INET
207 , sockType = Stream
208 , sockProto = defaultProtocol
209 , sockOpts = opts
210 }
211 (SockAddrInet port (tupleToHostAddress addr))
213-- | Like 'connections' but binds on the specified IPv4 address of the machine
214-- and listens for TCP connections on the specified port.
215--
216-- /Internal/
217{-# INLINE connectionsOnAddr #-}
218connectionsOnAddr
219 :: MonadAsync m
220 => (Word8, Word8, Word8, Word8)
221 -> PortNumber
222 -> SerialT m Socket
223connectionsOnAddr = connectionsOnAddrWith []
225-- | Like 'connections' but binds on the IPv4 address @0.0.0.0@ i.e. on all
226-- IPv4 addresses/interfaces of the machine and listens for TCP connections on
227-- the specified port.
228--
229-- > connectionsOnPort = connectionsOnAddr (0,0,0,0)
230--
231-- /Internal/
232{-# INLINE connectionsOnPort #-}
233connectionsOnPort :: MonadAsync m => PortNumber -> SerialT m Socket
234connectionsOnPort = connectionsOnAddr (0,0,0,0)
236-- | Like 'connections' but binds on the localhost IPv4 address @127.0.0.1@.
237-- The server can only be accessed from the local host, it cannot be accessed
238-- from other hosts on the network.
239--
240-- > connectionsOnLocalHost = connectionsOnAddr (127,0,0,1)
241--
242-- /Internal/
243{-# INLINE connectionsOnLocalHost #-}
244connectionsOnLocalHost :: MonadAsync m => PortNumber -> SerialT m Socket
245connectionsOnLocalHost = connectionsOnAddr (127,0,0,1)
247-------------------------------------------------------------------------------
248-- TCP Clients
249-------------------------------------------------------------------------------
251-- | Connect to the specified IP address and port number.
252--
253-- @since 0.7.0
254connect :: (Word8, Word8, Word8, Word8) -> PortNumber -> IO Socket
255connect addr port = do
256 sock <- socket AF_INET Stream defaultProtocol
257 Net.connect sock $ SockAddrInet port (Net.tupleToHostAddress addr)
258 return sock
260-- | Connect to a remote host using IP address and port and run the supplied
261-- action on the resulting socket. 'withConnectionM' makes sure that the
262-- socket is closed on normal termination or in case of an exception. If
263-- closing the socket raises an exception, then this exception will be raised
264-- by 'withConnectionM'.
265--
266-- /Internal/
267{-# INLINABLE withConnectionM #-}
268withConnectionM :: (MonadMask m, MonadIO m)
269 => (Word8, Word8, Word8, Word8) -> PortNumber -> (Socket -> m ()) -> m ()
270withConnectionM addr port =
271 bracket (liftIO $ connect addr port) (liftIO . Net.close)
273-------------------------------------------------------------------------------
274-- Connect (unfolds)
275-------------------------------------------------------------------------------
277-- | Transform an 'Unfold' from a 'Socket' to an unfold from a remote IP
278-- address and port. The resulting unfold opens a socket, uses it using the
279-- supplied unfold and then makes sure that the socket is closed on normal
280-- termination or in case of an exception. If closing the socket raises an
281-- exception, then this exception will be raised by 'usingConnection'.
282--
283-- /Internal/
284{-# INLINABLE usingConnection #-}
285usingConnection :: (MonadCatch m, MonadIO m)
286 => Unfold m Socket a
287 -> Unfold m ((Word8, Word8, Word8, Word8), PortNumber) a
288usingConnection =
289 UF.bracket (\(addr, port) -> liftIO $ connect addr port)
290 (liftIO . Net.close)
292-------------------------------------------------------------------------------
293-- Connect (streams)
294-------------------------------------------------------------------------------
296-- | @'withConnection' addr port act@ opens a connection to the specified IPv4
297-- host address and port and passes the resulting socket handle to the
298-- computation @act@. The handle will be closed on exit from 'withConnection',
299-- whether by normal termination or by raising an exception. If closing the
300-- handle raises an exception, then this exception will be raised by
301-- 'withConnection' rather than any exception raised by 'act'.
302--
303-- /Internal/
304{-# INLINABLE withConnection #-}
305withConnection :: (IsStream t, MonadCatch m, MonadIO m)
306 => (Word8, Word8, Word8, Word8) -> PortNumber -> (Socket -> t m a) -> t m a
307withConnection addr port =
308 S.bracket (liftIO $ connect addr port) (liftIO . Net.close)
310-------------------------------------------------------------------------------
311-- Read Addr to Stream
312-------------------------------------------------------------------------------
314-- | Read a stream from the supplied IPv4 host address and port number.
315--
316-- @since 0.7.0
317{-# INLINE read #-}
318read :: (MonadCatch m, MonadIO m)
319 => Unfold m ((Word8, Word8, Word8, Word8), PortNumber) Word8
320read = UF.concat (usingConnection ISK.readChunks) A.read
322-- | Read a stream from the supplied IPv4 host address and port number.
323--
324-- @since 0.7.0
325{-# INLINE toBytes #-}
326toBytes :: (IsStream t, MonadCatch m, MonadIO m)
327 => (Word8, Word8, Word8, Word8) -> PortNumber -> t m Word8
328toBytes addr port = AS.concat $ withConnection addr port ISK.toChunks
330-------------------------------------------------------------------------------
331-- Writing
332-------------------------------------------------------------------------------
334-- | Write a stream of arrays to the supplied IPv4 host address and port
335-- number.
336--
337-- @since 0.7.0
338{-# INLINE fromChunks #-}
339fromChunks
340 :: (MonadCatch m, MonadAsync m)
341 => (Word8, Word8, Word8, Word8)
342 -> PortNumber
343 -> SerialT m (Array Word8)
344 -> m ()
345fromChunks addr port xs =
346 S.drain $ withConnection addr port (\sk -> S.yieldM $ ISK.fromChunks sk xs)
348-- | Write a stream of arrays to the supplied IPv4 host address and port
349-- number.
350--
351-- @since 0.7.0
352{-# INLINE writeChunks #-}
353writeChunks
354 :: (MonadAsync m, MonadCatch m)
355 => (Word8, Word8, Word8, Word8)
356 -> PortNumber
357 -> Fold m (Array Word8) ()
358writeChunks addr port = Fold step initial extract
359 where
360 initial = do
361 skt <- liftIO (connect addr port)
362 fld <- FL.initialize (SK.writeChunks skt)
363 `MC.onException` (liftIO $ Net.close skt)
364 return (fld, skt)
365 step (fld, skt) x = do
366 r <- FL.runStep fld x `MC.onException` (liftIO $ Net.close skt)
367 return (r, skt)
368 extract ((Fold _ initial1 extract1), skt) = do
369 liftIO $ Net.close skt
370 initial1 >>= extract1
372-- | Like 'write' but provides control over the write buffer. Output will
373-- be written to the IO device as soon as we collect the specified number of
374-- input elements.
375--
376-- @since 0.7.0
377{-# INLINE fromBytesWithBufferOf #-}
378fromBytesWithBufferOf
379 :: (MonadCatch m, MonadAsync m)
380 => Int
381 -> (Word8, Word8, Word8, Word8)
382 -> PortNumber
383 -> SerialT m Word8
384 -> m ()
385fromBytesWithBufferOf n addr port m = fromChunks addr port $ AS.arraysOf n m
387-- | Like 'write' but provides control over the write buffer. Output will
388-- be written to the IO device as soon as we collect the specified number of
389-- input elements.
390--
391-- @since 0.7.0
392{-# INLINE writeWithBufferOf #-}
393writeWithBufferOf
394 :: (MonadAsync m, MonadCatch m)
395 => Int
396 -> (Word8, Word8, Word8, Word8)
397 -> PortNumber
398 -> Fold m Word8 ()
399writeWithBufferOf n addr port =
400 FL.lchunksOf n (writeNUnsafe n) (writeChunks addr port)
402-- | Write a stream to the supplied IPv4 host address and port number.
403--
404-- @since 0.7.0
405{-# INLINE fromBytes #-}
406fromBytes :: (MonadCatch m, MonadAsync m)
407 => (Word8, Word8, Word8, Word8) -> PortNumber -> SerialT m Word8 -> m ()
408fromBytes = fromBytesWithBufferOf defaultChunkSize
410-- | Write a stream to the supplied IPv4 host address and port number.
411--
412-- @since 0.7.0
413{-# INLINE write #-}
414write :: (MonadAsync m, MonadCatch m)
415 => (Word8, Word8, Word8, Word8) -> PortNumber -> Fold m Word8 ()
416write = writeWithBufferOf defaultChunkSize
418-------------------------------------------------------------------------------
419-- Transformations
420-------------------------------------------------------------------------------
422{-# INLINABLE withInputConnect #-}
423withInputConnect
424 :: (IsStream t, MonadCatch m, MonadAsync m)
425 => (Word8, Word8, Word8, Word8)
426 -> PortNumber
427 -> SerialT m Word8
428 -> (Socket -> t m a)
429 -> t m a
430withInputConnect addr port input f = S.bracket pre post handler
432 where
434 pre = do
435 sk <- liftIO $ connect addr port
436 tid <- fork (ISK.fromBytes sk input)
437 return (sk, tid)
439 handler (sk, _) = f sk
441 -- XXX kill the thread immediately?
442 post (sk, _) = liftIO $ Net.close sk
444-- | Send an input stream to a remote host and produce the output stream from
445-- the host. The server host just acts as a transformation function on the
446-- input stream. Both sending and receiving happen asynchronously.
447--
448-- /Internal/
449--
450{-# INLINABLE transformBytesWith #-}
451transformBytesWith
452 :: (IsStream t, MonadAsync m, MonadCatch m)
453 => (Word8, Word8, Word8, Word8)
454 -> PortNumber
455 -> SerialT m Word8
456 -> t m Word8
457transformBytesWith addr port input =
458 withInputConnect addr port input ISK.toBytes