Support direct WebSocket connections in ASGI server

miguelgrinberg · May 22, 2020 · c0e2817 · c0e2817
1 parent f7bbd97
commit c0e2817
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 3 deletions.
diff --git a/engineio/async_drivers/asgi.py b/engineio/async_drivers/asgi.py
@@ -135,7 +135,7 @@ async def read(self, length=None):
             if event['type'] == 'http.request':
                 payload += event.get('body') or b''
     elif event['type'] == 'websocket.connect':
-        await send({'type': 'websocket.accept'})
+        pass
     else:
         return {}
 
@@ -187,6 +187,14 @@ async def read(self, length=None):
 
 async def make_response(status, headers, payload, environ):
     headers = [(h[0].encode('utf-8'), h[1].encode('utf-8')) for h in headers]
+    if 'HTTP_SEC_WEBSOCKET_VERSION' in environ:
+        if status.startswith('200 '):
+            await environ['asgi.send']({'type': 'websocket.accept',
+                                        'headers': headers})
+        else:
+            await environ['asgi.send']({'type': 'websocket.close'})
+        return
+
     await environ['asgi.send']({'type': 'http.response.start',
                                 'status': int(status.split(' ')[0]),
                                 'headers': headers})
@@ -207,6 +215,7 @@ def __init__(self, handler):
     async def __call__(self, environ):
         self.asgi_receive = environ['asgi.receive']
         self.asgi_send = environ['asgi.send']
+        await self.asgi_send({'type': 'websocket.accept'})
         await self.handler(self)
 
     async def close(self):

diff --git a/tests/asyncio/test_async_asgi.py b/tests/asyncio/test_async_asgi.py
@@ -347,7 +347,7 @@ def test_translate_websocket_request(self):
                         (b'content-length', b'123')],
             'path': '/foo/bar',
             'query_string': b'baz=1'}, receive, send))
-        send.mock.assert_called_once_with({'type': 'websocket.accept'})
+        send.mock.assert_not_called()
 
     def test_translate_unknown_request(self):
         receive = AsyncMock(return_value={'type': 'http.foo'})
@@ -358,7 +358,6 @@ def test_translate_unknown_request(self):
             'query_string': b'baz=1'}, receive, send))
         self.assertEqual(environ, {})
 
-    # @mock.patch('async_aiohttp.aiohttp.web.Response')
     def test_make_response(self):
         environ = {
             'asgi.send': AsyncMock()
@@ -370,3 +369,23 @@ def test_make_response(self):
              'headers': [(b'foo', b'bar')]})
         environ['asgi.send'].mock.assert_any_call(
             {'type': 'http.response.body', 'body': b'payload'})
+
+    def test_make_response_websocket_accept(self):
+        environ = {
+            'asgi.send': AsyncMock(),
+            'HTTP_SEC_WEBSOCKET_VERSION': 'foo',
+        }
+        _run(async_asgi.make_response('200 OK', [('foo', 'bar')],
+                                      b'payload', environ))
+        environ['asgi.send'].mock.assert_called_with(
+            {'type': 'websocket.accept', 'headers': [(b'foo', b'bar')]})
+
+    def test_make_response_websocket_reject(self):
+        environ = {
+            'asgi.send': AsyncMock(),
+            'HTTP_SEC_WEBSOCKET_VERSION': 'foo',
+        }
+        _run(async_asgi.make_response('401 UNAUTHORIZED', [('foo', 'bar')],
+                                      b'payload', environ))
+        environ['asgi.send'].mock.assert_called_with(
+            {'type': 'websocket.close'})