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 响应，因为 FastAPI 的 ExceptionMiddleware 似乎无法处理此类情况。

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

Answer 1

握手完成后，协议将从 HTTP 更改为 WebSocket< /代码>。如果您尝试在 websocket 端点内引发 HTTP 异常，您会发现这是不可能的，或者返回 HTTP 响应（例如，return JSONResponse( ...status_code=404))，您将收到内部服务器错误，即 ASGI 可调用已返回但未发送握手。

选项 1

因此，如果您希望在协议升级之前拥有某种检查机制，则需要使用 中间件，如下所示。在中间件内部，您不能引发异常，但可以返回响应（即 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 实例来执行此操作，您可以具有与上面相同的逻辑，但是，在 is_user_allowed() 中传递 websocket 实例> 方法，以及捕获 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() 方法来终止连接。如果您愿意，可以使用如下所示的内容。但是，except 块中的 return 语句将引发内部服务器错误（即，ASGI 可调用返回而不发送握手。），如上所述早些时候。

@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)

Answer 2

我有同样的问题。在某些情况下，您希望使用常规 HTTP 标头回复传入的 Websocket 升级请求。例如，有时特殊的标头用于在工作平衡器等中配置网络路由。

我花了一些时间才弄清楚，不同协议的处理程序似乎与 FastAPI 一起生活在更基本的级别上。对路由的不同协议的处理似乎是从 ASGI 服务器（例如 uvicorn）传播到 starlette 到 FastAPI 。因此，就我而言，解决方案是简单地使用 --ws none 或 ws="none" 禁用 uvicorn 中的 Websocket 协议。因此，我使用 .get() 处理程序的常规路由会收到 websocket 升级请求，并可以使用 HTTP 标头进行回复。

这对于我的应用程序来说没问题，其中 FastAPI 充当网关，不需要实现 Websocket 本身。不幸的是，我不知道这个技术堆栈有处理混合情况的解决方案。

Answer 3

我设法在 HTTP 异常上使用 exception_handler，并在 websocket 尚未打开时直接向 ASGI 发送“websocket.http.response.start”信号。

（刚刚在 uvicorn==0.27.0 和 fastapi==0.103.1 上测试）

from fastapi import FastAPI, HTTPException
from fastapi.websockets import WebSocket
from fastapi.exception_handlers import http_exception_handler

app = FastAPI()

@app.exception_handler(HTTPException)
async def custom_http_exception_handler(request, exc):
    if not isinstance(request, WebSocket):
        return await http_exception_handler(request, exc)
    websocket = request
    if websocket.application_state == 1:  # 1 = CONNECTED
        await websocket.close(code=1008, reason=str(exc))  # 1008 = WS_1008_POLICY_VIOLATION
        return
    headers = getattr(exc, "headers") or []
    await websocket._send({"type": "websocket.http.response.start", "status": exc.status_code, "headers": headers})