NodeJS中为第三方WebSocket聊天应用实现带认证的代理

嘿，我帮你梳理出一个可行的Node.js方案，完美适配你的需求——先验证令牌再代理WebSocket连接，咱们一步步来：

第一步：安装依赖 #

首先得装ws模块（Node.js生态中处理WebSocket的常用库），在项目目录执行：

npm install ws

第二步：核心代码实现 #

以下是完整的代理服务器代码，你可以根据实际情况替换占位内容：

const http = require('http');
const WebSocket = require('ws');
const { URL } = require('url');

// 替换成你的实际聊天服务器WebSocket地址
const TARGET_CHAT_SERVER_WS = 'ws://your-chat-server:port/chat-path';
// 代理服务器监听端口
const PROXY_PORT = 3000;

// 1. 创建HTTP服务器：同时处理令牌验证GET接口和WebSocket升级请求
const server = http.createServer(async (req, res) => {
  // 处理你的令牌验证GET接口（如果需要单独对外提供验证能力的话）
  if (req.url.startsWith('/api/validate-token')) {
    const token = req.headers['x-api-token'] || req.query.token;
    // 替换成你已经掌握的令牌验证逻辑
    const isValid = await yourExistingValidationLogic(token);

    if (isValid) {
      res.writeHead(200, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ success: true, message: 'Token valid' }));
    } else {
      res.writeHead(401, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ success: false, message: 'Invalid token' }));
    }
    return;
  }

  // 其他非预期请求返回404
  res.writeHead(404);
  res.end('Not Found');
});

// 2. 创建WebSocket服务器：不单独监听端口，复用HTTP服务器的升级事件
const wss = new WebSocket.Server({ noServer: true });

// 3. 处理WebSocket升级请求：先验令牌，再代理
server.on('upgrade', async (request, socket, head) => {
  // 提取令牌：支持从请求头(X-API-Token)或URL参数(token=xxx)获取
  const token = request.headers['x-api-token'] || 
                new URL(request.url, `http://${request.headers.host}`).searchParams.get('token');

  // 无令牌直接拒绝连接
  if (!token) {
    socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
    socket.destroy();
    return;
  }

  try {
    // 执行令牌验证逻辑
    const isValid = await yourExistingValidationLogic(token);
    // 如果你的验证是调用远程GET接口，可替换成下方的示例函数：
    // const isValid = await validateTokenViaRemoteAPI(token);

    if (!isValid) {
      socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
      socket.destroy();
      return;
    }

    // 4. 验证通过，代理连接到目标聊天服务器
    const targetWsClient = new WebSocket(TARGET_CHAT_SERVER_WS);

    // 目标服务器连接成功后，完成客户端的WebSocket握手
    targetWsClient.on('open', () => {
      wss.handleUpgrade(request, socket, head, (clientWs) => {
        // 双向转发消息：客户端 ↔ 目标聊天服务器
        clientWs.on('message', (data) => targetWsClient.send(data));
        targetWsClient.on('message', (data) => clientWs.send(data));

        // 处理连接关闭，避免资源泄漏
        clientWs.on('close', () => targetWsClient.close());
        targetWsClient.on('close', () => clientWs.close());

        // 错误处理
        clientWs.on('error', (err) => {
          console.error('客户端WebSocket错误:', err);
          targetWsClient.close();
        });
        targetWsClient.on('error', (err) => {
          console.error('目标聊天服务器WebSocket错误:', err);
          clientWs.close();
        });
      });
    });

    // 目标服务器连接失败时返回503
    targetWsClient.on('error', (err) => {
      console.error('连接目标聊天服务器失败:', err);
      socket.write('HTTP/1.1 503 Service Unavailable\r\n\r\n');
      socket.destroy();
    });

  } catch (err) {
    console.error('令牌验证过程出错:', err);
    socket.write('HTTP/1.1 500 Internal Server Error\r\n\r\n');
    socket.destroy();
  }
});

// 启动代理服务器
server.listen(PROXY_PORT, () => {
  console.log(`WebSocket代理服务器已启动，监听端口: ${PROXY_PORT}`);
});

// 【可选】如果你的令牌验证是调用远程GET接口，用这个示例函数
async function validateTokenViaRemoteAPI(token) {
  return new Promise((resolve, reject) => {
    const authApiOptions = {
      hostname: 'your-auth-server-domain',
      port: 80,
      path: `/api/validate-token?token=${encodeURIComponent(token)}`,
      method: 'GET'
    };

    const req = http.request(authApiOptions, (res) => {
      resolve(res.statusCode === 200);
    });

    req.on('error', (err) => reject(err));
    req.end();
  });
}

关键逻辑说明 #

• 复用HTTP服务器：既处理你的令牌验证GET接口，又处理WebSocket升级请求，不用额外开端口
• 令牌前置验证：在WebSocket握手的关键时刻（upgrade事件）先校验令牌，不通过直接拒绝连接，避免无效连接占用资源
• 双向消息转发：验证通过后，创建到目标服务器的WebSocket客户端连接，把客户端和目标服务器的消息双向绑定，实现透明代理
• 完善的异常处理：覆盖了无令牌、验证失败、目标服务器连接失败等多种场景，避免服务器挂死或资源泄漏

测试建议 #

1. 启动代理服务器，客户端连接ws://localhost:3000，记得在请求头携带X-API-Token或URL后追加?token=你的有效令牌
2. 测试验证失败场景：传无效令牌，客户端会收到401响应并被断开连接
3. 测试正常通信场景：用有效令牌连接，应该能和目标聊天服务器正常收发消息

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