使用Gin框架搭配WebSocket完成实时聊天

前言

在写项目的时候，需要完成实时聊天的功能，于是简单的学习下WebSocket，想知道WebSocket是什么的小伙伴可以去网上别的地方学习一下。

要实现实时聊天，网上的大部分内容都是SpringBoot和WebSocket完成的，但是我使用Go写的，所以让我们来学习一下Gin框架搭配WebSocket完成实时聊天功能。

实时聊天

实时聊天的数据可以存到数据库，redis中，在这个小demo中没写相应的代码

聊天功能

package main

import (
	"encoding/json"
	"log"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"
)

var (
	// 使用 sync.Map 来保存 WebSocket 连接的映射
	// key 是 userID，value 是 WebSocket 连接
	webSocketMap sync.Map
	// 记录在线人数
	onlineCount int
)

type Message struct {
	FromUserID string `json:"fromUserId"`
	ToUserID   string `json:"toUserId"`
	// 其他字段...
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	upgrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool { return true },
	}

	// 从 URL 中获取 userId
	// userId 为发送者
	// toUserId 为接收者
	userId := r.URL.Query().Get("userId")
	toUserId := r.URL.Query().Get("toUserId")

	// 升级 HTTP 连接为 WebSocket 连接
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("WebSocket upgrade error:", err)
		return
	}

	// 注册 WebSocket 连接
	webSocketMap.Store(userId, conn)
	addOnlineCount()

	log.Println("用户连接:", userId, ",当前在线人数为:", getOnlineCount())

	// 发送消息
	sendMessage(userId, "连接成功")

	// 接收消息并处理
	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			// 关闭连接并从映射中删除
			conn.Close()
			webSocketMap.Delete(userId)
			subOnlineCount()

			log.Println("用户退出:", userId, ",当前在线人数为:", getOnlineCount())
			break
		}

		log.Println("用户消息:", userId, ",报文:", string(message))

		// 解析 JSON
		var msg Message
		if err := json.Unmarshal(message, &msg); err != nil {
			log.Println("Error parsing JSON:", err)
			continue
		}

		// 追加发送人(防止串改)
		msg.FromUserID = userId
		msg.ToUserID = toUserId

		// 传送给对应 ToUserID 用户的 WebSocket
		if msg.ToUserID != "" {
			sendMessage(msg.ToUserID, "收到来自用户 "+userId+" 的消息："+string(message))
		} else {
			sendMessage(userId, "消息发送失败，目标用户不存在")
		}
	}
}

func sendMessage(userID, message string) {
	if conn, ok := webSocketMap.Load(userID); ok {
		conn.(*websocket.Conn).WriteMessage(websocket.TextMessage, []byte(message))
	}
}

func getOnlineCount() int {
	return onlineCount
}

func addOnlineCount() {
	onlineCount++
}

func subOnlineCount() {
	onlineCount--
}

func main() {
	http.HandleFunc("/ws", handleWebSocket)
	http.ListenAndServe(":8080", nil)
}

测试

使用https://websocket.jsonin.com/进行测试（两个网页）

第一个测试URL：ws://127.0.0.1:8080/ws?userId=1&toUserId=2 表示用户1要给用户2发送信息

第二个测试URL：ws://127.0.0.1:8080/ws?userId=2&toUserId=1 表示用户2要给用户1发送信息

第一个连接成功：

第二个连接成功：

控制台信息：

发送信息

now，发送信息测试，在第一个测试中直接点击发送（好像只能发送json格式的，反正我知道数字类型的发送不到第二个测试）

第二个测试中接收到第一个测试中发来的信息

在测试二中同上。

控制台中信息