golang使用Websocket实例

本项目使用 github.com/gorilla/websocket 库 安装：

go get github.com/gorilla/websocket

服务端代码

package main

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

func WebSocketServer() {
	addr := "localhost:8002"
	http.HandleFunc("/wshandler", WebSocketUpgrade)
	log.Println("Starting websocket server at " + addr)

	go func() {
		err := http.ListenAndServe(addr, nil)
		if err != nil {
			log.Fatal(err)
		}
	}()

	log.Println("WebSocket 服务器正在运行。按Ctrl+C退出")
	select {}
}

func WebSocketUpgrade(resp http.ResponseWriter, req *http.Request) {
	// 初始化 Upgrader
	upgrader := websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	} // 使用默认的选项
	// 第三个参数是响应头,默认会初始化
	conn, err := upgrader.Upgrade(resp, req, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()

	// 读取客户端的发送额消息,并返回
	fmt.Println("websocket 连接成功")
	// fmt.Println(conn)
	go ReadMessage(conn)
	select {}
}

// 读取客户端发送的消息,并返回
func ReadMessage(conn *websocket.Conn) {
	for {
		// 消息类型:文本消息和二进制消息
		messageType, msg, err := conn.ReadMessage()
		if err != nil {
			log.Println(err)
			return
		}
		fmt.Println("receive msg:", string(msg))

		err = conn.WriteMessage(messageType, []byte("The server has received the message"))
		if err != nil {
			log.Println("write error:", err)
			return
		}
	}
}

func main() {
	WebSocketServer()
}

客户端

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"time"
	"github.com/gorilla/websocket"
)

func main() {
	wsUrl := "ws://localhost:8002/wshandler"
	// 连接WebSocket服务器
	conn, _, err := websocket.DefaultDialer.Dial(wsUrl, nil)
	if err != nil {
		log.Println(err)
		return
	}

	log.Println("连接成功")

	go func() {
		for {
			// 读取消息
			messageType, message, err := conn.ReadMessage()
			if err != nil {
				log.Fatal(err)
			}
			log.Println("消息类型:", messageType, "消息数据：", string(message))
		}
	}()

	// 发送消息
	controlMessage := map[string]interface{}{
		"type": "control",
		"data": "hello word",
	}
	jsonData, err := json.Marshal(controlMessage)
	if err != nil {
		fmt.Println("Error marshalling JSON:", err)
		return
	}

	// 间隔2秒发送一次消息
	ticket := time.NewTicker(2 * time.Second)
	defer ticket.Stop()
	for {
		<-ticket.C
		// 2. 发送消息
		err = conn.WriteMessage(websocket.TextMessage, []byte(string(jsonData)))
		if err != nil {
			log.Fatal(err)
		}
		log.Println("发送消息成功")
	}

}