前言
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) } } }
