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