how to reject a websocket request

[xiispace]

i write the code below

class DemoWs(WebsocketEndpoint):
    async def on_connect(self, websocket):
        if should_disconnect(websocket):
            await websocket.close()
        await websocket.accept()

but the exception happened

Traceback (most recent call last):
  File "/srv/wechat/service/venv/lib/python3.6/site-packages/starlette/endpoints.py", line 68, in __call__
    raise exc from None
  File "/srv/wechat/service/venv/lib/python3.6/site-packages/starlette/endpoints.py", line 59, in __call__
    message = await websocket.receive()
  File "/srv/wechat/service/venv/lib/python3.6/site-packages/starlette/websockets.py", line 40, in receive
    message = await self._receive()
  File "/srv/wechat/service/venv/lib/python3.6/site-packages/uvicorn/protocols/websockets/websockets_impl.py", line 226, in asgi_receive
    data = await self.recv()
  File "/srv/wechat/service/venv/lib/python3.6/site-packages/websockets/protocol.py", line 419, in recv
    return_when=asyncio.FIRST_COMPLETED,
  File "/usr/lib64/python3.6/asyncio/tasks.py", line 311, in wait
    fs = {ensure_future(f, loop=loop) for f in set(fs)}
  File "/usr/lib64/python3.6/asyncio/tasks.py", line 311, in <setcomp>
    fs = {ensure_future(f, loop=loop) for f in set(fs)}
  File "/usr/lib64/python3.6/asyncio/tasks.py", line 526, in ensure_future
    raise TypeError('An asyncio.Future, a coroutine or an awaitable is '
TypeError: An asyncio.Future, a coroutine or an awaitable is required

[unmade]

Did you try to accept it first, then close?

class DemoWs(WebsocketEndpoint):
    async def on_connect(self, websocket):
        await websocket.accept()
        if should_disconnect(websocket):
            await websocket.close()

[tomchristie]

You’ve forgotten to return after the close, so I guess you’re falling through to the accept block. (Though the exception isn’t very clear)

[xiispace]

@unmade @tomchristie thanks to your reply. I found the exception was throw in https://github.com/encode/starlette/blob/master/starlette/endpoints.py#L65
..
https://github.com/aaugustin/websockets/blob/master/src/websockets/protocol.py#L438
if i close the websocket before accept, the websocket will abort handshake and the self.transfer_data_task will be None, so receive from websocket will raise a Exception.
I don't know that's the design or bug in websockets, so i make a workaround for this problem.

class DemoWs(WebsocketEndpoint):
    async def on_connect(self, websocket):
        if should_disconnect(websocket):
            await websocket.close()
            return
        await websocket.accept()

    async def __call__(self, receive: Receive, send: Send) -> None:
        websocket = WebSocket(self.scope, receive=receive, send=send)
        await self.on_connect(websocket)

        close_code = status.WS_1000_NORMAL_CLOSURE

        try:
            while True:
                #workaround, catach the Exception if websocket handshake not finished
                try:
                    message = await websocket.receive()
                except TypeError:
                    break
                if message["type"] == "websocket.receive":
                    data = await self.decode(websocket, message)
                    await self.on_receive(websocket, data)
                elif message["type"] == "websocket.disconnect":
                    close_code = int(message.get("code", status.WS_1000_NORMAL_CLOSURE))
                    break
        except Exception as exc:
            close_code = status.WS_1011_INTERNAL_ERROR
            raise exc from None
        finally:
            await self.on_disconnect(websocket, close_code)

[bofeng]

Had same problem here, feel like we should be allowed to directly close it before accept, like the code @xiispace posted in the on_connect function .

Other than the workaround method from @xiispace to overwrite the method, any progress on starlette side to solve this?

[tomchristie]

feel like we should be allowed to directly close it before accept, like the code @xiispace posted in the on_connect function .

The problem with the snippet posted is that it's missing a return statement.

class DemoWs(WebsocketEndpoint):
    async def on_connect(self, websocket):
        if should_disconnect(websocket):
            await websocket.close()
            # return
        await websocket.accept()

[bofeng]

Hi @tomchristie Thank you for the response, but adding "return" won't solve the problem.

The exception is thrown here: https://github.com/encode/starlette/blob/master/starlette/endpoints.py#L65 when websocket tries to receive after it is closed.

Adding "return" will return to line https://github.com/encode/starlette/blob/master/starlette/endpoints.py#L59 , but it will still try to do the following receive, which throws the exception here https://github.com/encode/starlette/blob/master/starlette/endpoints.py#L74

[tomchristie]

Gotcha, so we should be checking if the websocket was closed at

starlette/starlette/endpoints.py

Line 59 in c294ba6

await self.on_connect(websocket)

and returning if so, right?

[bofeng]

Seems so. But I am not sure in this case, should it still call the on_disconnect hook (I guess in this case it is still "connected" but not "accepted"?) .. if so, it needs to call the on_disconnect hook before return. or, if "connected" is the status after calling accept, then this sock is not treated as "connected", no need to call on_disconnect hook