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

git.iko.soy/streamly

on 604f430 help sign in

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

456 lines 15 kB history raw

1#include "inline.hs"

3-- |

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

6--

7-- License : BSD3

8-- Maintainer : streamly@composewell.com

9-- Stability : experimental

10-- Portability : GHC

11--

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

14module Streamly.Internal.Network.Inet.TCP

15 (

16 -- * TCP Servers

17 -- ** Unfolds

18 acceptOnAddr

19 , acceptOnAddrWith

20 , acceptOnPort

21 , acceptOnPortWith

22 , acceptOnPortLocal

24 -- ** Streams

25 , connectionsOnAddr

26 , connectionsOnAddrWith

27 , connectionsOnPort

28 , connectionsOnLocalHost

30 -- * TCP clients

31 -- | IP Address based operations.

32 , connect

33 , withConnectionM

35 -- ** Unfolds

36 , usingConnection

37 , read

39 -- ** Streams

40 , withConnection

41 -- *** Source

42 , toBytes

43 -- , readUtf8

44 -- , readLines

45 -- , readFrames

46 -- , readByChunks

48 -- -- * Array Read

49 -- , readArrayUpto

50 -- , readArrayOf

52 -- , readChunksUpto

53 -- , readChunksOf

54 -- , readChunks

56 -- *** Sink

57 , write

58 -- , writeUtf8

59 -- , writeUtf8ByLines

60 -- , writeByFrames

61 , writeWithBufferOf

62 , putBytes

63 , putBytesWithBufferOf

65 -- -- * Array Write

66 -- , writeArray

67 , writeChunks

68 , putChunks

70 -- ** Transformation

71 , processBytes

72 {-

73 -- ** Sink Servers

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

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

77 -- or data coming from multiple files.

79 -- handle connections concurrently using a specified fold

80 -- , handleConnections

82 -- handle frames concurrently using a specified fold

83 , handleFrames

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

86 -- created by a specified fold.

87 , mergeFrames

89 -- * UDP Servers

90 , datagrams

91 , datagramsOn

92 -}

93 )

94where

96import Control.Exception (onException)

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

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

99import Data.Word (Word8)

100import Network.Socket

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

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

103 socket)

104import Prelude hiding (read)

105

106import Streamly.Internal.Data.Array.Foreign.Type (Array(..), defaultChunkSize, writeNUnsafe)

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

108import Streamly.Internal.Data.Stream.IsStream (MonadAsync)

109import Streamly.Internal.Data.SVar (fork)

110import Streamly.Internal.Data.Stream.Serial (SerialT)

111import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)

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

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

114import Streamly.Internal.Network.Socket (SockSpec(..), accept, connections)

115

116import qualified Control.Monad.Catch as MC

117import qualified Network.Socket as Net

118

119import qualified Streamly.Internal.Data.Unfold as UF

120import qualified Streamly.Internal.Data.Array.Foreign as A

121import qualified Streamly.Internal.Data.Array.Stream.Foreign as AS

122import qualified Streamly.Internal.Data.Fold.Type as FL

123import qualified Streamly.Internal.Data.Stream.IsStream as S

124import qualified Streamly.Network.Socket as SK

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

126

127-------------------------------------------------------------------------------

128-- Accept (unfolds)

129-------------------------------------------------------------------------------

130

131{-# INLINE acceptOnAddrWith #-}

132acceptOnAddrWith

133 :: MonadIO m

134 => [(SocketOption, Int)]

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

136acceptOnAddrWith opts = UF.lmap f accept

137 where

138 f (addr, port) =

139 (maxListenQueue

140 , SockSpec

141 { sockFamily = AF_INET

142 , sockType = Stream

143 , sockProto = defaultProtocol -- TCP

144 , sockOpts = opts

145 }

146 , SockAddrInet port (tupleToHostAddress addr)

147 )

148

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

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

151-- connections are accepted.

152--

153-- @since 0.7.0

154{-# INLINE acceptOnAddr #-}

155acceptOnAddr

156 :: MonadIO m

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

158acceptOnAddr = acceptOnAddrWith []

159

160{-# INLINE acceptOnPortWith #-}

161acceptOnPortWith :: MonadIO m

162 => [(SocketOption, Int)]

163 -> Unfold m PortNumber Socket

164acceptOnPortWith opts = UF.supplyFirst (0,0,0,0) (acceptOnAddrWith opts)

165

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

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

168-- the specified port.

169--

170-- > acceptOnPort = UF.supplyFirst acceptOnAddr (0,0,0,0)

171--

172-- @since 0.7.0

173{-# INLINE acceptOnPort #-}

174acceptOnPort :: MonadIO m => Unfold m PortNumber Socket

175acceptOnPort = UF.supplyFirst (0,0,0,0) acceptOnAddr

176

177-- | Like 'acceptOnAddr' but binds on the localhost IPv4 address @127.0.0.1@.

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

179-- from other hosts on the network.

180--

181-- > acceptOnPortLocal = UF.supplyFirst acceptOnAddr (127,0,0,1)

182--

183-- @since 0.7.0

184{-# INLINE acceptOnPortLocal #-}

185acceptOnPortLocal :: MonadIO m => Unfold m PortNumber Socket

186acceptOnPortLocal = UF.supplyFirst (127,0,0,1) acceptOnAddr

187

188-------------------------------------------------------------------------------

189-- Accept (streams)

190-------------------------------------------------------------------------------

191

192{-# INLINE connectionsOnAddrWith #-}

193connectionsOnAddrWith

194 :: MonadAsync m

195 => [(SocketOption, Int)]

196 -> (Word8, Word8, Word8, Word8)

197 -> PortNumber

198 -> SerialT m Socket

199connectionsOnAddrWith opts addr port =

200 connections maxListenQueue SockSpec

201 { sockFamily = AF_INET

202 , sockType = Stream

203 , sockProto = defaultProtocol

204 , sockOpts = opts

205 }

206 (SockAddrInet port (tupleToHostAddress addr))

207

208-- | Like 'connections' but binds on the specified IPv4 address of the machine

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

210--

211-- /Pre-release/

212{-# INLINE connectionsOnAddr #-}

213connectionsOnAddr

214 :: MonadAsync m

215 => (Word8, Word8, Word8, Word8)

216 -> PortNumber

217 -> SerialT m Socket

218connectionsOnAddr = connectionsOnAddrWith []

219

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

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

222-- the specified port.

223--

224-- > connectionsOnPort = connectionsOnAddr (0,0,0,0)

225--

226-- /Pre-release/

227{-# INLINE connectionsOnPort #-}

228connectionsOnPort :: MonadAsync m => PortNumber -> SerialT m Socket

229connectionsOnPort = connectionsOnAddr (0,0,0,0)

230

231-- | Like 'connections' but binds on the localhost IPv4 address @127.0.0.1@.

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

233-- from other hosts on the network.

234--

235-- > connectionsOnLocalHost = connectionsOnAddr (127,0,0,1)

236--

237-- /Pre-release/

238{-# INLINE connectionsOnLocalHost #-}

239connectionsOnLocalHost :: MonadAsync m => PortNumber -> SerialT m Socket

240connectionsOnLocalHost = connectionsOnAddr (127,0,0,1)

241

242-------------------------------------------------------------------------------

243-- TCP Clients

244-------------------------------------------------------------------------------

245

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

247-- socket or throws an exception.

248--

249-- @since 0.7.0

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{-# INLINABLE usingConnection #-}

282usingConnection :: (MonadCatch m, MonadAsync m)

283 => Unfold m Socket a

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

285usingConnection =

286 UF.bracket (\(addr, port) -> liftIO $ connect addr port)

287 (liftIO . Net.close)

288

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

290-- Connect (streams)

291-------------------------------------------------------------------------------

292

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

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

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

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

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

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

299--

300-- /Pre-release/

301{-# INLINABLE withConnection #-}

302withConnection :: (IsStream t, MonadCatch m, MonadAsync m)

303 => (Word8, Word8, Word8, Word8) -> PortNumber -> (Socket -> t m a) -> t m a

304withConnection addr port =

305 S.bracket (liftIO $ connect addr port) (liftIO . 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-- @since 0.7.0

314{-# INLINE read #-}

315read :: (MonadCatch m, MonadAsync m)

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

317read = UF.many (usingConnection ISK.readChunks) A.read

318

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

320--

321-- @since 0.7.0

322{-# INLINE toBytes #-}

323toBytes :: (IsStream t, MonadCatch m, MonadAsync m)

324 => (Word8, Word8, Word8, Word8) -> PortNumber -> t m Word8

325toBytes addr port = AS.concat $ withConnection addr port ISK.toChunks

326

327-------------------------------------------------------------------------------

328-- Writing

329-------------------------------------------------------------------------------

330

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

332-- number.

333--

334-- @since 0.7.0

335{-# INLINE putChunks #-}

336putChunks

337 :: (MonadCatch m, MonadAsync m)

338 => (Word8, Word8, Word8, Word8)

339 -> PortNumber

340 -> SerialT m (Array Word8)

341 -> m ()

342putChunks addr port xs =

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

344

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

346-- number.

347--

348-- @since 0.7.0

349{-# INLINE writeChunks #-}

350writeChunks

351 :: (MonadAsync m, MonadCatch m)

352 => (Word8, Word8, Word8, Word8)

353 -> PortNumber

354 -> Fold m (Array Word8) ()

355writeChunks addr port = Fold step initial extract

356 where

357 initial = do

358 skt <- liftIO (connect addr port)

359 fld <- FL.initialize (SK.writeChunks skt) `MC.onException` liftIO (Net.close skt)

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

361 step (Tuple' fld skt) x = do

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

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

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

365 liftIO $ Net.close skt

366 res <- initial1

367 case res of

368 FL.Partial fs -> extract1 fs

369 FL.Done fb -> return fb

370

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

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

373-- input elements.

374--

375-- @since 0.7.0

376{-# INLINE putBytesWithBufferOf #-}

377putBytesWithBufferOf

378 :: (MonadCatch m, MonadAsync m)

379 => Int

380 -> (Word8, Word8, Word8, Word8)

381 -> PortNumber

382 -> SerialT m Word8

383 -> m ()

384putBytesWithBufferOf n addr port m = putChunks addr port $ AS.arraysOf n m

385

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

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

388-- input elements.

389--

390-- @since 0.7.0

391{-# INLINE writeWithBufferOf #-}

392writeWithBufferOf

393 :: (MonadAsync m, MonadCatch m)

394 => Int

395 -> (Word8, Word8, Word8, Word8)

396 -> PortNumber

397 -> Fold m Word8 ()

398writeWithBufferOf n addr port =

399 FL.chunksOf n (writeNUnsafe n) (writeChunks addr port)

400

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

402--

403-- @since 0.7.0

404{-# INLINE putBytes #-}

405putBytes :: (MonadCatch m, MonadAsync m)

406 => (Word8, Word8, Word8, Word8) -> PortNumber -> SerialT m Word8 -> m ()

407putBytes = putBytesWithBufferOf defaultChunkSize

408

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

410--

411-- @since 0.7.0

412{-# INLINE write #-}

413write :: (MonadAsync m, MonadCatch m)

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

415write = writeWithBufferOf defaultChunkSize

416

417-------------------------------------------------------------------------------

418-- Transformations

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

420

421{-# INLINABLE withInputConnect #-}

422withInputConnect

423 :: (IsStream t, MonadCatch m, MonadAsync m)

424 => (Word8, Word8, Word8, Word8)

425 -> PortNumber

426 -> SerialT m Word8

427 -> (Socket -> t m a)

428 -> t m a

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

430

431 where

432

433 pre = do

434 sk <- liftIO $ connect addr port

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

436 return (sk, tid)

437

438 handler (sk, _) = f sk

439

440 -- XXX kill the thread immediately?

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

442

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

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

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

446--

447-- /Pre-release/

448--

449{-# INLINABLE processBytes #-}

450processBytes

451 :: (IsStream t, MonadAsync m, MonadCatch m)

452 => (Word8, Word8, Word8, Word8)

453 -> PortNumber

454 -> SerialT m Word8

455 -> t m Word8

456processBytes addr port input = withInputConnect addr port input ISK.toBytes