Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

567 строки
23KB

  1. import asyncio
  2. try:
  3. import aiohttp
  4. except ImportError: # pragma: no cover
  5. aiohttp = None
  6. import six
  7. try:
  8. import websockets
  9. except ImportError: # pragma: no cover
  10. websockets = None
  11. from . import client
  12. from . import exceptions
  13. from . import packet
  14. from . import payload
  15. class AsyncClient(client.Client):
  16. """An Engine.IO client for asyncio.
  17. This class implements a fully compliant Engine.IO web client with support
  18. for websocket and long-polling transports, compatible with the asyncio
  19. framework on Python 3.5 or newer.
  20. :param logger: To enable logging set to ``True`` or pass a logger object to
  21. use. To disable logging set to ``False``. The default is
  22. ``False``.
  23. :param json: An alternative json module to use for encoding and decoding
  24. packets. Custom json modules must have ``dumps`` and ``loads``
  25. functions that are compatible with the standard library
  26. versions.
  27. :param request_timeout: A timeout in seconds for requests. The default is
  28. 5 seconds.
  29. """
  30. def is_asyncio_based(self):
  31. return True
  32. async def connect(self, url, headers={}, transports=None,
  33. engineio_path='engine.io'):
  34. """Connect to an Engine.IO server.
  35. :param url: The URL of the Engine.IO server. It can include custom
  36. query string parameters if required by the server.
  37. :param headers: A dictionary with custom headers to send with the
  38. connection request.
  39. :param transports: The list of allowed transports. Valid transports
  40. are ``'polling'`` and ``'websocket'``. If not
  41. given, the polling transport is connected first,
  42. then an upgrade to websocket is attempted.
  43. :param engineio_path: The endpoint where the Engine.IO server is
  44. installed. The default value is appropriate for
  45. most cases.
  46. Note: this method is a coroutine.
  47. Example usage::
  48. eio = engineio.Client()
  49. await eio.connect('http://localhost:5000')
  50. """
  51. if self.state != 'disconnected':
  52. raise ValueError('Client is not in a disconnected state')
  53. valid_transports = ['polling', 'websocket']
  54. if transports is not None:
  55. if isinstance(transports, six.text_type):
  56. transports = [transports]
  57. transports = [transport for transport in transports
  58. if transport in valid_transports]
  59. if not transports:
  60. raise ValueError('No valid transports provided')
  61. self.transports = transports or valid_transports
  62. self.queue = self.create_queue()
  63. return await getattr(self, '_connect_' + self.transports[0])(
  64. url, headers, engineio_path)
  65. async def wait(self):
  66. """Wait until the connection with the server ends.
  67. Client applications can use this function to block the main thread
  68. during the life of the connection.
  69. Note: this method is a coroutine.
  70. """
  71. if self.read_loop_task:
  72. await self.read_loop_task
  73. async def send(self, data, binary=None):
  74. """Send a message to a client.
  75. :param data: The data to send to the client. Data can be of type
  76. ``str``, ``bytes``, ``list`` or ``dict``. If a ``list``
  77. or ``dict``, the data will be serialized as JSON.
  78. :param binary: ``True`` to send packet as binary, ``False`` to send
  79. as text. If not given, unicode (Python 2) and str
  80. (Python 3) are sent as text, and str (Python 2) and
  81. bytes (Python 3) are sent as binary.
  82. Note: this method is a coroutine.
  83. """
  84. await self._send_packet(packet.Packet(packet.MESSAGE, data=data,
  85. binary=binary))
  86. async def disconnect(self, abort=False):
  87. """Disconnect from the server.
  88. :param abort: If set to ``True``, do not wait for background tasks
  89. associated with the connection to end.
  90. Note: this method is a coroutine.
  91. """
  92. if self.state == 'connected':
  93. await self._send_packet(packet.Packet(packet.CLOSE))
  94. await self.queue.put(None)
  95. self.state = 'disconnecting'
  96. await self._trigger_event('disconnect', run_async=False)
  97. if self.current_transport == 'websocket':
  98. await self.ws.close()
  99. if not abort:
  100. await self.read_loop_task
  101. self.state = 'disconnected'
  102. try:
  103. client.connected_clients.remove(self)
  104. except ValueError: # pragma: no cover
  105. pass
  106. self._reset()
  107. def start_background_task(self, target, *args, **kwargs):
  108. """Start a background task.
  109. This is a utility function that applications can use to start a
  110. background task.
  111. :param target: the target function to execute.
  112. :param args: arguments to pass to the function.
  113. :param kwargs: keyword arguments to pass to the function.
  114. This function returns an object compatible with the `Thread` class in
  115. the Python standard library. The `start()` method on this object is
  116. already called by this function.
  117. Note: this method is a coroutine.
  118. """
  119. return asyncio.ensure_future(target(*args, **kwargs))
  120. async def sleep(self, seconds=0):
  121. """Sleep for the requested amount of time.
  122. Note: this method is a coroutine.
  123. """
  124. return await asyncio.sleep(seconds)
  125. def create_queue(self):
  126. """Create a queue object."""
  127. q = asyncio.Queue()
  128. q.Empty = asyncio.QueueEmpty
  129. return q
  130. def create_event(self):
  131. """Create an event object."""
  132. return asyncio.Event()
  133. def _reset(self):
  134. if self.http: # pragma: no cover
  135. asyncio.ensure_future(self.http.close())
  136. super()._reset()
  137. async def _connect_polling(self, url, headers, engineio_path):
  138. """Establish a long-polling connection to the Engine.IO server."""
  139. if aiohttp is None: # pragma: no cover
  140. self.logger.error('aiohttp not installed -- cannot make HTTP '
  141. 'requests!')
  142. return
  143. self.base_url = self._get_engineio_url(url, engineio_path, 'polling')
  144. self.logger.info('Attempting polling connection to ' + self.base_url)
  145. r = await self._send_request(
  146. 'GET', self.base_url + self._get_url_timestamp(), headers=headers,
  147. timeout=self.request_timeout)
  148. if r is None:
  149. self._reset()
  150. raise exceptions.ConnectionError(
  151. 'Connection refused by the server')
  152. if r.status != 200:
  153. raise exceptions.ConnectionError(
  154. 'Unexpected status code {} in server response'.format(
  155. r.status))
  156. try:
  157. p = payload.Payload(encoded_payload=await r.read())
  158. except ValueError:
  159. six.raise_from(exceptions.ConnectionError(
  160. 'Unexpected response from server'), None)
  161. open_packet = p.packets[0]
  162. if open_packet.packet_type != packet.OPEN:
  163. raise exceptions.ConnectionError(
  164. 'OPEN packet not returned by server')
  165. self.logger.info(
  166. 'Polling connection accepted with ' + str(open_packet.data))
  167. self.sid = open_packet.data['sid']
  168. self.upgrades = open_packet.data['upgrades']
  169. self.ping_interval = open_packet.data['pingInterval'] / 1000.0
  170. self.ping_timeout = open_packet.data['pingTimeout'] / 1000.0
  171. self.current_transport = 'polling'
  172. self.base_url += '&sid=' + self.sid
  173. self.state = 'connected'
  174. client.connected_clients.append(self)
  175. await self._trigger_event('connect', run_async=False)
  176. for pkt in p.packets[1:]:
  177. await self._receive_packet(pkt)
  178. if 'websocket' in self.upgrades and 'websocket' in self.transports:
  179. # attempt to upgrade to websocket
  180. if await self._connect_websocket(url, headers, engineio_path):
  181. # upgrade to websocket succeeded, we're done here
  182. return
  183. self.ping_loop_task = self.start_background_task(self._ping_loop)
  184. self.write_loop_task = self.start_background_task(self._write_loop)
  185. self.read_loop_task = self.start_background_task(
  186. self._read_loop_polling)
  187. async def _connect_websocket(self, url, headers, engineio_path):
  188. """Establish or upgrade to a WebSocket connection with the server."""
  189. if websockets is None: # pragma: no cover
  190. self.logger.error('websockets package not installed')
  191. return False
  192. websocket_url = self._get_engineio_url(url, engineio_path,
  193. 'websocket')
  194. if self.sid:
  195. self.logger.info(
  196. 'Attempting WebSocket upgrade to ' + websocket_url)
  197. upgrade = True
  198. websocket_url += '&sid=' + self.sid
  199. else:
  200. upgrade = False
  201. self.base_url = websocket_url
  202. self.logger.info(
  203. 'Attempting WebSocket connection to ' + websocket_url)
  204. # get the cookies from the long-polling connection so that they can
  205. # also be sent the the WebSocket route
  206. cookies = None
  207. if self.http:
  208. cookies = '; '.join(["{}={}".format(cookie.key, cookie.value)
  209. for cookie in self.http._cookie_jar])
  210. headers = headers.copy()
  211. headers['Cookie'] = cookies
  212. try:
  213. ws = await websockets.connect(
  214. websocket_url + self._get_url_timestamp(),
  215. extra_headers=headers)
  216. except (websockets.exceptions.InvalidURI,
  217. websockets.exceptions.InvalidHandshake,
  218. OSError):
  219. if upgrade:
  220. self.logger.warning(
  221. 'WebSocket upgrade failed: connection error')
  222. return False
  223. else:
  224. raise exceptions.ConnectionError('Connection error')
  225. if upgrade:
  226. p = packet.Packet(packet.PING, data='probe').encode(
  227. always_bytes=False)
  228. try:
  229. await ws.send(p)
  230. except Exception as e: # pragma: no cover
  231. self.logger.warning(
  232. 'WebSocket upgrade failed: unexpected send exception: %s',
  233. str(e))
  234. return False
  235. try:
  236. p = await ws.recv()
  237. except Exception as e: # pragma: no cover
  238. self.logger.warning(
  239. 'WebSocket upgrade failed: unexpected recv exception: %s',
  240. str(e))
  241. return False
  242. pkt = packet.Packet(encoded_packet=p)
  243. if pkt.packet_type != packet.PONG or pkt.data != 'probe':
  244. self.logger.warning(
  245. 'WebSocket upgrade failed: no PONG packet')
  246. return False
  247. p = packet.Packet(packet.UPGRADE).encode(always_bytes=False)
  248. try:
  249. await ws.send(p)
  250. except Exception as e: # pragma: no cover
  251. self.logger.warning(
  252. 'WebSocket upgrade failed: unexpected send exception: %s',
  253. str(e))
  254. return False
  255. self.current_transport = 'websocket'
  256. if self.http: # pragma: no cover
  257. await self.http.close()
  258. self.logger.info('WebSocket upgrade was successful')
  259. else:
  260. try:
  261. p = await ws.recv()
  262. except Exception as e: # pragma: no cover
  263. raise exceptions.ConnectionError(
  264. 'Unexpected recv exception: ' + str(e))
  265. open_packet = packet.Packet(encoded_packet=p)
  266. if open_packet.packet_type != packet.OPEN:
  267. raise exceptions.ConnectionError('no OPEN packet')
  268. self.logger.info(
  269. 'WebSocket connection accepted with ' + str(open_packet.data))
  270. self.sid = open_packet.data['sid']
  271. self.upgrades = open_packet.data['upgrades']
  272. self.ping_interval = open_packet.data['pingInterval'] / 1000.0
  273. self.ping_timeout = open_packet.data['pingTimeout'] / 1000.0
  274. self.current_transport = 'websocket'
  275. self.state = 'connected'
  276. client.connected_clients.append(self)
  277. await self._trigger_event('connect', run_async=False)
  278. self.ws = ws
  279. self.ping_loop_task = self.start_background_task(self._ping_loop)
  280. self.write_loop_task = self.start_background_task(self._write_loop)
  281. self.read_loop_task = self.start_background_task(
  282. self._read_loop_websocket)
  283. return True
  284. async def _receive_packet(self, pkt):
  285. """Handle incoming packets from the server."""
  286. packet_name = packet.packet_names[pkt.packet_type] \
  287. if pkt.packet_type < len(packet.packet_names) else 'UNKNOWN'
  288. self.logger.info(
  289. 'Received packet %s data %s', packet_name,
  290. pkt.data if not isinstance(pkt.data, bytes) else '<binary>')
  291. if pkt.packet_type == packet.MESSAGE:
  292. await self._trigger_event('message', pkt.data, run_async=True)
  293. elif pkt.packet_type == packet.PONG:
  294. self.pong_received = True
  295. elif pkt.packet_type == packet.CLOSE:
  296. await self.disconnect(abort=True)
  297. elif pkt.packet_type == packet.NOOP:
  298. pass
  299. else:
  300. self.logger.error('Received unexpected packet of type %s',
  301. pkt.packet_type)
  302. async def _send_packet(self, pkt):
  303. """Queue a packet to be sent to the server."""
  304. if self.state != 'connected':
  305. return
  306. await self.queue.put(pkt)
  307. self.logger.info(
  308. 'Sending packet %s data %s',
  309. packet.packet_names[pkt.packet_type],
  310. pkt.data if not isinstance(pkt.data, bytes) else '<binary>')
  311. async def _send_request(
  312. self, method, url, headers=None, body=None,
  313. timeout=None): # pragma: no cover
  314. if self.http is None or self.http.closed:
  315. self.http = aiohttp.ClientSession()
  316. http_method = getattr(self.http, method.lower())
  317. try:
  318. return await http_method(
  319. url, headers=headers, data=body,
  320. timeout=aiohttp.ClientTimeout(total=timeout))
  321. except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
  322. self.logger.info('HTTP %s request to %s failed with error %s.',
  323. method, url, exc)
  324. async def _trigger_event(self, event, *args, **kwargs):
  325. """Invoke an event handler."""
  326. run_async = kwargs.pop('run_async', False)
  327. ret = None
  328. if event in self.handlers:
  329. if asyncio.iscoroutinefunction(self.handlers[event]) is True:
  330. if run_async:
  331. return self.start_background_task(self.handlers[event],
  332. *args)
  333. else:
  334. try:
  335. ret = await self.handlers[event](*args)
  336. except asyncio.CancelledError: # pragma: no cover
  337. pass
  338. except:
  339. self.logger.exception(event + ' async handler error')
  340. if event == 'connect':
  341. # if connect handler raised error we reject the
  342. # connection
  343. return False
  344. else:
  345. if run_async:
  346. async def async_handler():
  347. return self.handlers[event](*args)
  348. return self.start_background_task(async_handler)
  349. else:
  350. try:
  351. ret = self.handlers[event](*args)
  352. except:
  353. self.logger.exception(event + ' handler error')
  354. if event == 'connect':
  355. # if connect handler raised error we reject the
  356. # connection
  357. return False
  358. return ret
  359. async def _ping_loop(self):
  360. """This background task sends a PING to the server at the requested
  361. interval.
  362. """
  363. self.pong_received = True
  364. self.ping_loop_event.clear()
  365. while self.state == 'connected':
  366. if not self.pong_received:
  367. self.logger.info(
  368. 'PONG response has not been received, aborting')
  369. if self.ws:
  370. await self.ws.close()
  371. await self.queue.put(None)
  372. break
  373. self.pong_received = False
  374. await self._send_packet(packet.Packet(packet.PING))
  375. try:
  376. await asyncio.wait_for(self.ping_loop_event.wait(),
  377. self.ping_interval)
  378. except (asyncio.TimeoutError,
  379. asyncio.CancelledError): # pragma: no cover
  380. pass
  381. self.logger.info('Exiting ping task')
  382. async def _read_loop_polling(self):
  383. """Read packets by polling the Engine.IO server."""
  384. while self.state == 'connected':
  385. self.logger.info(
  386. 'Sending polling GET request to ' + self.base_url)
  387. r = await self._send_request(
  388. 'GET', self.base_url + self._get_url_timestamp(),
  389. timeout=max(self.ping_interval, self.ping_timeout) + 5)
  390. if r is None:
  391. self.logger.warning(
  392. 'Connection refused by the server, aborting')
  393. await self.queue.put(None)
  394. break
  395. if r.status != 200:
  396. self.logger.warning('Unexpected status code %s in server '
  397. 'response, aborting', r.status)
  398. await self.queue.put(None)
  399. break
  400. try:
  401. p = payload.Payload(encoded_payload=await r.read())
  402. except ValueError:
  403. self.logger.warning(
  404. 'Unexpected packet from server, aborting')
  405. await self.queue.put(None)
  406. break
  407. for pkt in p.packets:
  408. await self._receive_packet(pkt)
  409. self.logger.info('Waiting for write loop task to end')
  410. await self.write_loop_task
  411. self.logger.info('Waiting for ping loop task to end')
  412. self.ping_loop_event.set()
  413. await self.ping_loop_task
  414. if self.state == 'connected':
  415. await self._trigger_event('disconnect', run_async=False)
  416. try:
  417. client.connected_clients.remove(self)
  418. except ValueError: # pragma: no cover
  419. pass
  420. self._reset()
  421. self.logger.info('Exiting read loop task')
  422. async def _read_loop_websocket(self):
  423. """Read packets from the Engine.IO WebSocket connection."""
  424. while self.state == 'connected':
  425. p = None
  426. try:
  427. p = await self.ws.recv()
  428. except websockets.exceptions.ConnectionClosed:
  429. self.logger.info(
  430. 'Read loop: WebSocket connection was closed, aborting')
  431. await self.queue.put(None)
  432. break
  433. except Exception as e:
  434. self.logger.info(
  435. 'Unexpected error "%s", aborting', str(e))
  436. await self.queue.put(None)
  437. break
  438. if isinstance(p, six.text_type): # pragma: no cover
  439. p = p.encode('utf-8')
  440. pkt = packet.Packet(encoded_packet=p)
  441. await self._receive_packet(pkt)
  442. self.logger.info('Waiting for write loop task to end')
  443. await self.write_loop_task
  444. self.logger.info('Waiting for ping loop task to end')
  445. self.ping_loop_event.set()
  446. await self.ping_loop_task
  447. if self.state == 'connected':
  448. await self._trigger_event('disconnect', run_async=False)
  449. try:
  450. client.connected_clients.remove(self)
  451. except ValueError: # pragma: no cover
  452. pass
  453. self._reset()
  454. self.logger.info('Exiting read loop task')
  455. async def _write_loop(self):
  456. """This background task sends packages to the server as they are
  457. pushed to the send queue.
  458. """
  459. while self.state == 'connected':
  460. # to simplify the timeout handling, use the maximum of the
  461. # ping interval and ping timeout as timeout, with an extra 5
  462. # seconds grace period
  463. timeout = max(self.ping_interval, self.ping_timeout) + 5
  464. packets = None
  465. try:
  466. packets = [await asyncio.wait_for(self.queue.get(), timeout)]
  467. except (self.queue.Empty, asyncio.TimeoutError,
  468. asyncio.CancelledError):
  469. self.logger.error('packet queue is empty, aborting')
  470. break
  471. if packets == [None]:
  472. self.queue.task_done()
  473. packets = []
  474. else:
  475. while True:
  476. try:
  477. packets.append(self.queue.get_nowait())
  478. except self.queue.Empty:
  479. break
  480. if packets[-1] is None:
  481. packets = packets[:-1]
  482. self.queue.task_done()
  483. break
  484. if not packets:
  485. # empty packet list returned -> connection closed
  486. break
  487. if self.current_transport == 'polling':
  488. p = payload.Payload(packets=packets)
  489. r = await self._send_request(
  490. 'POST', self.base_url, body=p.encode(),
  491. headers={'Content-Type': 'application/octet-stream'},
  492. timeout=self.request_timeout)
  493. for pkt in packets:
  494. self.queue.task_done()
  495. if r is None:
  496. self.logger.warning(
  497. 'Connection refused by the server, aborting')
  498. break
  499. if r.status != 200:
  500. self.logger.warning('Unexpected status code %s in server '
  501. 'response, aborting', r.status)
  502. self._reset()
  503. break
  504. else:
  505. # websocket
  506. try:
  507. for pkt in packets:
  508. await self.ws.send(pkt.encode(always_bytes=False))
  509. self.queue.task_done()
  510. except websockets.exceptions.ConnectionClosed:
  511. self.logger.info(
  512. 'Write loop: WebSocket connection was closed, '
  513. 'aborting')
  514. break
  515. self.logger.info('Exiting write loop task')