Go语言实现Websocket服务端

发表于 | 分类于
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>