如何在aiohttp中实现WebSocket未认证连接的超时关闭机制？

在aiohttp里实现带认证的WebSocket服务器其实思路很清晰，虽然它没有像Tornado那样直接暴露on_open方法，但我们可以在WebSocket处理协程里模拟这个逻辑，再结合asyncio的任务管理来实现认证超时。我给你梳理一套完整的实现方案：

1. 模拟on_open逻辑 #

aiohttp中，当客户端发起WebSocket连接并完成握手后，我们定义的WebSocket处理协程会被触发，此时就相当于Tornado里on_open的时机——可以在这里做连接初始化操作，比如记录连接日志、初始化连接相关的状态。

2. 实现认证超时机制 #

我们可以用asyncio.create_task创建一个超时等待任务，同时等待客户端发送的第一条认证消息：

• 如果在超时时间内收到合法的认证消息，就取消超时任务，进入正常通信流程
• 如果超时未收到消息，或者认证失败，直接关闭WebSocket连接

import asyncio
from aiohttp import web

async def websocket_handler(request):
    # 完成WebSocket握手，这就是我们的"on_open"时机
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    print(f"新WebSocket连接建立: {request.remote}")

    # 配置认证超时时间（比如5秒）
    AUTH_TIMEOUT = 5

    # 定义超时任务：超时后自动关闭连接
    async def auth_timeout_task():
        await asyncio.sleep(AUTH_TIMEOUT)
        if not ws.closed:
            await ws.send_json({"error": "认证超时，请在连接后5秒内发送认证信息"})
            await ws.close()

    # 将超时任务加入事件循环
    timeout_task = asyncio.create_task(auth_timeout_task())

    try:
        # 等待客户端发送的第一条消息（必须是认证消息）
        msg = await ws.receive()
        if msg.type == web.WSMsgType.TEXT:
            # 替换成你的实际认证逻辑，比如校验token、API key等
            try:
                auth_data = msg.json()
                if auth_data.get("token") == "your_valid_auth_token":
                    # 认证成功：取消超时任务，通知客户端
                    timeout_task.cancel()
                    await ws.send_json({"status": "认证成功"})

                    # 进入正常消息处理流程
                    async for msg in ws:
                        if msg.type == web.WSMsgType.TEXT:
                            await ws.send_json({"response": f"收到消息: {msg.data}"})
                        elif msg.type == web.WSMsgType.ERROR:
                            print(f"WebSocket连接异常关闭: {ws.exception()}")
                else:
                    # 认证失败：关闭连接
                    await ws.send_json({"error": "无效的认证信息"})
                    await ws.close()
            except ValueError:
                # 收到的不是合法JSON，直接关闭
                await ws.send_json({"error": "认证消息格式错误，需为JSON格式"})
                await ws.close()
        else:
            # 收到非文本消息（比如二进制），直接关闭
            await ws.close()
    except asyncio.CancelledError:
        # 超时任务被取消时会触发此异常，无需额外处理
        pass
    finally:
        # 确保超时任务被正确清理，避免残留资源
        if not timeout_task.done():
            timeout_task.cancel()
            try:
                await timeout_task
            except asyncio.CancelledError:
                pass
        print(f"WebSocket连接关闭: {request.remote}")

# 启动服务器
app = web.Application()
app.add_routes([web.get('/ws', websocket_handler)])

if __name__ == '__main__':
    web.run_app(app, host='0.0.0.0', port=8080)

• 连接初始化：await ws.prepare(request)完成WebSocket握手后，就可以执行连接初始化逻辑，比如记录客户端IP、初始化连接状态变量。
• 超时任务管理：通过asyncio.create_task启动超时任务，它会在后台等待指定时间，不会阻塞我们等待客户端消息的逻辑；认证成功后用timeout_task.cancel()取消该任务即可。
• 认证逻辑扩展：你可以把示例中的token校验替换成实际业务逻辑，比如从数据库查询有效token、校验签名等。
• 并发安全：如果需要管理所有已认证的连接（比如广播消息），可以用一个全局集合存储已认证的ws对象，但要记得用asyncio.Lock保护集合的读写操作，避免并发问题。

内容的提问来源于stack exchange，提问作者bman