Go语言实现Websocket服务端

项目初始化

创建项目xh-control-ws

进入目录下进行初始化

1
2
go mod init xh-control-ws
go mod tidy

安装依赖

安装依赖库

1
go get github.com/gorilla/websocket

基本示例

创建文件main.go

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
package main

import (
"github.com/gorilla/websocket"
"log"
"net/http"
)

var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
// 解决跨域问题
CheckOrigin: func(r *http.Request) bool {
return true
},
}

func main() {
http.HandleFunc("/ws", echoHandler)
log.Fatal(http.ListenAndServe(":10088", nil))
}

func echoHandler(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Println(err)
return
}

defer conn.Close()
// 设置关闭处理函数
conn.SetCloseHandler(func(code int, text string) error {
//处理下线逻辑
log.Printf("Connection closed [%s]\n", conn.RemoteAddr().String())
return nil
})
log.Printf("Connection open [%s]:\n", conn.RemoteAddr().String())
for {
messageType, message, err := conn.ReadMessage()
if err != nil {
if strings.Contains(err.Error(), "websocket: close") {
//处理下线逻辑
log.Printf("Connection closed [%s]\n", conn.RemoteAddr().String())
}else{
log.Println(err)
}
break
}
if messageType == websocket.TextMessage {
//接收消息
log.Printf("Received [%s]:\nmessage:%s", conn.RemoteAddr().String(), message)
//发送消息
sendMsg := "服务器返回:" + string(message)
err = conn.WriteMessage(messageType, []byte(sendMsg))
log.Printf("Send [%s]:\nmessage:%s", conn.RemoteAddr().String(), sendMsg)
if err != nil {
log.Println(err)
break
}
}
}
}

注意在读消息的时候判断了错误信息,当用户下线的时候也会有错误 ,我们额外判断一下。

1
2
3
4
5
6
7
8
9
10
messageType, message, err := conn.ReadMessage()
if err != nil {
if strings.Contains(err.Error(), "close") {
//处理下线逻辑
log.Printf("Connection closed [%s]\n", conn.RemoteAddr().String())
}else{
log.Println(err)
}
break
}

注意

SetCloseHandler只能监听到正常关闭的事件,非正常关闭的可以使用上面的方式判断

协程发送消息

我们想到的是这样的,但是这样会报错

错误的写法:

1
2
3
4
5
6
7
8
9
func sendMsg(conn *websocket.Conn, msg string) {
go func() {
err := conn.WriteMessage(websocket.TextMessage, []byte(msg))
if err != nil {
return
}
log.Printf("Send [%s]:\nmessage:%s", conn.RemoteAddr().String(), msg)
}()
}

正确的写法:

添加一个对象

1
2
3
4
type SendMsgModel struct {
Conn *websocket.Conn `json:"conn"` //概要
Msg string `json:"msg"` //用户ID
}

代码中

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
var broadcast = make(chan model.SendMsgModel)

func main() {
http.HandleFunc("/ws", echoHandler)
// 协程调用
go handleMessages()
log.Fatal(http.ListenAndServe(":10088", nil))
}

func handleMessages() {
for {
// 从消息通道中读取消息
message := <-broadcast
err := message.Conn.WriteMessage(websocket.TextMessage, []byte(message.Msg))
if err != nil {
fmt.Println(err)
}
}
}

func sendMsg(conn *websocket.Conn, msg string) {
// 要发送的消息放入通道中
broadcast <- model.SendMsgModel{Conn: conn, Msg: msg}
}

连接的标识

1
fmt.Sprintf("%p", conn)

连接的IP及端口

1
conn.RemoteAddr().String()

这个在服务器上获取到的远程的IP一直是127.0.0.1,即使Nginx配置了header也不起作用,所以不建议作为连接的唯一标识。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
map $http_upgrade $connection_upgrade {
default upgrade;
'' close;
}

server {
location / {
proxy_pass http://mcwstest;
proxy_http_version 1.1;
proxy_set_header Upgrade $http_upgrade;
proxy_set_header Connection $connection_upgrade;
proxy_set_header Host $host;
proxy_set_header X-Real-IP $remote_addr;
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
}
}

前端测试

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>WS测试</title>
</head>
<body>
<script>
let ws = new WebSocket("ws://localhost:10088/ws");
ws.onopen = function () {
console.log("Connected");
ws.send("Hello, World!");
};

ws.onmessage = function (event) {
console.log("Received message: " + event.data);
};

ws.onclose = function () {
console.log("Disconnected");
};
</script>
</body>
</html>