Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1""" 

2UNet module provide a asyncio Transport Layer to comunicate peers. 

3 

4This module is mainly intended to send pub/sub messages to remote peers. 

5 

6TODO: encapsulate pub/sub messages (key, payload) 

7 

8""" 

9import asyncio 

10import socket 

11import random 

12import re 

13import struct 

14import ujson as json 

15import pickle 

16import base64 

17import logging 

18import zlib 

19import uuid 

20 

21from asyncio.protocols import DatagramProtocol 

22 

23# own libraries 

24from gutools.tools import add_to_loop, new_uid, parse_uri, serializable_container, rebuild 

25from gutools.uagents import UAgent, Domestic, domestic 

26from gutools.uobjects import UObject, SmartContainer, Answer 

27 

28 

29from loganalyzer.loganalyzer import setup_config 

30 

31 

32class Request(object): 

33 """Encapsulate a request with related response.""" 

34 def __init__(self, uri): 

35 self.uri = uri 

36 self.uid = new_uid() 

37 self.response = Response() 

38 

39class Response(Answer): 

40 """Encapsulate a response. 

41 Inherits from Answer.""" 

42 pass 

43 

44class Session(object): 

45 "TBD" 

46 pass 

47 

48class UProtocolFactory(): 

49 """Protocol Factory used by TL.""" 

50 def __init__(self, uid, link, tl): 

51 self.uid = uid 

52 self.link = link 

53 self.tl = tl 

54 self.types = { 

55 'broadcast': UProtocol, 

56 'direct': UProtocol, 

57 } 

58 

59 def __call__(self): 

60 # TODO: implemente a real factory here based on config or (hash) uid 

61 proto_klass = self.types.get(self.link) 

62 return proto_klass(self.uid, self.tl) 

63 

64 

65class UProtocol(DatagramProtocol, Domestic): 

66 """ 

67 DatagramProtocol to comunicate with a peer. 

68 

69 It will keep input and output queues in order to provide a 

70 reliable *stream* of datagrams. 

71 

72 User can use direct functions like `send(key, data, to, address)` 

73 or publishing in the same realm with key='/net/publish' 

74 

75 - `uid`: the uid of the peer-connection. 

76 - `realm`: the realm to attend pub/sub. 

77 

78 Special keys: 

79 

80 - `MISSING`: used for requesting missing datagrams 

81 - `WATERMARK`: update peer watermark and drop packets received by peer. 

82 

83 (sender, to, payload) 

84 """ 

85 def __init__(self, uid, tl): 

86 DatagramProtocol.__init__(self) 

87 Domestic.__init__(self) 

88 self.uid = uid 

89 self.tl = tl 

90 

91 self._output = dict() 

92 self._input = dict() 

93 self._out_sequence = dict() 

94 self._in_sequence = dict() 

95 

96 self._oob_handlers = { 

97 'MISSING': self._resend_missing, 

98 'WATERMARK': self._update_watermark, 

99 } 

100 

101 def connection_made(self, transport): 

102 self.transport = transport 

103 addr = self.transport.get_extra_info('socket').getsockname() 

104 logging.debug("connection_made: {}: {}".format(self.uid, addr)) 

105 

106 def connection_lost(self, exc): 

107 logging.debug(f'connection_lost: {exc}') 

108 

109 def datagram_received(self, raw, addr): 

110 baddr = self.transport.get_extra_info('socket').getsockname() 

111 raw = zlib.decompress(raw) 

112 sender, to, sequence, payload = json.decode(raw.decode()) 

113 if sender == self.uid: 

114 logging.debug(f'>> [{self.uid}] Skip <{payload}> self-loop-message from {addr} in {baddr}') 

115 return 

116 

117 if to in (self.uid, ): # direct message (broadcast or direct socket) 

118 logging.debug(f">> [{self.uid}] Received <{payload}> from {sender}:{addr} to '{to}' in {baddr}") 

119 

120 in_sequence = self._in_sequence.get(sender, 0) 

121 if sequence < in_sequence: 

122 return # already dispatched message (is an echo) 

123 

124 inputs = self._input.get(sender) 

125 if inputs is None: 

126 inputs = self._input[sender] = dict() 

127 inputs[sequence] = (payload, addr) 

128 

129 while in_sequence in inputs: 

130 (key, data), addr = inputs.pop(in_sequence) 

131 in_sequence += 1 

132 # TODO: include sender in published message? 

133 self.tl.hub.publish(key, data) 

134 

135 self._in_sequence[sender] = in_sequence 

136 # self.request_missing() 

137 

138 elif to in ('', '*'): # anyone or OOB message, do not manage sequences 

139 logging.debug(f">> [{self.uid}] Received <{payload}> from {sender}:{addr} to '{to}' in {baddr}") 

140 (key, data) = payload 

141 func = self._oob_handlers.get(key) 

142 if func: 

143 func(key, data, sender) 

144 else: 

145 self.tl.hub.publish(key, data) 

146 else: 

147 logging.warn(f">> [{self.uid}] Ignoring <{payload}> from {sender}:{addr} to '{to}' in {baddr}") 

148 

149 def send(self, key, data, to, address): 

150 """Send a (key, data) pub/sub message to a specific peers or 

151 broadcasting. 

152 

153 - If `to` is empty or '*' then message is consider broadcast. 

154 - Else is added to the specific peer queue. 

155 - `address` could be an <broadcast> address or direct one. 

156 """ 

157 logging.debug(f"send: {self.uid} -> {to}: {address}") 

158 

159 # get peer sequence 

160 if to not in ('', '*'): 

161 sequence = self._out_sequence.setdefault(to, 0) 

162 self._out_sequence[to] += 1 

163 else: 

164 sequence = -1 

165 

166 # encode raw message 

167 payload = (key, data) 

168 raw = json.encode((self.uid, to, sequence, payload)).encode() 

169 raw = zlib.compress(raw) 

170 

171 # save message for replaying when network fails 

172 if to not in ('', '*'): 

173 output = self._output.get(to) 

174 if output is None: 

175 output = self._output[to] = dict() 

176 output[sequence] = (raw, address) 

177 

178 self.transport.sendto(raw, address) 

179 

180 @domestic(restart=2) 

181 def request_missing(self): 

182 "domestic task for request missing datagrams" 

183 for sender, inputs in self._input.items(): 

184 keys = inputs.keys() 

185 if keys: 

186 m = max(keys) 

187 missing = set(range(self._in_sequence[sender], m)) 

188 missing.difference_update(keys) 

189 _, address = inputs[m] 

190 if missing: 

191 logging.debug(f'{self.uid} Request MISSING: {missing}') 

192 self.send('MISSING', missing, '', ('<broadcast>', 9999)) 

193 else: 

194 watermark = self._in_sequence[sender] - 1 

195 logging.debug(f'{self.uid} Sending WATERMARK: {watermark}') 

196 self.send('WATERMARK', watermark, '', ('<broadcast>', 9999)) 

197 foo = 1 

198 

199 def _resend_missing(self, key, missing, sender): 

200 "resent some missing datagrams" 

201 logging.debug(f'{self.uid} Resend MISSING: {missing}') 

202 output = self._output.get(sender) 

203 if output and missing: 

204 for sequence in missing: 

205 info = output.get(sequence) 

206 info and self.transport.sendto(*info) 

207 watermark = min(missing) - 1 

208 self._update_watermark(key, watermark, sender) 

209 

210 def _update_watermark(self, key, watermark, sender): 

211 "Update the watermark and remove the already sent messages" 

212 output = self._output.get(sender, {}) 

213 for k in [k for k in output.keys() if k <= watermark]: 

214 output.pop(k) 

215 

216 @domestic() 

217 def _shake_channels(self): 

218 """Domestic function that tries to *shake* peers channels taking a 

219 proactive role sending a random packet or the last sent packet 

220 to the peer. 

221 

222 This behavior pretend to update remote queues in order the remote 

223 peers can request for updated information. 

224 

225 This is due *missing packets* request or the *last packet* maybe can 

226 be lost so peers may think there is nothing else to do, but 

227 actually they have missing information that may never 

228 receive unless the sender sent more packets that hopefuly update the 

229 last packet sequence creating a *hole* that can be filled. 

230 """ 

231 if random.random() > 0.5: 

232 # send Last Packet (to update missing in remote site) 

233 for sender, sequence in self._out_sequence.items(): 

234 info = self._output[sender].get(sequence - 1) 

235 if info: 

236 logging.debug(f'{self.uid} Sending Last: {sequence-1}') 

237 self.transport.sendto(*info) 

238 else: 

239 # send a random packet because if more probable that 

240 # remote peer already know the end of the queue 

241 # so we're proactivel send a possible 'missing' packet 

242 # due missing requests may also be lost 

243 for sender, info in self._output.items(): 

244 if info: 

245 sequence = random.choice(list(info.keys())) 

246 info = info[sequence] 

247 logging.debug(f'{self.uid} Sending Random Packet: {sequence}') 

248 self.transport.sendto(*info) 

249 

250 

251 

252class TL(UAgent): 

253 """An asyncio Transport Layer 

254 

255 This class send pub/sub message to remote peers. 

256 

257 - Message sendings are reliable. 

258 - Uses two channels for UDP and UDP broadcasting. 

259 

260 User can simply publish to `/net/publish` channel in the same `realm` 

261 or use `send()` method. 

262 """ 

263 def __init__(self, uid=None, app=None, uri='>udp://:9090', loop=None): 

264 super(TL, self).__init__(uid, app) 

265 self.uri = uri 

266 self._uri = parse_uri(uri, bind='') 

267 self.channels = dict() 

268 "active transports and protocols" 

269 

270 self._loop = None # set in start() 

271 

272 "Future that points out when all channels are ready for use." 

273 

274 self.app.hub.subscribe('/net/publish', self.send_pubsub) 

275 self.app.hub.subscribe('/net/receive', self.receive_pubsub) 

276 

277 foo = 1 

278 

279 async def main(self): 

280 """`main` asyncio method for TL agent. 

281 

282 - open channels 

283 - wait until agent is stopped 

284 - close channels 

285 """ 

286 # open communications channels 

287 await self.start() 

288 self.channels_ready.set_result(True) 

289 

290 # now coroutine holds until Agent is stopped 

291 await self.running 

292 

293 # close channels 

294 await self.stop() 

295 

296 def stop(self): 

297 for link, (transport, protocol, _, _) in self.channels.items(): 

298 protocol.stop_domestics() # del protocol ? 

299 # transport.close() # done in self.stop 

300 

301 self.running.set_result(False) 

302 

303 

304 def send_pubsub(self, key, message): 

305 """Send a pub/sub (key, message) by broadcasting to all peers in 

306 the same <broadcast> address. 

307 

308 TODO: need to review encapsulation. Do not use URecord. 

309 """ 

310 # payload = (key, message) 

311 raw = pickle.dumps(message) 

312 raw = base64.urlsafe_b64encode(raw) 

313 # raw = zlib.compress(raw) 

314 transport, protocol, baddr, saddr = self.channels['broadcast'] 

315 protocol.send(key='/net/receive', data=raw, to='*', address=saddr) 

316 

317 def receive_pubsub(self, key, raw): 

318 """Inject a remote (key, data) pub/sub message""" 

319 # raw = zlib.decompress(raw) 

320 raw = base64.urlsafe_b64decode(raw) 

321 key, data = pickle.loads(raw) 

322 self.hub.publish(key, data) 

323 

324 async def start(self): 

325 """Create open channels for <broadcast> and direct communications. 

326 

327 - Create a UDP socket with specific options. 

328 - Bind the socket. 

329 - Create a asyncio transport and protocol. 

330 - Add to self.channels dict 

331 """ 

332 

333 # TODO: are asyncio.events instead futures? 

334 self.running = self._loop.create_future() 

335 self.channels_ready = self._loop.create_future() 

336 

337 channels = { 

338 'broadcast': { 

339 'options': [(socket.SO_BROADCAST, 1), (socket.SO_REUSEADDR, 1), ], 

340 'bind': ('', 9999), 

341 }, 

342 'direct': { 

343 'options': [(socket.SO_BROADCAST, 1), (socket.SO_REUSEADDR, 1), ], 

344 } 

345 } 

346 address = list(self._uri['address']) 

347 for i, (link, info) in enumerate(channels.items()): 

348 sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, 

349 socket.SOL_UDP) 

350 for option, value in info['options']: 

351 sock.setsockopt(socket.SOL_SOCKET, option, value) 

352 

353 addr = info.get('bind', tuple(address)) 

354 sock.bind(addr) 

355 # print(self.uri, link, addr) 

356 # address[1] += 1 # next link will be binded to the next port 

357 

358 factory = UProtocolFactory( 

359 uid=self.uid, link=link, tl=self 

360 ) 

361 

362 transport, protocol = await self._loop.create_datagram_endpoint( 

363 lambda: factory(), 

364 sock=sock, 

365 ) 

366 if link in ('broadcast', ): 

367 saddr = tuple([f'<{link}>', addr[1]]) 

368 else: 

369 saddr = addr 

370 self.channels[link] = (transport, protocol, addr, saddr) 

371 foo = 1 

372 

373 async def stop(self): 

374 """Close all active channels by closing its transports.""" 

375 for link, (transport, _, _, _) in self.channels.items(): 

376 transport.close() 

377 

378 

379# HEADER_FORMAT = '>L' 

380# HEADER_SIZE = struct.calcsize(HEADER_FORMAT) 

381# class UStreamAgent(UAgent): 

382 # def __init__(self, uid=None, realm='', uri='tcp://localhost:9090'): 

383 # super().__init__(uid, realm) 

384 # self.uri = uri 

385 # self._uri = parse.urlparse(uri) 

386 # self.address = split_netloc(self._uri.netloc) 

387 

388 # self.reader = None 

389 # self.writer = None 

390 

391 # self.running = asyncio.Future() 

392 # self.channels_ready = asyncio.Future() 

393 

394 # self.msg = b'' 

395 

396 # async def _next_message(self, reader, writer): 

397 # context = {} 

398 # while len(self.msg) < HEADER_SIZE: 

399 # data = await reader.read(0x1000) 

400 # self.msg += data 

401 

402 # size, = struct.unpack_from(HEADER_FORMAT, self.msg) 

403 # cut = size + HEADER_SIZE 

404 # while len(self.msg) < cut: 

405 # data = await reader.read(0x1000) 

406 # self.msg += data 

407 

408 # data, self.msg = self.msg[4:cut], self.msg[cut:] 

409 # data = json.decode(data) 

410 # return data, context 

411 

412 # async def _send(self, data, writer): 

413 # msg = bytes(json.encode(data), encoding='UTF-8') 

414 # msg = struct.pack(HEADER_FORMAT, len(msg)) + msg 

415 # writer.write(msg) 

416 # await writer.drain() 

417 

418 # async def main(self): 

419 # await self.start() 

420 # self.channels_ready.set_result(True) 

421 # while not self.running.done(): 

422 # await asyncio.sleep(1000) 

423 # message, context = await self._next_message(self.reader, self.writer) 

424 # self.hande(message, context) 

425 

426 # def _handle(self, message, context): 

427 # print(f"receive: {message}") 

428 # if message == 'quit': 

429 # self.running.set_result(False) 

430 

431 # async def start(self): 

432 # raise NotImplementedError('Must be overriden') 

433 

434 

435# class UClient(UStreamAgent): 

436 # """Encapsulate sending objects to UServer""" 

437 

438 # def __init__(self, uid=None, realm='', uri='tcp://localhost:9090'): 

439 # super().__init__(uid, realm, uri) 

440 

441 # async def start(self): 

442 # return asyncio.open_connection(*self.address) 

443 

444 # # async def main(self): 

445 # # self.running = asyncio.Future() 

446 # # self.reader, self.writer = await asyncio.open_connection(*self.address) 

447 # # self.channels_ready.set_result(True) 

448 # # while not self.running.done(): 

449 # # message = await self._next_message(self.reader, self.writer) 

450 # # print(f"Server send: {message}") 

451 # # if message == 'quit': 

452 # # self.running.set_result(False) 

453 

454 # async def send(self, data): 

455 # while not self.channels_ready.done(): 

456 # await asyncio.sleep(0.1) 

457 # assert self.writer 

458 # await self._send(data, self.writer) 

459 

460# class UServer(UStreamAgent): 

461 # """Give access to exposed methods using TCP as well""" 

462 

463 # async def start(self): 

464 # scheme = self._uri.scheme 

465 # if scheme in ('tcp', ): 

466 # self.server = await asyncio.start_server(self.tcp_handle, *self.address) 

467 # # self.reader, self.writer = server.sockets 

468 # elif scheme in ('udp', ): 

469 # foo = 1 

470 # else: 

471 # raise RuntimeError(f'Unknown scheme {scheme}') 

472 

473 

474 # async def main(self): 

475 # self.running = asyncio.Future() 

476 # async with await asyncio.start_server(self.tcp_handle, *self.address) as server: 

477 # addr = server.sockets[0].getsockname() 

478 # print(f'Serving on {addr}: quit to shutdown server') 

479 # await self.running 

480 # foo = 1 

481 

482 # async def tcp_handle(self, reader, writer): 

483 # addr = writer.get_extra_info('peername') 

484 # message = f"{addr!r} is connected !!!!" 

485 # print(message) 

486 # while not self.running.done(): 

487 # message = await self._next_message(reader, writer) 

488 # print(f"Client send: {message}") 

489 

490 # await self._send(message, writer) 

491 # if message == "exit": 

492 # message = f"{addr!r} wants to close the connection." 

493 # print(message) 

494 # break 

495 # if message == 'quit': 

496 # message = f"{addr!r} wants to shutdown the server" 

497 # self.running.set_result(False) 

498 # writer.close() 

499 

500def test_Request(): 

501 uri = 'tl://121212/uagent/dbrecords/foo' 

502 req = Request(uri) 

503 

504 foo = 1 

505 

506 

507def test_parse_uri(): 

508 uri = parse_uri('http://myhost/mypath') 

509 assert uri['scheme'] == 'http' 

510 assert uri['fservice'] == 'myhost' 

511 assert uri['host'] == 'myhost' 

512 assert uri['path'] == '/mypath' 

513 

514 uri = parse_uri('http://agp@myhost/mypath') 

515 assert uri['scheme'] == 'http' 

516 assert uri['fservice'] == 'agp@myhost' 

517 assert uri['host'] == 'myhost' 

518 assert uri['path'] == '/mypath' 

519 

520 uri = parse_uri('http://agp:mypassword@myhost/mypath') 

521 assert uri['scheme'] == 'http' 

522 assert uri['fservice'] == 'agp:mypassword@myhost' 

523 assert uri['host'] == 'myhost' 

524 assert uri['path'] == '/mypath' 

525 assert uri['user'] == 'agp' 

526 assert uri['password'] == 'mypassword' 

527 

528 uri = parse_uri('http://agp:mypassword@myhost/mypath?myquery=1#42') 

529 assert uri['scheme'] == 'http' 

530 assert uri['fservice'] == 'agp:mypassword@myhost' 

531 assert uri['host'] == 'myhost' 

532 assert uri['path'] == '/mypath' 

533 assert uri['user'] == 'agp' 

534 assert uri['password'] == 'mypassword' 

535 assert uri['query'] == 'myquery=1' 

536 assert uri['fragment'] == '42' 

537 

538 uri = parse_uri('>http://agp:mypassword@myhost/mypath?myquery=1&foo=3#42') 

539 assert uri['direction'] == '>' 

540 assert uri['fscheme'] == '>http' 

541 assert uri['scheme'] == 'http' 

542 assert uri['fservice'] == 'agp:mypassword@myhost' 

543 assert uri['host'] == 'myhost' 

544 assert uri['path'] == '/mypath' 

545 assert uri['user'] == 'agp' 

546 assert uri['password'] == 'mypassword' 

547 assert uri['query'] == 'myquery=1&foo=3' 

548 assert uri['fragment'] == '42' 

549 

550 foo = 1 

551 

552if __name__ == '__main__': 

553 test_parse_uri() 

554 test_Request()