import asyncio

import time

import six

from . import exceptions

from . import packet

from . import payload

from . import socket

class AsyncSocket(socket.Socket):

def create_queue(self):

return asyncio.Queue()

async def poll(self):

"""Wait for packets to send to the client."""

try:

packets = [await asyncio.wait_for(self.queue.get(),

self.server.ping_timeout)]

self.queue.task_done()

except (asyncio.TimeoutError, asyncio.CancelledError):

raise exceptions.QueueEmpty()

if packets == [None]:

return []

try:

packets.append(self.queue.get_nowait())

self.queue.task_done()

except asyncio.QueueEmpty:

pass

return packets

async def receive(self, pkt):

"""Receive packet from the client."""

self.server.logger.info('%s: Received packet %s data %s',

self.sid, packet.packet_names[pkt.packet_type],

pkt.data if not isinstance(pkt.data, bytes)

else '<binary>')

if pkt.packet_type == packet.PING:

self.last_ping = time.time()

await self.send(packet.Packet(packet.PONG, pkt.data))

elif pkt.packet_type == packet.MESSAGE:

await self.server._trigger_event('message', self.sid, pkt.data)

elif pkt.packet_type == packet.UPGRADE:

await self.send(packet.Packet(packet.NOOP))

elif pkt.packet_type == packet.CLOSE:

await self.close(wait=False, abort=True)

else:

raise exceptions.UnknownPacketError()

async def send(self, pkt):

"""Send a packet to the client."""

if self.closed:

raise IOError('Socket is closed')

if time.time() - self.last_ping > self.server.ping_timeout:

self.server.logger.info('%s: Client is gone, closing socket',

self.sid)

return await self.close(wait=False, abort=True)

self.server.logger.info('%s: Sending packet %s data %s',

self.sid, packet.packet_names[pkt.packet_type],

pkt.data if not isinstance(pkt.data, bytes)

else '<binary>')

await self.queue.put(pkt)

async def handle_get_request(self, environ):

"""Handle a long-polling GET request from the client."""

connections = [

s.strip()

for s in environ.get('HTTP_CONNECTION', '').lower().split(',')]

transport = environ.get('HTTP_UPGRADE', '').lower()

if 'upgrade' in connections and transport in self.upgrade_protocols:

self.server.logger.info('%s: Received request to upgrade to %s',

self.sid, transport)

return await getattr(self, '_upgrade_' + transport)(environ)

try:

packets = await self.poll()

except exceptions.QueueEmpty as e:

await self.close(wait=False)

raise e

return packets

async def handle_post_request(self, environ):

"""Handle a long-polling POST request from the client."""

length = int(environ.get('CONTENT_LENGTH', '0'))

if length > self.server.max_http_buffer_size:

raise exceptions.ContentTooLongError()

else:

body = await environ['wsgi.input'].read(length)

p = payload.Payload(encoded_payload=body)

for pkt in p.packets:

await self.receive(pkt)

async def close(self, wait=True, abort=False):

"""Close the socket connection."""

if not self.closed and not self.closing:

self.closing = True

reraise_exc = None

try:

await self.server._trigger_event('disconnect', self.sid)

except Exception as e:

reraise_exc = e

if not abort:

await self.send(packet.Packet(packet.CLOSE))

self.closed = True

if wait:

await self.queue.join()

if reraise_exc:

raise reraise_exc

async def _upgrade_websocket(self, environ):

"""Upgrade the connection from polling to websocket."""

if self.upgraded:

raise IOError('Socket has been upgraded already')

if self.server._async['websocket'] is None or \

self.server._async['websocket_class'] is None:

# the selected async mode does not support websocket

return self.server._bad_request()

websocket_class = getattr(self.server._async['websocket'],

self.server._async['websocket_class'])

ws = websocket_class(self._websocket_handler)

return await ws(environ)

async def _websocket_handler(self, ws):

"""Engine.IO handler for websocket transport."""

if self.connected:

# the socket was already connected, so this is an upgrade

await self.queue.join() # flush the queue first

pkt = await ws.wait()

if pkt != packet.Packet(packet.PING,

data=six.text_type('probe')).encode(

always_bytes=False):

self.server.logger.info(

'%s: Failed websocket upgrade, no PING packet', self.sid)

return

await ws.send(packet.Packet(

packet.PONG,

data=six.text_type('probe')).encode(always_bytes=False))

await self.send(packet.Packet(packet.NOOP))

pkt = await ws.wait()

decoded_pkt = packet.Packet(encoded_packet=pkt)

if decoded_pkt.packet_type != packet.UPGRADE:

self.upgraded = False

self.server.logger.info(

('%s: Failed websocket upgrade, expected UPGRADE packet, '

'received %s instead.'),

self.sid, pkt)

return

self.upgraded = True

else:

self.connected = True

self.upgraded = True

# start separate writer thread

async def writer():

while True:

packets = None

try:

packets = await self.poll()

except exceptions.QueueEmpty:

break

if not packets:

# empty packet list returned -> connection closed

break

try:

for pkt in packets:

await ws.send(pkt.encode(always_bytes=False))

except:

break

writer_task = asyncio.ensure_future(writer())

self.server.logger.info(

'%s: Upgrade to websocket successful', self.sid)

reraise_exc = None

while True:

p = None

try:

p = await asyncio.wait_for(ws.wait(), self.server.ping_timeout)

except:

break

if p is None:

# connection closed by client

break

if isinstance(p, six.text_type): # pragma: no cover

p = p.encode('utf-8')

pkt = packet.Packet(encoded_packet=p)

try:

await self.receive(pkt)

except exceptions.UnknownPacketError:

pass

except Exception as e:

# if we get an unexpected exception (such as something in an

# application event handler) we close the connection properly

# and then reraise the exception

reraise_exc = e

break

await self.queue.put(None) # unlock the writer task so it can exit

await asyncio.wait_for(writer_task, timeout=None)

await self.close(wait=True, abort=True)

if reraise_exc:

raise reraise_exc