src/Streamly/Internal/Network/Inet/TCP.hs · streamly

git.iko.soy/streamly

on e01057f help sign in

streamly/src/Streamly/Internal/Network/Inet/TCP.hs

458 lines 15 kB history raw

1-- |

2-- Module : Streamly.Internal.Network.Inet.TCP

4--

5-- License : BSD3

6-- Maintainer : streamly@composewell.com

7-- Stability : experimental

8-- Portability : GHC

9--

10-- Combinators to build Inet/TCP clients and servers.

12module Streamly.Internal.Network.Inet.TCP

13 (

14 -- * TCP Servers

15 -- ** Streams

16 acceptOnAddr

17 , acceptOnAddrWith

18 , acceptOnPort

19 -- , acceptOnPortWith

20 , acceptOnPortLocal

22 -- ** Unfolds

23 , acceptorOnAddr

24 , acceptorOnAddrWith

25 , acceptorOnPort

26 , acceptorOnPortWith

27 , acceptorOnPortLocal

29 -- * TCP clients

30 -- | IP Address based operations.

31 , connect

32 , withConnectionM

34 -- ** Unfolds

35 , usingConnection

36 , reader

38 -- ** Streams

39 , withConnection

40 -- *** Source

41 , read

42 -- , readUtf8

43 -- , readLines

44 -- , readFrames

45 -- , readByChunks

47 -- -- * Array Read

48 -- , readArrayUpto

49 -- , readArrayOf

51 -- , readChunksUpto

52 -- , readChunksOf

53 -- , readChunks

55 -- *** Sink

56 , write

57 -- , writeUtf8

58 -- , writeUtf8ByLines

59 -- , writeByFrames

60 , writeWithBufferOf

61 , putBytes

62 , putBytesWithBufferOf

64 -- -- * Array Write

65 -- , writeArray

66 , writeChunks

67 , putChunks

69 -- ** Transformation

70 , pipeBytes

71 {-

72 -- ** Sink Servers

74 -- These abstractions can be applied to any setting where we need to do a

75 -- sink processing of multiple streams e.g. output from multiple processes

76 -- or data coming from multiple files.

78 -- handle connections concurrently using a specified fold

79 -- , handleConnections

81 -- handle frames concurrently using a specified fold

82 , handleFrames

84 -- merge frames from all connection into a single stream. Frames can be

85 -- created by a specified fold.

86 , mergeFrames

88 -- * UDP Servers

89 , datagrams

90 , datagramsOn

91 -}

92 )

93where

95#include "inline.hs"

97import Control.Exception (onException)

98import Control.Monad.Catch (MonadCatch, MonadMask, bracket)

99import Control.Monad.IO.Class (MonadIO(..))

100import Data.Word (Word8)

101import Network.Socket

102 (Socket, PortNumber, SocketOption(..), Family(..), SockAddr(..),

103 SocketType(..), defaultProtocol, maxListenQueue, tupleToHostAddress,

104 socket)

105import Prelude hiding (read)

106

107import Streamly.Internal.Control.Concurrent (MonadAsync)

108import Streamly.Internal.Control.ForkLifted (fork)

109import Streamly.Internal.Data.Array.Type (Array(..), writeNUnsafe)

110import Streamly.Internal.Data.Fold.Type (Fold(..))

111import Streamly.Data.Stream (Stream)

112import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))

113import Streamly.Internal.Data.Unboxed (Unbox)

114import Streamly.Internal.Data.Unfold.Type (Unfold(..))

115import Streamly.Internal.Network.Socket (SockSpec(..), accept, acceptor)

116import Streamly.Internal.System.IO (defaultChunkSize)

117

118import qualified Control.Monad.Catch as MC

119import qualified Network.Socket as Net

120

121import qualified Streamly.Data.Array as A

122import qualified Streamly.Data.Fold as FL

123import qualified Streamly.Data.Unfold as UF

124import qualified Streamly.Internal.Data.Array.Type as A

125import qualified Streamly.Internal.Data.Unfold as UF (first, bracketIO)

126import qualified Streamly.Internal.Data.Fold.Type as FL (Step(..), reduce)

127import qualified Streamly.Internal.Data.Stream.StreamD as S

128import qualified Streamly.Internal.Data.Stream.Exception.Lifted as S

129import qualified Streamly.Internal.Network.Socket as ISK

130

131-------------------------------------------------------------------------------

132-- Accept (unfolds)

133-------------------------------------------------------------------------------

134

135{-# INLINE acceptorOnAddrWith #-}

136acceptorOnAddrWith

137 :: MonadIO m

138 => [(SocketOption, Int)]

139 -> Unfold m ((Word8, Word8, Word8, Word8), PortNumber) Socket

140acceptorOnAddrWith opts = UF.lmap f acceptor

141 where

142 f (addr, port) =

143 (maxListenQueue

144 , SockSpec

145 { sockFamily = AF_INET

146 , sockType = Stream

147 , sockProto = defaultProtocol -- TCP

148 , sockOpts = opts

149 }

150 , SockAddrInet port (tupleToHostAddress addr)

151 )

152

153-- | Unfold a tuple @(ipAddr, port)@ into a stream of connected TCP sockets.

154-- @ipAddr@ is the local IP address and @port@ is the local port on which

155-- connections are accepted.

156--

157{-# INLINE acceptorOnAddr #-}

158acceptorOnAddr

159 :: MonadIO m

160 => Unfold m ((Word8, Word8, Word8, Word8), PortNumber) Socket

161acceptorOnAddr = acceptorOnAddrWith []

162

163{-# INLINE acceptorOnPortWith #-}

164acceptorOnPortWith :: MonadIO m

165 => [(SocketOption, Int)]

166 -> Unfold m PortNumber Socket

167acceptorOnPortWith opts = UF.first (0,0,0,0) (acceptorOnAddrWith opts)

168

169-- | Like 'acceptorOnAddr' but binds on the IPv4 address @0.0.0.0@ i.e. on all

170-- IPv4 addresses/interfaces of the machine and listens for TCP connections on

171-- the specified port.

172--

173-- > acceptorOnPort = UF.first acceptorOnAddr (0,0,0,0)

174--

175{-# INLINE acceptorOnPort #-}

176acceptorOnPort :: MonadIO m => Unfold m PortNumber Socket

177acceptorOnPort = UF.first (0,0,0,0) acceptorOnAddr

178

179-- | Like 'acceptorOnAddr' but binds on the localhost IPv4 address @127.0.0.1@.

180-- The server can only be accessed from the local host, it cannot be accessed

181-- from other hosts on the network.

182--

183-- > acceptorOnPortLocal = UF.first acceptorOnAddr (127,0,0,1)

184--

185{-# INLINE acceptorOnPortLocal #-}

186acceptorOnPortLocal :: MonadIO m => Unfold m PortNumber Socket

187acceptorOnPortLocal = UF.first (127,0,0,1) acceptorOnAddr

188

189-------------------------------------------------------------------------------

190-- Accept (streams)

191-------------------------------------------------------------------------------

192

193{-# INLINE acceptOnAddrWith #-}

194acceptOnAddrWith

195 :: MonadIO m

196 => [(SocketOption, Int)]

197 -> (Word8, Word8, Word8, Word8)

198 -> PortNumber

199 -> Stream m Socket

200acceptOnAddrWith opts addr port =

201 accept maxListenQueue SockSpec

202 { sockFamily = AF_INET

203 , sockType = Stream

204 , sockProto = defaultProtocol

205 , sockOpts = opts

206 }

207 (SockAddrInet port (tupleToHostAddress addr))

208

209-- | Like 'accept' but binds on the specified IPv4 address of the machine

210-- and listens for TCP connections on the specified port.

211--

212-- /Pre-release/

213{-# INLINE acceptOnAddr #-}

214acceptOnAddr

215 :: MonadIO m

216 => (Word8, Word8, Word8, Word8)

217 -> PortNumber

218 -> Stream m Socket

219acceptOnAddr = acceptOnAddrWith []

220

221-- | Like 'accept' but binds on the IPv4 address @0.0.0.0@ i.e. on all

222-- IPv4 addresses/interfaces of the machine and listens for TCP connections on

223-- the specified port.

224--

225-- > acceptOnPort = acceptOnAddr (0,0,0,0)

226--

227-- /Pre-release/

228{-# INLINE acceptOnPort #-}

229acceptOnPort :: MonadIO m => PortNumber -> Stream m Socket

230acceptOnPort = acceptOnAddr (0,0,0,0)

231

232-- | Like 'accept' but binds on the localhost IPv4 address @127.0.0.1@.

233-- The server can only be accessed from the local host, it cannot be accessed

234-- from other hosts on the network.

235--

236-- > acceptOnPortLocal = acceptOnAddr (127,0,0,1)

237--

238-- /Pre-release/

239{-# INLINE acceptOnPortLocal #-}

240acceptOnPortLocal :: MonadIO m => PortNumber -> Stream m Socket

241acceptOnPortLocal = acceptOnAddr (127,0,0,1)

242

243-------------------------------------------------------------------------------

244-- TCP Clients

245-------------------------------------------------------------------------------

246

247-- | Connect to the specified IP address and port number. Returns a connected

248-- socket or throws an exception.

249--

250connect :: (Word8, Word8, Word8, Word8) -> PortNumber -> IO Socket

251connect addr port = do

252 sock <- socket AF_INET Stream defaultProtocol

253 Net.connect sock (SockAddrInet port (Net.tupleToHostAddress addr))

254 `onException` Net.close sock

255 return sock

256

257-- | Connect to a remote host using IP address and port and run the supplied

258-- action on the resulting socket. 'withConnectionM' makes sure that the

259-- socket is closed on normal termination or in case of an exception. If

260-- closing the socket raises an exception, then this exception will be raised

261-- by 'withConnectionM'.

262--

263-- /Pre-release/

264{-# INLINABLE withConnectionM #-}

265withConnectionM :: (MonadMask m, MonadIO m)

266 => (Word8, Word8, Word8, Word8) -> PortNumber -> (Socket -> m ()) -> m ()

267withConnectionM addr port =

268 bracket (liftIO $ connect addr port) (liftIO . Net.close)

269

270-------------------------------------------------------------------------------

271-- Connect (unfolds)

272-------------------------------------------------------------------------------

273

274-- | Transform an 'Unfold' from a 'Socket' to an unfold from a remote IP

275-- address and port. The resulting unfold opens a socket, uses it using the

276-- supplied unfold and then makes sure that the socket is closed on normal

277-- termination or in case of an exception. If closing the socket raises an

278-- exception, then this exception will be raised by 'usingConnection'.

279--

280-- /Pre-release/

281{-# INLINE usingConnection #-}

282usingConnection :: (MonadCatch m, MonadAsync m)

283 => Unfold m Socket a

284 -> Unfold m ((Word8, Word8, Word8, Word8), PortNumber) a

285usingConnection = UF.bracketIO (\(addr, port) -> connect addr port) Net.close

286

287-------------------------------------------------------------------------------

288-- Connect (streams)

289-------------------------------------------------------------------------------

290

291-- | @'withConnection' addr port act@ opens a connection to the specified IPv4

292-- host address and port and passes the resulting socket handle to the

293-- computation @act@. The handle will be closed on exit from 'withConnection',

294-- whether by normal termination or by raising an exception. If closing the

295-- handle raises an exception, then this exception will be raised by

296-- 'withConnection' rather than any exception raised by 'act'.

297--

298-- /Pre-release/

299{-# INLINE withConnection #-}

300withConnection :: (MonadCatch m, MonadAsync m)

301 => (Word8, Word8, Word8, Word8)

302 -> PortNumber

303 -> (Socket -> Stream m a)

304 -> Stream m a

305withConnection addr port = S.bracketIO (connect addr port) Net.close

306

307-------------------------------------------------------------------------------

308-- Read Addr to Stream

309-------------------------------------------------------------------------------

310

311-- | Read a stream from the supplied IPv4 host address and port number.

312--

313{-# INLINE reader #-}

314reader :: (MonadCatch m, MonadAsync m)

315 => Unfold m ((Word8, Word8, Word8, Word8), PortNumber) Word8

316reader = UF.many A.reader (usingConnection ISK.chunkReader)

317

318{-# INLINE concatChunks #-}

319concatChunks :: (Monad m, Unbox a) => Stream m (Array a) -> Stream m a

320concatChunks = S.unfoldMany A.reader

321

322-- | Read a stream from the supplied IPv4 host address and port number.

323--

324-- /Pre-release/

325{-# INLINE read #-}

326read :: (MonadCatch m, MonadAsync m)

327 => (Word8, Word8, Word8, Word8) -> PortNumber -> Stream m Word8

328read addr port = concatChunks $ withConnection addr port ISK.readChunks

329

330-------------------------------------------------------------------------------

331-- Writing

332-------------------------------------------------------------------------------

333

334-- | Write a stream of arrays to the supplied IPv4 host address and port

335-- number.

336--

337-- /Pre-release/

338{-# INLINE putChunks #-}

339putChunks

340 :: (MonadCatch m, MonadAsync m)

341 => (Word8, Word8, Word8, Word8)

342 -> PortNumber

343 -> Stream m (Array Word8)

344 -> m ()

345putChunks addr port xs =

346 S.fold FL.drain

347 $ withConnection addr port (\sk -> S.fromEffect $ ISK.putChunks sk xs)

348

349-- | Write a stream of arrays to the supplied IPv4 host address and port

350-- number.

351--

352{-# INLINE writeChunks #-}

353writeChunks

354 :: (MonadIO 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.reduce (ISK.writeChunks skt)

363 `MC.onException` liftIO (Net.close skt)

364 return $ FL.Partial (Tuple' fld skt)

365 step (Tuple' fld skt) x = do

366 r <- FL.addOne x fld `MC.onException` liftIO (Net.close skt)

367 return $ FL.Partial (Tuple' r skt)

368 extract (Tuple' (Fold _ initial1 extract1) skt) = do

369 liftIO $ Net.close skt

370 res <- initial1

371 case res of

372 FL.Partial fs -> extract1 fs

373 FL.Done fb -> return fb

374

375-- | Like 'write' but provides control over the write buffer. Output will

376-- be written to the IO device as soon as we collect the specified number of

377-- input elements.

378--

379-- /Pre-release/

380{-# INLINE putBytesWithBufferOf #-}

381putBytesWithBufferOf

382 :: (MonadCatch m, MonadAsync m)

383 => Int

384 -> (Word8, Word8, Word8, Word8)

385 -> PortNumber

386 -> Stream m Word8

387 -> m ()

388putBytesWithBufferOf n addr port m = putChunks addr port $ A.chunksOf n m

389

390-- | Like 'write' but provides control over the write buffer. Output will

391-- be written to the IO device as soon as we collect the specified number of

392-- input elements.

393--

394{-# INLINE writeWithBufferOf #-}

395writeWithBufferOf

396 :: (MonadIO m, MonadCatch m)

397 => Int

398 -> (Word8, Word8, Word8, Word8)

399 -> PortNumber

400 -> Fold m Word8 ()

401writeWithBufferOf n addr port =

402 FL.groupsOf n (writeNUnsafe n) (writeChunks addr port)

403

404-- | Write a stream to the supplied IPv4 host address and port number.

405--

406-- /Pre-release/

407{-# INLINE putBytes #-}

408putBytes :: (MonadCatch m, MonadAsync m)

409 => (Word8, Word8, Word8, Word8) -> PortNumber -> Stream m Word8 -> m ()

410putBytes = putBytesWithBufferOf defaultChunkSize

411

412-- | Write a stream to the supplied IPv4 host address and port number.

413--

414{-# INLINE write #-}

415write :: (MonadIO m, MonadCatch m)

416 => (Word8, Word8, Word8, Word8) -> PortNumber -> Fold m Word8 ()

417write = writeWithBufferOf defaultChunkSize

418

419-------------------------------------------------------------------------------

420-- Transformations

421-------------------------------------------------------------------------------

422

423{-# INLINE withInputConnect #-}

424withInputConnect

425 :: (MonadCatch m, MonadAsync m)

426 => (Word8, Word8, Word8, Word8)

427 -> PortNumber

428 -> Stream m Word8

429 -> (Socket -> Stream m a)

430 -> Stream m a

431withInputConnect addr port input f = S.bracket pre post handler

432

433 where

434

435 pre = do

436 sk <- liftIO $ connect addr port

437 tid <- fork (ISK.putBytes sk input)

438 return (sk, tid)

439

440 handler (sk, _) = f sk

441

442 -- XXX kill the thread immediately?

443 post (sk, _) = liftIO $ Net.close sk

444

445-- | Send an input stream to a remote host and produce the output stream from

446-- the host. The server host just acts as a transformation function on the

447-- input stream. Both sending and receiving happen asynchronously.

448--

449-- /Pre-release/

450--

451{-# INLINE pipeBytes #-}

452pipeBytes

453 :: (MonadAsync m, MonadCatch m)

454 => (Word8, Word8, Word8, Word8)

455 -> PortNumber

456 -> Stream m Word8

457 -> Stream m Word8

458pipeBytes addr port input = withInputConnect addr port input ISK.read