Go语言实现Websocket服务端

项目初始化

创建项目xh-control-ws

进入目录下进行初始化

go mod init xh-control-ws
go mod tidy

安装依赖

安装依赖库

go get github.com/gorilla/websocket

基本示例

创建文件main.go

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]:\ncode %d and reason: %s\n", conn.RemoteAddr().String(), code, text)
		return nil
	})
	log.Printf("Connection open [%s]:\n", conn.RemoteAddr().String())
	for {
		messageType, message, err := conn.ReadMessage()
		if err != nil {
			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
			}
		}
	}
}

协程发送消息

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

错误的写法：

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)
	}()
}

正确的写法：

添加一个对象

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

代码中

var broadcast = make(chan model.SendMsgModel)

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

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}
}

如果想多个协程处理，handleMessages()调用多次即可，是不会导致处理信息重复的。

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

连接相关

连接的标识

conn.RemoteAddr().String()

连接关闭的回调

// 设置关闭处理函数
conn.SetCloseHandler(func(code int, text string) error {
    // 连接关闭时的处理逻辑
    log.Printf("Connection closed [%s]:\ncode %d and reason: %s\n", conn.RemoteAddr().String(), code, text)
    return nil
})

前端测试

<!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>

类型转换

在Go语言中，可以使用[]byte和string类型之间进行转换。

这两种类型之间的转换可以通过类型转换或者使用标准库中的函数来完成。

[]byte到string的转换：

可以使用string()函数将[]byte转换为string：

byteSlice := []byte{'H', 'e', 'l', 'l', 'o'}
str := string(byteSlice)

string到[]byte的转换：

可以使用[]byte()函数将string转换为[]byte：

str := "Hello"
byteSlice := []byte(str)

需要注意的是，在Go中，string是不可变的，而[]byte是可变的。

因此，将string转换为[]byte后，可以修改[]byte中的内容，但是不能直接修改string的内容。

示例：

str := "Hello"
byteSlice := []byte(str)
byteSlice[0] = 'h'
modifiedStr := string(byteSlice)
fmt.Println(modifiedStr) // 输出 "hello"

但是，直接修改string中的内容是不被允许的：

str := "Hello"
// 下面这行代码会导致编译错误，因为string的内容不可修改
// str[0] = 'h'

这些方法可以方便地在string和[]byte之间进行转换。