Timeout WebSocket connections in aiohttp - python-3.x

Question

My WebSocket server implementation is open to the world, but the client is required to send an authenticate message after the connection was established or the server should close the connection.
How can I implement this in aiohttp? It seems, I need to do the following things:
Create an on_open method for every socket connection: I can't find a way (similarly to on_open in Tornado) to create such event.
Create a timer: asyncio's sleep or call_back methods of the main event loop may be used. But I can't find a way to send the WebSocketResponse to the callback function:
await asyncio.sleep(10, timer, loop=request.app.loop)
Closing the connection if not authenticated
This is what I had before with Tornado:
def open(self, *args, **kwargs):
self.timeout = ioloop.IOLoop.instance().add_timeout(
datetime.timedelta(seconds=60),
self._close_on_timeout
)
def remove_timeout_timer(self):
ioloop.IOLoop.instance().remove_timeout(self.timeout)
self.timeout = None
def on_message(self, message):
if message = 'AUTHENTICATE':
self.authenticated = True
self.remove_timeout_timer
def _close_on_timeout(self):
if not self.authenticated:
if self.ws_connection:
self.close()
Here is what I have using aiohttp for setting up a timer:
async def ensure_client_logged(ws):
await asyncio.sleep(3) # wait 3 seconds
await ws.send_str('hello')
async def ws_handler(request):
ws = web.WebSocketResponse()
asyncio.ensure_future(ensure_client_logged(ws), loop=request.app.loop)
But the code is running in a blocking way, meaning the server becomes unresponsive while it is sleeping.
Can someone please point me in the right direction?

Answer 1

You need to establish a deadline for the authentication procedure. asyncio.wait_for is a convenient way to do that:
async def ws_handler(request):
loop = asyncio.get_event_loop()
ws = web.WebSocketResponse()
loop.create_task(handle_client(ws))
async def handle_client(ws):
try:
authenticated = await asyncio.wait_for(_authenticate(ws), 10)
except asyncio.TimeoutError:
authenticated = False
if not authenticated:
ws.close()
return
# continue talking to the client
async def _authenticate(ws):
# implement authentication here, without worrying about
# timeout - the coroutine will be automatically canceled
# once the timeout elapses
...
return True # if successfully authenticated

Answer 2

This is a full working example for the benefits of the future users:
from aiohttp import web
import asyncio
async def wait_for_authentication(ws, app):
async for msg in ws:
if msg.type == web.WSMsgType.TEXT and msg.data == 'AUTHENTICATE': # Implement your own authentication
await ws.send_str('WELCOME')
return True
else:
await ws.send_str('NOT AUTHENTICATED')
async def authenticate(ws, app) -> bool:
try:
authenticated = await asyncio.wait_for(wait_for_authentication(ws, app), 5)
except asyncio.TimeoutError:
authenticated = False
if not authenticated:
await ws.send_str('The AUTHENTICATE command was not received. Closing the connection...')
await ws.close()
return False
async def ws_handler(request):
ws = web.WebSocketResponse()
await ws.prepare(request)
await request.app.loop.create_task(authenticate(ws, request.app))
async for msg in ws:
if msg.type != web.WSMsgType.TEXT:
continue
await ws.send_str(msg.data)
def init():
app = web.Application()
app.router.add_get('/', ws_handler)
return app
web.run_app(init())