amirouche/aiohttp+websockets.md

## aiohttp+websockets.md

      
    Raw
  

              aiohttp+websockets.md
            
          
    ** aiohttp now supports in its webframework websocket **
http://aiohttp.readthedocs.org/en/v0.14.1/web.html#websockets
This is a quick hack (ported from django-c10k) to get websocket working along side classic http with aiohttp web server. I think it would be better to inherit the aiohttp.web.RequestHandler and add the code to handle the upgrade in RequestHandler.start instead of overriding RequestHandler.transport.close in WebsocketResponse.
Anyway, it seems like it works.
requirements:
pip3 install aiohttp websockets


## app.py
import asyncio
from functools import wraps

from aiohttp import web
from websockets import handshake
from websockets import WebSocketCommonProtocol


INDEX = open('index.html').read().encode('utf-8')


class WebSocketResponse(web.Response):

    def __init__(self, request, switch_protocols):
        http11 = request._version == (1, 1)
        get_header = lambda k: dict(request.headers)[k.upper()]
        key = handshake.check_request(get_header)
        if not http11 or not key:
            super('Invalid WebSocket handshake.\n', status=400)
        else:
            headers = dict()
            set_header = headers.__setitem__
            handshake.build_response(set_header, key)
            self.switch_protocols = switch_protocols
            super().__init__(status=101, headers=headers)
            self._keep_alive = True
            request.transport.close = switch_protocols


@asyncio.coroutine
def handle(request):
    return web.Response(body=INDEX)


def websocket(handler):

    @asyncio.coroutine
    @wraps(handler)
    def wrapper(request, *args, **kwargs):
        transport = request.transport
        http_protocol = transport._protocol

        @asyncio.coroutine
        def run_ws_handler(ws):
            yield from handler(ws, request, *args, **kwargs)
            yield from ws.close()

        def switch_protocols():
            ws_protocol = WebSocketCommonProtocol()
            transport._protocol = ws_protocol
            ws_protocol.connection_made(transport)

            # Ensure aiohttp doesn't interfere.
            http_protocol.transport = None

            asyncio.async(run_ws_handler(ws_protocol))

        return WebSocketResponse(request, switch_protocols)

    return wrapper


@websocket
@asyncio.coroutine
def ws(ws, request):
    while ws.open:
        message = yield from ws.recv()
        if message:
            print(ws, request, message)
            yield from ws.send("Hello from aiohttp powered websocket!")


@asyncio.coroutine
def init(loop):
    app = web.Application(loop=loop)
    app.router.add_route('GET', '/', handle)
    app.router.add_route('GET', '/ws', ws)

    server = yield from loop.create_server(
        app.make_handler(),
        '127.0.0.1',
        8080
    )
    print("Server started at http://127.0.0.1:8080")
    return server


loop = asyncio.get_event_loop()
loop.run_until_complete(init(loop))
loop.run_forever()

## index.html
<!DOCTYPE html>
<html>
  <head>
    <meta charset="utf-8" />
    <title>WebSocket Test</title>
    <script language="javascript" type="text/javascript">
      var wsUri = "ws://127.0.0.1:8080/ws"; var output;
      function init() { output = document.getElementById("output");
      testWebSocket(); }
      function testWebSocket() { websocket = new WebSocket(wsUri);
      websocket.onopen = function(evt) { onOpen(evt) };
      websocket.onclose = function(evt) { onClose(evt) };
      websocket.onmessage = function(evt) { onMessage(evt) };
      websocket.onerror = function(evt) { onError(evt) }; }
      function onOpen(evt) { writeToScreen("CONNECTED"); doSend("WebSocket rocks"); }
      function onClose(evt) { writeToScreen("DISCONNECTED"); }
      function onMessage(evt) { console.log(evt); writeToScreen('<span style="color: blue;">RESPONSE: ' + evt.data+'</span>'); websocket.close(); }
      function onError(evt) { writeToScreen('<span style="color: red;">ERROR:</span> ' + evt.data); }
      function doSend(message) { writeToScreen("SENT: " + message);  websocket.send(message); }
      function writeToScreen(message) { var pre = document.createElement("p"); pre.style.wordWrap = "break-word"; pre.innerHTML = message; output.appendChild(pre); }  window.addEventListener("load", init, false);
    </script>
  </head>
  <body>
    <h2>WebSocket Test</h2>  <div id="output"></div>
  </body>
</html>
	import asyncio
	from functools import wraps

	from aiohttp import web
	from websockets import handshake
	from websockets import WebSocketCommonProtocol


	INDEX = open('index.html').read().encode('utf-8')


	class WebSocketResponse(web.Response):

	def __init__(self, request, switch_protocols):
	http11 = request._version == (1, 1)
	get_header = lambda k: dict(request.headers)[k.upper()]
	key = handshake.check_request(get_header)
	if not http11 or not key:
	super('Invalid WebSocket handshake.\n', status=400)
	else:
	headers = dict()
	set_header = headers.__setitem__
	handshake.build_response(set_header, key)
	self.switch_protocols = switch_protocols
	super().__init__(status=101, headers=headers)
	self._keep_alive = True
	request.transport.close = switch_protocols


	@asyncio.coroutine
	def handle(request):
	return web.Response(body=INDEX)


	def websocket(handler):

	@asyncio.coroutine
	@wraps(handler)
	def wrapper(request, args, *kwargs):
	transport = request.transport
	http_protocol = transport._protocol

	@asyncio.coroutine
	def run_ws_handler(ws):
	yield from handler(ws, request, args, *kwargs)
	yield from ws.close()

	def switch_protocols():
	ws_protocol = WebSocketCommonProtocol()
	transport._protocol = ws_protocol
	ws_protocol.connection_made(transport)

	# Ensure aiohttp doesn't interfere.
	http_protocol.transport = None

	asyncio.async(run_ws_handler(ws_protocol))

	return WebSocketResponse(request, switch_protocols)

	return wrapper


	@websocket
	@asyncio.coroutine
	def ws(ws, request):
	while ws.open:
	message = yield from ws.recv()
	if message:
	print(ws, request, message)
	yield from ws.send("Hello from aiohttp powered websocket!")


	@asyncio.coroutine
	def init(loop):
	app = web.Application(loop=loop)
	app.router.add_route('GET', '/', handle)
	app.router.add_route('GET', '/ws', ws)

	server = yield from loop.create_server(
	app.make_handler(),
	'127.0.0.1',
	8080
	)
	print("Server started at http://127.0.0.1:8080")
	return server


	loop = asyncio.get_event_loop()
	loop.run_until_complete(init(loop))
	loop.run_forever()
	<!DOCTYPE html>
	<html>
	<head>
	<meta charset="utf-8" />
	<title>WebSocket Test</title>
	<script language="javascript" type="text/javascript">
	var wsUri = "ws://127.0.0.1:8080/ws"; var output;
	function init() { output = document.getElementById("output");
	testWebSocket(); }
	function testWebSocket() { websocket = new WebSocket(wsUri);
	websocket.onopen = function(evt) { onOpen(evt) };
	websocket.onclose = function(evt) { onClose(evt) };
	websocket.onmessage = function(evt) { onMessage(evt) };
	websocket.onerror = function(evt) { onError(evt) }; }
	function onOpen(evt) { writeToScreen("CONNECTED"); doSend("WebSocket rocks"); }
	function onClose(evt) { writeToScreen("DISCONNECTED"); }
	function onMessage(evt) { console.log(evt); writeToScreen('<span style="color: blue;">RESPONSE: ' + evt.data+'</span>'); websocket.close(); }
	function onError(evt) { writeToScreen('<span style="color: red;">ERROR:</span> ' + evt.data); }
	function doSend(message) { writeToScreen("SENT: " + message); websocket.send(message); }
	function writeToScreen(message) { var pre = document.createElement("p"); pre.style.wordWrap = "break-word"; pre.innerHTML = message; output.appendChild(pre); } window.addEventListener("load", init, false);
	</script>
	</head>
	<body>
	<h2>WebSocket Test</h2> <div id="output"></div>
	</body>
	</html>