FastAPI：拒绝带有 HTTP 响应的 WebSocket 连接

Question

在基于 FastAPI 的 Web 应用程序中，我有一个 WebSocket 端点，仅当满足某些条件时才应允许连接，否则它应返回答复HTTP 404而不是升级与HTTP 101.

据我了解，协议完全支持这一点，但我找不到任何方法可以使用 FastAPI 或 Starlette 来做到这一点。

如果我有类似的东西：

@router.websocket("/foo")
async def ws_foo(request: WebSocket):
    if _user_is_allowed(request):
        await request.accept()
        _handle_ws_connection(request)
    else:
        raise HTTPException(status_code=404)

该异常不会转换为 404 响应，因为 FastAPIExceptionMiddleware似乎无法处理此类情况。

是否有任何本机/内置方式支持这种“拒绝”流程？

Answer 1

握手完成后，协议将从 更改HTTP为WebSocket。如果您尝试HTTP在 websocket 端点内部引发异常，您会发现这是不可能的，或者返回响应HTTP（例如return JSONResponse(...status_code=404)），您将收到内部服务器错误，即ASGI callable returned without sending handshake。

选项1

因此，如果您希望在协议升级之前拥有某种检查机制，则需要使用 a Middleware，如下所示。在中间件内部，您不能引发异常，但可以返回响应（即，Response，JSONResponse，PlainTextResponse等），这实际上是FastAPI在幕后处理异常的方式。作为参考，请查看这篇文章以及此处的讨论。

async def is_user_allowed(request: Request):
    # if conditions are not met, return False
    print(request['headers'])
    print(request.client)
    return False

@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    if not await is_user_allowed(request):
        return JSONResponse(content={"message": "User not allowed"}, status_code=404)
    response = await call_next(request)
    return response

或者，如果您愿意，您可以使用is_user_allowed()引发需要用try-except块捕获的自定义异常的方法：

class UserException(Exception):
    def __init__(self, message):
        self.message = message
        super().__init__(message)

async def is_user_allowed(request: Request):
    # if conditions are not met, raise UserException
    raise UserException(message="User not allowed.")
    
@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    try:
        await is_user_allowed(request)
    except UserException as e:
        return JSONResponse(content={"message": f'{e.message}'}, status_code=404)
    response = await call_next(request)
    return response

选项2

但是，如果您需要使用websocket实例来执行此操作，则可以具有与上面相同的逻辑，但是，websocket在is_user_allowed()方法中传递实例，并捕获 websocket 端点内的异常（受this启发）。

@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    await ws.accept()
    try:
        await is_user_allowed(ws)
        await handle_conn(ws)
    except UserException as e:
        await ws.send_text(e.message) # optionally send a message to the client before closing the connection
        await ws.close()

但是，在上面，您必须首先接受连接，以便close()在引发异常时可以调用该方法来终止连接。如果您愿意，可以使用如下所示的内容。然而，如前所述，块return内的该语句except将引发内部服务器错误（即）。ASGI callable returned without sending handshake.

@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    try:
        await is_user_allowed(ws)
    except UserException as e:
        return
    await ws.accept()
    await handle_conn(ws)