章
目
录
近期，我在项目复盘时发现，之前搭建的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即时通讯系统的安全性。在实际项目中，大家可以根据具体需求对代码进行进一步优化和扩展。
