uWebSockets.js中WebSocket升级时UserData传递的正确方式

在使用uWebSockets.js构建WebSocket服务时，开发者经常会遇到需要在WebSocket连接建立后访问原始HTTP请求信息的需求。本文深入探讨了这一场景下的正确实现方式，并解释了常见的误区。

问题背景

许多开发者尝试在WebSocket升级过程中将HttpRequest对象直接存储为UserData，期望在后续的WebSocket通信中能够继续访问请求信息。这种直觉性的做法看似合理，但实际上存在根本性的问题。

错误示范

典型的错误实现方式如下：

res.upgrade(
  { req },  // 直接将HttpRequest对象存入UserData
  req.getHeader('sec-websocket-key'),
  req.getHeader('sec-websocket-protocol'),
  req.getHeader('sec-websocket-extensions'),
  context
);

然后在message处理器中尝试访问：

message(ws: WebSocket<HttpRequest>, message) {
  const req = ws.getUserData();
  console.log(req.getHeader('x-forwarded-for'));
}

这种实现会导致运行时错误，因为HttpRequest对象在升级完成后就已经失效，无法继续使用。

正确解决方案

正确的做法是在升级阶段提取所需的请求信息，而不是存储整个HttpRequest对象：

res.upgrade(
  {
    ip: req.getHeader('x-forwarded-for'),
    // 其他需要保存的请求信息
  },
  req.getHeader('sec-websocket-key'),
  req.getHeader('sec-websocket-protocol'),
  req.getHeader('sec-websocket-extensions'),
  context
);

然后在WebSocket处理器中访问这些预存的数据：

message(ws: WebSocket<{ip: string}>, message) {
  console.log(ws.getUserData().ip);
}

技术原理

HttpRequest对象之所以不能在WebSocket连接建立后继续使用，是因为：

1. 该对象仅在HTTP请求处理期间有效
2. WebSocket协议建立后，底层连接已从HTTP协议切换为WebSocket协议
3. 原始HTTP请求的上下文在协议切换后已被释放

最佳实践建议

1. 提前提取数据：在upgrade阶段提取所有需要的请求信息
2. 最小化存储：只保存真正需要的数据，避免内存浪费
3. 类型安全：为UserData定义明确的TypeScript类型
4. 错误处理：考虑请求头可能不存在的情况

interface WSUserData {
  ip?: string;
  protocol?: string;
  // 其他自定义字段
}

res.upgrade(
  {
    ip: req.getHeader('x-forwarded-for') || 'unknown',
    protocol: req.getHeader('sec-websocket-protocol')
  } as WSUserData,
  // ...其他参数
);

通过遵循这些原则，开发者可以安全高效地在WebSocket连接中访问原始请求信息，同时避免潜在的内存问题和运行时错误。