[golang]gin框架接收websocket通信

前言

WebSocket是一种在单个TCP连接上进行全双工通信的协议。WebSocket让客户端和服务端之间的数据交换变得非常简单，且允许服务器主动向客户端推送数据，并且之后客户端和服务端所有的通信都依靠这个专用协议进行。

本文使用gin框架编写服务端应用，配置路由接收websocket请求并处理。同时实现一个websocket命令行客户端用于与服务端通信。

服务端

下面代码示例中，使用gin创建一个应用，并将自定义函数WebSocketHandler()注册到/ws路由。WebSocketHandler()功能非常简单，客户端发送什么就原样返回什么。

package main

import (
	"fmt"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

func WebSocketHandler(c *gin.Context) {
	// 获取WebSocket连接
	wsUpgrader := websocket.Upgrader{
		HandshakeTimeout: time.Second * 10,
		ReadBufferSize:   1024,
		WriteBufferSize:  1024,
	}
	ws, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer ws.Close()

	// 处理WebSocket消息
	for {
		messageType, p, err := ws.ReadMessage()
		if err != nil {
			fmt.Println(err)
			return
		}
		switch messageType {
		case websocket.TextMessage:
			fmt.Printf("处理文本消息, %s\n", string(p))
			ws.WriteMessage(websocket.TextMessage, p)
            // c.Writer.Write(p)
		case websocket.BinaryMessage:
			fmt.Println("处理二进制消息")
		case websocket.CloseMessage:
			fmt.Println("关闭websocket连接")
			return
		case websocket.PingMessage:
			fmt.Println("处理ping消息")
			ws.WriteMessage(websocket.PongMessage, []byte("ping"))
		case websocket.PongMessage:
			fmt.Println("处理pong消息")
			ws.WriteMessage(websocket.PongMessage, []byte("pong"))
		default:
			fmt.Printf("未知消息类型: %d\n", messageType)
			return
		}
	}

}

func NewServer() *gin.Engine {
	gin.SetMode(gin.DebugMode) // 设置运行模式
	gin.DisableConsoleColor()  // 禁用控制台输出的颜色
	router := gin.Default()
	return router
}

func main() {
	// 创建Gin应用
	app := NewServer()

	// 注册WebSocket路由
	app.GET("/ws", WebSocketHandler)

	// 启动应用
	err := app.Run("127.0.0.1:8080")
	if err != nil {
		panic(err)
	}
}

客户端

写个了命令行客户端用于连接websocket服务端，接收键盘输入，然后发送到服务端。使用flag解析命令行参数用于配置服务端连接。

package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/gorilla/websocket"
)

var (
	Addr  string
	Path  string
	Token string
)

func init() {
	flag.StringVar(&Addr, "addr", "localhost:8080", "WebSocket 服务器地址")
	flag.StringVar(&Path, "path", "/ws", "WebSocket接口路由")
	flag.StringVar(&Token, "token", "123456", "连接 WebSocket 服务器的令牌")
	flag.Parse()
}

func main() {
	header := make(http.Header)
	header.Set("token", Token)
	u := url.URL{Scheme: "ws", Host: Addr, Path: "/ws"}
	conn, _, err := websocket.DefaultDialer.Dial(u.String(), header)
	if err != nil {
		log.Fatalf("连接 WebSocket 服务器失败：%v", err)
		return
	}
	defer conn.Close()

	// 创建channel用于监听操作系统的中断信号
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt, syscall.SIGTERM, syscall.SIGINT)

	response := make(chan string, 8)
	defer close(response)
	// 启动一个 goroutine 用于接收 WebSocket 服务器的响应
	go func(resp chan string) {
		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				log.Printf("server> ERROR! %v\n", err)
				return
			}
			resp <- string(message)
		}
	}(response)

	// 读取用户的键盘输入，并发送到 WebSocket 服务器
	for {
		select {
		case <-interrupt: // 等待中断信号
			log.Println("收到中断信号，关闭 WebSocket 连接 ...")
			err := conn.WriteMessage(
				websocket.CloseMessage, 
				websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
			if err != nil {
				log.Printf("发送关闭消息失败：%v\n", err)
			}
			<-interrupt // 关闭websocket连接之前, 确保已经发送到服务端的消息能够被确认和处理
			return
		default:
			var input string
			fmt.Printf("%s client> ", time.Now().Format("2006-01-02 15:04:05"))
			fmt.Scanln(&input)
			if input == "exit" {
				log.Println("用户输入 exit, 关闭 WebSocket 连接 ...")
				err := conn.WriteMessage(
					websocket.CloseMessage, 
					websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
				if err != nil {
					log.Printf("发送关闭消息失败：%v", err)
					return
				}
				return
			}
			if len(input) == 0 {
				log.Println("输入消息为空")
				continue
			}
			err := conn.WriteMessage(websocket.TextMessage, []byte(input))
			if err != nil {
				log.Printf("发送消息失败：%v", err)
				continue
			}
			// 阻塞等待服务端响应
			resp := <-response
			log.Printf("server> %s\n", resp)
		}
	}
}

参考

• 掘金 - GIn整合WebSocket