WebSocket是一种在单个TCP连接上进行全双工通信的协议。该协议于2011年被IETF定为标准RFC 6455，并由RFC7936补充规范。WebSocket API也被W3C定为标准。这种通信方式使得客户端和服务器之间的数据交换变得更加简单，允许服务端主动向客户端推送数据。在WebSocket API中，客户端和服务器只需要完成一次握手，之后两者之间就可以直接创建持久性的连接，并进行双向数据传输。
这篇文章我就准备使用不到200行Go代码使用WebSocket实现一个简单的私聊系统，正文开始~
首先，我们需要安装gorilla/websocket包，它是Go语言中一个非常流行的WebSocket库：
go get -u github.com/gorilla/websocket其实发送私聊消息无非就是解决两个问题：
延伸一下，在我们的功能设计中就分为两步，一是确认身份，就需要客户端有一个向服务端注册的动作，二是发消息，而且在发消息的同时除了内容以外，还要知道发送者和接受者都是谁。
作为私聊系统的第一步，客户端需要先与服务端建立WebSocket连接，抽象出User和用户连接池，在注册时发送消息然后由客户端保存该User对应的WebSocket连接。
<img src="https://files.mdnice.com/user/56899/bd4a8194-c476-4fd4-8d0c-0406a71ce0e8.png" style="zoom:67%;" />
消息发送时要使用固定的结构体，使用JSON格式进行序列化，首先是客户端进行消息的构建，然后发送到服务器，此时服务器就类似于一个路由器，将之前保存的接收方的WebSocket连接拿出来进行消息的发送。
下面我们进行代码的实现。
首先是结构体，也就是消息类型的定义，Chat作为客户端和服务器交互的统一类型，承载了事件类型（注册或发送消息事件）和消息体，消息体中有接受者、发送者和消息内容等主要字段。
package model
type Event int
const (
Register Event = 1
SendMsg = 2
)
type Chat struct {
Event Event
Message Message
User User
}
type Message struct {
SendUser *User
Receiver *User
Content string
CreateTime string
}
type User struct {
UserId int64
UserName string
Address string
}服务端代码主要是提供一个WebSocket连接，根据接收的消息事件取出相关的消息，再进行后续的注册或发送消息逻辑。
var (
ws = websocket.Upgrader{}
userMap map[int64]*websocket.Conn
)
func init() {
userMap = make(map[int64]*websocket.Conn)
}
func main() {
http.HandleFunc("/privateChat", privateChat)
_ = http.ListenAndServe(":9900", nil)
}
func privateChat(w http.ResponseWriter, r *http.Request) {
c, err := ws.Upgrade(w, r, nil) //升级将 HTTP 服务器连接升级到 WebSocket 协议
if err != nil {
log.Printf("upgrade err:%s\n", err)
return
}
for {
mt, message, err := c.ReadMessage()
if err != nil {
log.Printf("read message err:%s\n", err)
continue
}
var chat model.Chat
if err := json.Unmarshal(message, &chat); err != nil {
log.Printf("unmarshal err:%s\n", err)
continue
}
switch chat.Event {
case model.Register:
//设置userID
chat.User.UserId = time.Now().Unix()
//保存用户连接
userMap[chat.User.UserId] = c
break
case model.SendMsg:
chat.Message.CreateTime = time.Now().Format("2006-01-02 15:04:05")
//拿到用户连接
ok := false
c, ok = userMap[chat.Message.Receiver.UserId]
if !ok {
//如果没有，拿到发送方用户的连接，告诉他不行
c, _ = userMap[chat.Message.SendUser.UserId]
chat.Message.Receiver = chat.Message.SendUser
chat.Message.Content = "发送失败"
}
break
default:
c, _ = userMap[chat.Message.SendUser.UserId]
chat.Message.Receiver = chat.Message.SendUser
chat.Message.Content = "消息类型不对"
}
log.Printf("now chat : %+v \n", chat)
//响应数据
bytes, err := json.Marshal(chat)
if err != nil {
log.Printf("marshal err:%s\n", err)
continue
}
if err := c.WriteMessage(mt, bytes); err != nil {
log.Printf("write message err:%s\n", err)
continue
}
}
}客户端代码首先是连接WebSocket服务端，然后再启动一个HTTP服务用于接收和发送具体的WebSocket消息。
const (
PrivateChatUrl = "ws://127.0.0.1:9900/privateChat"
)
var (
user model.User
name string
port int
)
func init() {
flag.StringVar(&name, "name", "", "user name")
flag.IntVar(&port, "port", 8801, "server port")
}
func main() {
flag.Parse()
c, _, err := websocket.DefaultDialer.Dial(PrivateChatUrl, nil)
if err != nil {
log.Fatal("dial:", err)
return
}
//注册
u := model.Chat{Event: model.Register, User: model.User{UserName: name}}
bytes, _ := json.Marshal(u)
_ = c.WriteMessage(websocket.TextMessage, bytes)
//读取消息监听
go func() {
for {
_, msg, err := c.ReadMessage()
if err != nil {
log.Printf("read message err:%s\n", err)
continue
}
var res model.Chat
if err := json.Unmarshal(msg, &res); err != nil {
log.Printf("unmarshal err:%s\n", err)
continue
}
switch res.Event {
case model.Register:
user = res.User
fmt.Printf("register success , now user %v \n", user)
case model.SendMsg:
fmt.Printf("用户 %s 在 %s 给你发送了一条消息:%s \n",
res.Message.SendUser.UserName, res.Message.CreateTime, res.Message.Content)
}
}
}()
http.HandleFunc("/send", func(w http.ResponseWriter, r *http.Request) {
uid := r.URL.Query().Get("uid")
content := r.URL.Query().Get("content")
msg, _ := json.Marshal(model.Chat{
Event: model.SendMsg,
Message: model.Message{
SendUser: &user,
Receiver: &model.User{UserId: cast.ToInt64(uid)},
Content: content,
},
})
_ = c.WriteMessage(websocket.TextMessage, msg)
_, _ = w.Write([]byte("ok"))
})
_ = http.ListenAndServe(fmt.Sprintf(":%d",port), nil)
}首先启动服务端：
go run server/main.go然后启动客户端，命名为ZhangSan，HTTP端口号为8801，进行注册：
go run client/main.go -port 8801 -name ZhangSan得到的相应结果：
register success , now user {1713606616 ZhangSan }然后启动另一个客户端，命名为LiSi，HTTP端口号为8802，进行注册：
go run client/main.go -port 8802 -name LiSi得到的相应结果：
register success , now user {1713606595 LiSi }现在我们知道了两个用户的uId，然后我们可以进行发送消息：
上图是用户LiSi给用户ZhangSan发送消息，我们查看接收方的响应：
完成~
不知道上面的代码大家有没有看出问题，没错，Go语言的map类型是线程不安全的，因此最好进行在操作时进行加锁。
//Lock
userMap[userId] = WebSocketConn //具体操作
//Unlock 除此之外，还有许多功能可以扩展，希望大家能给出建议~
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。
