优化 WebSocket 实现，提升代码健壮性与可维护性

时间：2025-03-18 18:29:05

WebSocket 技术因其低延迟、双向通信的特点，成为实现实时功能的理想选择。然而，在实际开发中，WebSocket 的实现往往面临安全性、异常处理和性能优化等挑战。本文将基于一个 Vue.js 项目中的 WebSocket 实现（WebsocketMixin.js），探讨其优化后的设计思路、改进点以及如何提升代码的安全性、性能和可维护性。

一、优化背景

原始代码实现了基本的 WebSocket 功能，包括连接初始化、事件监听和重连机制。然而，经过审查发现，代码存在以下问题：

  1. 安全性不足：
  • 缺乏对 token 和 userId 的有效性检查，可能导致连接失败或安全问题。
  • URL 构造部分未对 window._CONFIG['domianURL'] 和 this.socketUrl 进行校验，可能存在非法构造的风险。
  1. 异常处理不完善：
  • 没有对关键逻辑（如 Vue.ls.get 和 store.getters.userInfo.id）进行空值检查。
  • 在发送数据时捕获的错误信息不够准确，使用了非标准的 err.code。
  1. 重复逻辑：
  • URL 构造部分多次判断 this.socketUrl 的前缀和后缀，导致代码冗余。
  1. 性能优化空间：
  • 重连逻辑可能因异常导致无限循环，缺乏对重连失败的进一步处理。

为了解决这些问题，我们对代码进行了全面优化。


import store from '@/store/';
import { ACCESS_TOKEN } from '@/store/mutation-types';
import Vue from 'vue';

export const WebsocketMixin = {
  data() {
    return {
      websock: null,
      lockReconnect: false,
      socketUrl: '', // 确保在使用前初始化
    };
  },
  mounted() {
    this.initWebSocket();
  },
  destroyed() {
    this.websocketOnclose();
  },
  methods: {
    initWebSocket() {
      try {
        const token = Vue.ls.get(ACCESS_TOKEN);
        if (!token) {
          console.error("Token is missing or invalid.");
          return;
        }

        const userId = store.getters.userInfo?.id;
        if (!userId) {
          console.error("User ID is missing or invalid.");
          return;
        }

        const domainUrl = window._CONFIG['domianURL'];
        if (!domainUrl || !/^(https?:\/\/)/.test(domainUrl)) {
          console.error("Invalid domain URL configuration.");
          return;
        }

        const sanitizedSocketUrl = this.sanitizeSocketUrl(this.socketUrl);
        const wsProtocol = domainUrl.startsWith('https://') ? 'wss://' : 'ws://';
        const url = `${wsProtocol}${domainUrl.replace(/^(https?:\/\/)/, '')}${sanitizedSocketUrl}${userId}/${token}`;

        this.websock = new WebSocket(url);

        this.websock.onopen = this.websocketOnopen;
        this.websock.onerror = this.websocketOnerror;
        this.websock.onmessage = this.websocketOnmessage;
        this.websock.onclose = this.websocketOnclose;
      } catch (err) {
        console.error("Failed to initialize WebSocket:", err.message);
      }
    },

    sanitizeSocketUrl(socketUrl) {
      if (!socketUrl) return '/';
      let sanitized = socketUrl.trim();
      if (!sanitized.startsWith('/')) sanitized = '/' + sanitized;
      if (!sanitized.endsWith('/')) sanitized += '/';
      return sanitized;
    },

    websocketOnopen() {
      console.log("WebSocket连接成功");
    },

    websocketOnerror(e) {
      console.error("WebSocket连接发生错误:", e.message);
      this.reconnect();
    },

    websocketOnclose(e) {
      console.log("WebSocket连接已关闭:", e.code, e.reason);
      this.reconnect();
    },

    websocketSend(text) {
      if (!this.websock || this.websock.readyState !== WebSocket.OPEN) {
        console.warn("WebSocket is not open, cannot send message.");
        return;
      }
      try {
        this.websock.send(text);
      } catch (err) {
        console.error("Failed to send message:", err.message);
      }
    },

    reconnect() {
      if (this.lockReconnect) return;
      this.lockReconnect = true;

      setTimeout(() => {
        console.info("尝试重连...");
        this.initWebSocket();
        this.lockReconnect = false;
      }, 5000);
    },
  },
};

三、优化点解析

1. 安全性增强

  • Token 和用户信息的有效性检查：
  • 在 initWebSocket 方法中增加了对 token 和 userId 的校验，避免因缺失或无效数据导致的错误。
  • 如果 token 或 userId 为空，直接返回并记录错误日志。
  • URL 构造的安全性：
  • 提取了 sanitizeSocketUrl 方法，确保 this.socketUrl 的格式化逻辑不会引入多余字符或导致无效路径。
  • 验证了 window._CONFIG['domianURL'] 的合法性，防止构造出非法的 WebSocket URL。

2. 异常处理改进

  • 捕获初始化过程中的异常：
  • 在 initWebSocket 方法中增加了 try-catch 块，捕获并记录初始化过程中的异常。
  • 标准化错误信息：
  • 替换了 err.code 为更通用的 err.message，确保错误信息的准确性。

3. 性能与可维护性提升

  • 提取重复逻辑：
  • 将 URL 构造逻辑提取为独立方法 sanitizeSocketUrl，简化了主逻辑，提升了代码的可读性和可维护性。
  • 状态管理优化：
  • 在 websocketSend 方法中增加了对 WebSocket 状态的检查，避免无效发送。
  • 使用 lockReconnect 标志位防止重复重连。

4. 边界条件处理

  • 空值检查：
  • 对 token、userId 和 domainUrl 进行了严格的空值检查，避免潜在的运行时错误。
  • 默认值设置：
  • 如果 this.socketUrl 为空，默认设置为 /，确保 URL 构造的完整性。

四、总结

通过以上优化，代码的安全性、健壮性和可维护性均得到了显著提升。具体改进点包括：

  1. 增强了对 token 和 userId 的有效性检查，避免因缺失或无效数据导致的错误。
  2. 提取了 URL 构造逻辑，简化了主逻辑，提升了代码的可读性和可维护性。
  3. 改进了异常处理机制，捕获并记录初始化过程中的异常，确保程序的稳定性。
  4. 优化了重连逻辑，防止重复重连和无限循环。

相关文章