Pass reason when closing a WebSocket connection

miguelgrinberg · Jul 22, 2021 · 583c7db · 583c7db
1 parent 1a6e5b1
commit 583c7db
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/src/engineio/async_drivers/asgi.py b/src/engineio/async_drivers/asgi.py
@@ -192,7 +192,13 @@ async def make_response(status, headers, payload, environ):
             await environ['asgi.send']({'type': 'websocket.accept',
                                         'headers': headers})
         else:
-            await environ['asgi.send']({'type': 'websocket.close'})
+            if payload:
+                reason = payload.decode('utf-8') \
+                    if isinstance(payload, bytes) else str(payload)
+                await environ['asgi.send']({'type': 'websocket.close',
+                                            'reason': reason})
+            else:
+                await environ['asgi.send']({'type': 'websocket.close'})
         return
 
     await environ['asgi.send']({'type': 'http.response.start',

diff --git a/tests/asyncio/test_async_asgi.py b/tests/asyncio/test_async_asgi.py
@@ -492,6 +492,20 @@ def test_make_response_websocket_reject(self):
                 '401 UNAUTHORIZED', [('foo', 'bar')], b'payload', environ
             )
         )
+        environ['asgi.send'].mock.assert_called_with(
+            {'type': 'websocket.close', 'reason': 'payload'}
+        )
+
+    def test_make_response_websocket_reject_no_payload(self):
+        environ = {
+            'asgi.send': AsyncMock(),
+            'asgi.scope': {'type': 'websocket'},
+        }
+        _run(
+            async_asgi.make_response(
+                '401 UNAUTHORIZED', [('foo', 'bar')], None, environ
+            )
+        )
         environ['asgi.send'].mock.assert_called_with(
             {'type': 'websocket.close'}
         )