WebSocket即时通讯：前后端设计与基于token的鉴权实现

文
章
目
录

• 一、后端实现
• （一）路由鉴权代码解析
• （二）WebSocket连接处理代码
• 二、前端实现

近期，我在项目复盘时发现，之前搭建的WebSocket即时通讯功能存在权限验证缺失的问题，仅实现了简单的连接与消息发送功能。在面试过程中被问及WebSocket鉴权相关问题后，我决定深入研究并补上这一安全漏洞，下面和大家分享一下具体的实现过程。

一、后端实现

（一）路由鉴权代码解析

在后端使用Go语言进行开发时，路由鉴权是实现WebSocket权限验证的重要一步。代码如下：

defaultRoutes.GET(\"/ws\", func(ctx *gin.Context) {
    // 从请求参数中获取token
    t := ctx.Query(\"token\")
    // 解析token，获取token对象以及可能的错误信息
    token, _, err := middlewares.ParseToken(t)
    // 判断token是否有效，若无效则返回错误信息
    if err != nil ||!token.Valid {
        ctx.JSON(400, gin.H{
            \"message\": \"token无效\",
        })
    } else {
        // 若token有效，调用UserController的WS方法处理WebSocket连接
        controllers.UserController{}.WS(ctx.Writer, ctx.Request)
    }
})

这段代码的核心逻辑是，从请求的URL参数中提取token，然后通过middlewares.ParseToken函数对token进行解析和验证。如果token无效，比如解析出错或者本身不合法，就返回一个包含错误信息的JSON响应，告知客户端token无效。只有当token有效时，才会继续调用UserController的WS方法来处理WebSocket连接，从而确保只有通过鉴权的用户才能建立WebSocket连接。

（二）WebSocket连接处理代码

接下来是WebSocket连接相关的代码：

var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}
var conns []*websocket.Conn

这里定义了一个upgrader对象，用于将HTTP连接升级为WebSocket连接。CheckOrigin函数设置为始终返回true，表示允许来自任何源的连接。同时，定义了一个conns切片，用于存储所有已建立的WebSocket连接。

func (this UserController) WS(w http.ResponseWriter, r *http.Request) {
    // 将HTTP连接升级为WebSocket连接
    c, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        // 若升级过程出错，打印错误信息并返回
        println(\"upgrade错误:\", err)
        return
    }
    // 在函数结束时关闭WebSocket连接，确保资源正确释放
    defer c.Close()
    // 将新建立的连接添加到连接列表中
    conns = append(conns, c)
    for {
        // 持续读取WebSocket连接上的消息
        _, _, err := c.ReadMessage()
        if err != nil {
            // 若读取消息出错，打印错误信息并跳出循环
            println(\"read:\", err)
            break
        }
    }
}

在WS方法中，首先使用upgrader.Upgrade将HTTP连接升级为WebSocket连接。如果升级过程出现错误，打印错误信息并返回。成功升级后，将该连接添加到conns列表中，方便后续管理。然后进入一个无限循环，持续读取连接上的消息，一旦读取消息出现错误，打印错误信息并结束循环，关闭连接。

二、前端实现

前端部分使用JavaScript来建立WebSocket连接并进行相关操作。代码如下：

// 从sessionStorage中获取token
let token=sessionStorage.getItem(\"token\")
// 判断当前环境是开发环境还是生产环境
const env = process.env.NODE_ENV
// 根据不同环境构建WebSocket连接的URL，并将token作为参数附带上
const url = env == \'development\' ? \"ws://localhost:8088/ws?token=\" + token : \"ws://114.116.249.103:8088/ws?token=\" + token
// 创建WebSocket实例
const websocket = new WebSocket(url)
// 定义一个响应式变量，用于跟踪WebSocket连接状态
let socketState = ref(true)
// 连接成功的回调函数
websocket.onopen = (evt) => {
    console.log(\"链接成功\")
    socketState.value = true
}
// 接收到消息的回调函数
websocket.onmessage = (evt) => {
    // 根据接收到的消息内容执行不同的自定义功能
    if (evt.data == \"xxx1\") {
        refreshChartJL()
    } else if (evt.data == \"xxx2\") {
        refreshChartMusic()
    }
}
// 连接关闭的回调函数
websocket.onclose = () => {
    console.log(\"链接关闭\")
    socketState.value = false
}

在前端代码中，首先从sessionStorage中获取存储的token。然后根据当前的环境变量NODE_ENV来确定WebSocket连接的URL，在开发环境下连接到本地服务器，在生产环境下连接到正式服务器，并将token作为参数添加到URL中。接着创建WebSocket实例，并定义了onopen、onmessage和onclose三个回调函数，分别用于处理连接成功、接收到消息和连接关闭的情况。在onmessage回调函数中，根据接收到的不同消息内容，调用相应的自定义函数进行处理。

通过以上前后端的代码实现，完成了基于token的WebSocket鉴权功能，有效提升了WebSocket即时通讯系统的安全性。在实际项目中，大家可以根据具体需求对代码进行进一步优化和扩展。