[WIP] asyncio/aiohttp で websocket reverse proxy 書きたい

websocket_proxy.py

import asyncio
import aiohttp

from aiohttp import web


class WebsocketProxy(object):
    def __init__(self, upstream_url):
        self.upstream_url = upstream_url

    async def handler(self, request):
        session = aiohttp.ClientSession()
        async with session.ws_connect(self.upstream_url) as server_ws:
            client_ws = web.WebSocketResponse()
            await client_ws.prepare(request)

            print('pipe', client_ws, 'to', server_ws)
            done, pending = await asyncio.wait([
                self.pipe_ws(server_ws, client_ws),
                self.pipe_ws(client_ws, server_ws),
            ], return_when=asyncio.FIRST_COMPLETED)

            for f in pending:
                print('cancel', f)
                f.cancel()

            await server_ws.close()
        await client_ws.close()
        print('ws connection closed')
        session.close()
        return client_ws

    @staticmethod
    async def pipe_ws(ws_from, ws_to):
        async for msg in ws_from:
            #print(msg, 'received from', ws_from)
            if msg.tp == aiohttp.MsgType.binary:
                ws_to.send_bytes(msg.data)
            elif msg.tp == aiohttp.MsgType.text:
                ws_to.send_str(msg.data)
            elif msg.tp == aiohttp.MsgType.ping:
                ws_to.ping()
            elif msg.tp == aiohttp.MsgType.pong:
                ws_to.pong()
            else:
                assert msg.tp in [aiohttp.MsgType.closed,
                                  aiohttp.MsgType.error]
                break

        print('exit', ws_from, ws_to)


def main():
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--host', default='0.0.0.0')
    parser.add_argument('--port', type=int, default=8080)
    parser.add_argument('upstream_url')

    args = parser.parse_args()

    proxy = WebsocketProxy(args.upstream_url)
    app = web.Application()
    app.router.add_route('GET', '/', proxy.handler)
    web.run_app(app, host=args.host, port=args.port)


if __name__ == '__main__':
    main()