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

git.iko.soy/streamly

on 79df417 help sign in

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

462 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.Data.Array (Array)

110import Streamly.Internal.Data.Array (pinnedWriteNUnsafe)

111import Streamly.Internal.Data.Fold ( Fold(..) )

112import Streamly.Data.Stream (Stream)

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

114import Streamly.Internal.Data.Unbox (Unbox)

115import Streamly.Data.Unfold (Unfold)

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

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

118

119import qualified Control.Monad.Catch as MC

120import qualified Network.Socket as Net

121

122import qualified Streamly.Data.Array as A

123import qualified Streamly.Data.Fold as FL

124import qualified Streamly.Data.Stream as S

125import qualified Streamly.Data.Unfold as UF

126import qualified Streamly.Internal.Data.Array as A (pinnedChunksOf)

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

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

129

130import qualified Streamly.Internal.Data.Stream.Exception.Lifted as S (bracket)

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

132 (chunkReader, putBytes, putChunks, read, readChunks, writeChunks)

133

134-------------------------------------------------------------------------------

135-- Accept (unfolds)

136-------------------------------------------------------------------------------

137

138{-# INLINE acceptorOnAddrWith #-}

139acceptorOnAddrWith

140 :: MonadIO m

141 => [(SocketOption, Int)]

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

143acceptorOnAddrWith opts = UF.lmap f acceptor

144 where

145 f (addr, port) =

146 (maxListenQueue

147 , SockSpec

148 { sockFamily = AF_INET

149 , sockType = Stream

150 , sockProto = defaultProtocol -- TCP

151 , sockOpts = opts

152 }

153 , SockAddrInet port (tupleToHostAddress addr)

154 )

155

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

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

158-- connections are accepted.

159--

160{-# INLINE acceptorOnAddr #-}

161acceptorOnAddr

162 :: MonadIO m

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

164acceptorOnAddr = acceptorOnAddrWith []

165

166{-# INLINE acceptorOnPortWith #-}

167acceptorOnPortWith :: MonadIO m

168 => [(SocketOption, Int)]

169 -> Unfold m PortNumber Socket

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

171

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

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

174-- the specified port.

175--

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

177--

178{-# INLINE acceptorOnPort #-}

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

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

181

182-- | Like 'acceptorOnAddr' 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-- > acceptorOnPortLocal = UF.first acceptorOnAddr (127,0,0,1)

187--

188{-# INLINE acceptorOnPortLocal #-}

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

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

191

192-------------------------------------------------------------------------------

193-- Accept (streams)

194-------------------------------------------------------------------------------

195

196{-# INLINE acceptOnAddrWith #-}

197acceptOnAddrWith

198 :: MonadIO m

199 => [(SocketOption, Int)]

200 -> (Word8, Word8, Word8, Word8)

201 -> PortNumber

202 -> Stream m Socket

203acceptOnAddrWith opts addr port =

204 accept maxListenQueue SockSpec

205 { sockFamily = AF_INET

206 , sockType = Stream

207 , sockProto = defaultProtocol

208 , sockOpts = opts

209 }

210 (SockAddrInet port (tupleToHostAddress addr))

211

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

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

214--

215-- /Pre-release/

216{-# INLINE acceptOnAddr #-}

217acceptOnAddr

218 :: MonadIO m

219 => (Word8, Word8, Word8, Word8)

220 -> PortNumber

221 -> Stream m Socket

222acceptOnAddr = acceptOnAddrWith []

223

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

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

226-- the specified port.

227--

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

229--

230-- /Pre-release/

231{-# INLINE acceptOnPort #-}

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

233acceptOnPort = acceptOnAddr (0,0,0,0)

234

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

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

237-- from other hosts on the network.

238--

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

240--

241-- /Pre-release/

242{-# INLINE acceptOnPortLocal #-}

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

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

245

246-------------------------------------------------------------------------------

247-- TCP Clients

248-------------------------------------------------------------------------------

249

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

251-- socket or throws an exception.

252--

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

254connect addr port = do

255 sock <- socket AF_INET Stream defaultProtocol

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

257 `onException` Net.close sock

258 return sock

259

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-- /Pre-release/

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)

272

273-------------------------------------------------------------------------------

274-- Connect (unfolds)

275-------------------------------------------------------------------------------

276

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-- /Pre-release/

284{-# INLINE usingConnection #-}

285usingConnection :: (MonadCatch m, MonadAsync m)

286 => Unfold m Socket a

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

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

289

290-------------------------------------------------------------------------------

291-- Connect (streams)

292-------------------------------------------------------------------------------

293

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

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

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

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

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

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

300--

301-- /Pre-release/

302{-# INLINE withConnection #-}

303withConnection :: (MonadCatch m, MonadAsync m)

304 => (Word8, Word8, Word8, Word8)

305 -> PortNumber

306 -> (Socket -> Stream m a)

307 -> Stream m a

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

309

310-------------------------------------------------------------------------------

311-- Read Addr to Stream

312-------------------------------------------------------------------------------

313

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

315--

316{-# INLINE reader #-}

317reader :: (MonadCatch m, MonadAsync m)

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

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

320

321{-# INLINE concatChunks #-}

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

323concatChunks = S.unfoldMany A.reader

324

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

326--

327-- /Pre-release/

328{-# INLINE read #-}

329read :: (MonadCatch m, MonadAsync m)

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

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

332

333-------------------------------------------------------------------------------

334-- Writing

335-------------------------------------------------------------------------------

336

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

338-- number.

339--

340-- /Pre-release/

341{-# INLINE putChunks #-}

342putChunks

343 :: (MonadCatch m, MonadAsync m)

344 => (Word8, Word8, Word8, Word8)

345 -> PortNumber

346 -> Stream m (Array Word8)

347 -> m ()

348putChunks addr port xs =

349 S.fold FL.drain

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

351

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

353-- number.

354--

355{-# INLINE writeChunks #-}

356writeChunks

357 :: (MonadIO m, MonadCatch m)

358 => (Word8, Word8, Word8, Word8)

359 -> PortNumber

360 -> Fold m (Array Word8) ()

361writeChunks addr port = Fold step initial extract

362 where

363 initial = do

364 skt <- liftIO (connect addr port)

365 fld <- FL.reduce (ISK.writeChunks skt)

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

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

368 step (Tuple' fld skt) x = do

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

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

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

372 liftIO $ Net.close skt

373 res <- initial1

374 case res of

375 FL.Partial fs -> extract1 fs

376 FL.Done fb -> return fb

377

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

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

380-- input elements.

381--

382-- /Pre-release/

383{-# INLINE putBytesWithBufferOf #-}

384putBytesWithBufferOf

385 :: (MonadCatch m, MonadAsync m)

386 => Int

387 -> (Word8, Word8, Word8, Word8)

388 -> PortNumber

389 -> Stream m Word8

390 -> m ()

391putBytesWithBufferOf n addr port m =

392 putChunks addr port $ A.pinnedChunksOf n m

393

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

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

396-- input elements.

397--

398{-# INLINE writeWithBufferOf #-}

399writeWithBufferOf

400 :: (MonadIO m, MonadCatch m)

401 => Int

402 -> (Word8, Word8, Word8, Word8)

403 -> PortNumber

404 -> Fold m Word8 ()

405writeWithBufferOf n addr port =

406 FL.groupsOf n (pinnedWriteNUnsafe n) (writeChunks addr port)

407

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

409--

410-- /Pre-release/

411{-# INLINE putBytes #-}

412putBytes :: (MonadCatch m, MonadAsync m)

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

414putBytes = putBytesWithBufferOf defaultChunkSize

415

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

417--

418{-# INLINE write #-}

419write :: (MonadIO m, MonadCatch m)

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

421write = writeWithBufferOf defaultChunkSize

422

423-------------------------------------------------------------------------------

424-- Transformations

425-------------------------------------------------------------------------------

426

427{-# INLINE withInputConnect #-}

428withInputConnect

429 :: (MonadCatch m, MonadAsync m)

430 => (Word8, Word8, Word8, Word8)

431 -> PortNumber

432 -> Stream m Word8

433 -> (Socket -> Stream m a)

434 -> Stream m a

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

436

437 where

438

439 pre = do

440 sk <- liftIO $ connect addr port

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

442 return (sk, tid)

443

444 handler (sk, _) = f sk

445

446 -- XXX kill the thread immediately?

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

448

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

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

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

452--

453-- /Pre-release/

454--

455{-# INLINE pipeBytes #-}

456pipeBytes

457 :: (MonadAsync m, MonadCatch m)

458 => (Word8, Word8, Word8, Word8)

459 -> PortNumber

460 -> Stream m Word8

461 -> Stream m Word8

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