FastAPI：使用HTTP响应拒绝WebSocket连接

Question

FastAPI：使用HTTP响应拒绝WebSocket连接

3

在基于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)

由于FastAPI的ExceptionMiddleware似乎无法处理这种情况，因此异常不会转换为404响应。

是否有任何原生/内置的方法来支持这种“拒绝”流程？

- shevron

2个回答

0

我遇到了同样的问题。在某些情况下，您希望用常规的HTTP头回复一个传入的Websocket升级请求。例如，有时特殊的头部用于配置工作负载均衡器等网络路由。

我花了一些时间才弄清楚，不同协议的处理程序似乎在FastAPI的更基本级别上存在。路由的不同协议处理似乎是从ASGI服务器（例如uvicorn）传播到starlette再到FastAPI的。所以在我的情况下，解决方案就是简单地在uvicorn中禁用Websocket协议，使用--ws none或ws="none"）。结果是，我的常规路由使用.get()处理程序接收到了Websocket升级请求，并可以回复一个HTTP头。

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

- fungs

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Chris · Accepted Answer

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

选项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实例来完成这个任务，你可以采用与上面相同的逻辑，但是将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()方法来终止连接。如果您愿意，可以使用以下内容。但是，在except块中的return语句将抛出内部服务器错误（即，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)