如何使用aiohttp(python 3.6)实现websocket感知的反向代理

时间:2018-09-18 09:11:17

标签: python jupyter-notebook python-3.6 reverse-proxy aiohttp

我正在尝试使用aiohttp为jupyter笔记本实现特定于应用程序的反向代理。对于http请求,它工作正常,但是websocket转发不起作用。来自浏览器的请求到达并被转发,但是jupyter即将发布没有任何响应。我假设我的websocket客户端代码不响应jupyter的传入消息。

在Jupyter方面,唯一表明不对劲的消息是这样的消息:

WebSocket ping timeout after 90009 ms.

这是我编写代理的尝试

from aiohttp import web
from aiohttp import client
import aiohttp
import logging
import pprint

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


baseUrl = 'http://0.0.0.0:8888'
mountPoint = '/fakeUuid'

async def handler(req):
    proxyPath = req.match_info.get('proxyPath','no proxyPath placeholder defined')
    reqH = req.headers.copy()
    if reqH['connection'] == 'Upgrade' and reqH['upgrade'] == 'websocket' and req.method == 'GET':

      ws_server = web.WebSocketResponse()
      await ws_server.prepare(req)
      logger.info('##### WS_SERVER %s' % pprint.pformat(ws_server))

      client_session = aiohttp.ClientSession()
      async with client_session.ws_connect(baseUrl+req.path_qs,
        headers = { 'cookie': reqH['cookie'] },
      ) as ws_client:
        logger.info('##### WS_CLIENT %s' % pprint.pformat(ws_client))

        async for server_msg in ws_server:
          logger.info('>>> msg from browser: %s',pprint.pformat(server_msg))
          if server_msg.type == aiohttp.WSMsgType.TEXT:
            await ws_client.send_str(server_msg.data)
          else:
            await ws_client.send_bytes(server_msg.data)

        async for client_msg in ws_client:
          logger.info('>>> msg from jupyter: %s',pprint.pformat(client_msg))
          if client_msg.tp == aiohttp.WSMsgType.TEXT:
            await ws_server.send_str(client_msg.data)
          else:
            await ws_server.send_bytes(client_msg.data)

        return ws_server
    else:
      async with client.request(
          req.method,baseUrl+mountPoint+proxyPath,
          headers = reqH,
          allow_redirects=False,
          data = await req.read()
      ) as res:
          headers = res.headers.copy()
          body = await res.read()
          return web.Response(
            headers = headers,
            status = res.status,
            body = body
          )
      return ws_server

app = web.Application()
app.router.add_route('*',mountPoint + '{proxyPath:.*}', handler)
web.run_app(app,port=3984)

1 个答案:

答案 0 :(得分:0)

经验教训:两个async for阻塞了当前函数的流程。通过使用asyncio.wait运行它们,我可以使它们同时运行。生成的程序如下所示:

from aiohttp import web
from aiohttp import client
import aiohttp
import asyncio
import logging
import pprint

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


baseUrl = 'http://0.0.0.0:8888'
mountPoint = '/fakeUuid'

async def handler(req):
    proxyPath = req.match_info.get('proxyPath','no proxyPath placeholder defined')
    reqH = req.headers.copy()
    if reqH['connection'] == 'Upgrade' and reqH['upgrade'] == 'websocket' and req.method == 'GET':

      ws_server = web.WebSocketResponse()
      await ws_server.prepare(req)
      logger.info('##### WS_SERVER %s' % pprint.pformat(ws_server))

      client_session = aiohttp.ClientSession(cookies=req.cookies)
      async with client_session.ws_connect(
        baseUrl+req.path_qs,
      },
      ) as ws_client:
        logger.info('##### WS_CLIENT %s' % pprint.pformat(ws_client))

        async def wsforward(ws_from,ws_to):
          async for msg in ws_from:
            logger.info('>>> msg: %s',pprint.pformat(msg))
            mt = msg.type
            md = msg.data
            if mt == aiohttp.WSMsgType.TEXT:
              await ws_to.send_str(md)
            elif mt == aiohttp.WSMsgType.BINARY:
              await ws_to.send_bytes(md)
            elif mt == aiohttp.WSMsgType.PING:
              await ws_to.ping()
            elif mt == aiohttp.WSMsgType.PONG:
              await ws_to.pong()
            elif ws_to.closed:
              await ws_to.close(code=ws_to.close_code,message=msg.extra)
            else:
              raise ValueError('unexpecte message type: %s',pprint.pformat(msg))

        finished,unfinished = await asyncio.wait([wsforward(ws_server,ws_client),wsforward(ws_client,ws_server)],return_when=asyncio.FIRST_COMPLETED)

        return ws_server
    else:
      async with client.request(
          req.method,baseUrl+mountPoint+proxyPath,
          headers = reqH,
          allow_redirects=False,
          data = await req.read()
      ) as res:
          headers = res.headers.copy()
          body = await res.read()
          return web.Response(
            headers = headers,
            status = res.status,
            body = body
          )
      return ws_server

app = web.Application()
app.router.add_route('*',mountPoint + '{proxyPath:.*}', handler)
web.run_app(app,port=3984)