goland 實現websocket server的示例代碼

package main

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"net/http"
	uuid "github.com/satori/go.uuid"
)

//Client:單個websocket
type Client struct {
	Id      string
	Socket  *websocket.Conn
	Message chan []byte
}

var clientCount   uint // 客戶端數量

//從websocket中直接讀取數據
func (c *Client) Read() {
	defer func() {
		//客戶端關閉
		if err := c.Socket.Close(); err != nil {
			fmt.Printf("client [%s] disconnect err: %s", c.Id, err)
		}
		//關閉後直接註銷客戶端
		//WebsocketManager.UnRegisterClient(c)
		clientCount--
		fmt.Printf("client [%s],客戶端關閉：[%s],Count [%d]\n", c.Id, websocket.CloseMessage, clientCount)
	}()

	for {
		messageType, message, err := c.Socket.ReadMessage()
		//讀取數據失敗
		if err != nil || messageType == websocket.CloseMessage {
			fmt.Printf("client [%s],數據讀取失敗或通道關閉：[%s],客戶端連接狀態：[%s]\n", c.Id, err.Error(), websocket.CloseMessage)
			break
		}

		// TODO 解析發送過來的參數
		//var data ReadData
		//err = json.Unmarshal(message, &data)
		//if err != nil {
		//  fmt.Println("數據解析失敗")
		//	return
		//}

		// TODO 前端請求返回數據到指定客戶端

		// 簡單測試
		c.Message <- message


	}
}

//寫入數據到websocket中
func (c *Client) Write() {
	defer func() {
		//客戶端關閉
		if err := c.Socket.Close(); err != nil {
			fmt.Printf("client [%s] disconnect err: %s \n", c.Id, err)
			return
		}
		//關閉後直接註銷客戶端
		//WebsocketManager.UnRegisterClient(c)
		clientCount--
		fmt.Printf("client [%s],客戶端關閉：[%s]\n", c.Id, websocket.CloseMessage)
	}()

	for {
		select {
		case message, ok := <-c.Message:

			if !ok {
				//數據寫入失敗，關閉通道
				fmt.Printf("client [%s],客戶端連接狀態：[%s]\n", c.Id, websocket.CloseMessage)
				_ = c.Socket.WriteMessage(websocket.CloseMessage, []byte{})
				//消息通道關閉後直接註銷客戶端
				return
			}

			err := c.Socket.WriteMessage(websocket.TextMessage, message)
			if err != nil {
				fmt.Printf("client [%s] write message err: %s \n", c.Id, err)
				return
			}
		}
	}
}


// 方法二: 通過對象創建  客戶端連接
func WsClient(context *gin.Context) {
	upGrande := websocket.Upgrader{
		//設置允許跨域
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
		//設置請求協議
		Subprotocols: []string{context.GetHeader("Sec-WebSocket-Protocol")},
	}
	//創建連接
	conn, err := upGrande.Upgrade(context.Writer, context.Request, nil)
	if err != nil {
		fmt.Printf("websocket connect error: %s", context.Param("channel"))
		//format.NewResponseJson(context).Error(51001)
		return
	}
	//生成唯一標識client_id
	var uuid = uuid.NewV4().String()
	client := &Client{
		Id:      uuid,
		Socket:  conn,
		Message: make(chan []byte, 1024),
	}
	//註冊
	//ws.WebsocketManager.RegisterClient(client)
	clientCount++

	//起協程，實時接收和回復數據
	go client.Read()
	go client.Write()
}

// 方法一: 直接創建客戶端
func NewConnection(c *gin.Context) {
	// 定義同源檢查，這裡隻作簡單試驗不校驗
	upGrader := websocket.Upgrader{
	   CheckOrigin: func(r *http.Request) bool {
	       return true
	   },
	}
	 ws, err := upGrader.Upgrade(c.Writer, c.Request, nil)
	//ws, err := websocket.Upgrade(c.Writer, c.Request, nil, 1024, 1024)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"msg": "服務端錯誤",
		})
		return
	}


	var message = make(chan []byte)

	go func() {
		defer ws.Close()

		for {
			fmt.Println("start transfer message")
			msgType, msg, err := ws.ReadMessage()
			if err != nil || websocket.CloseMessage == msgType {
				fmt.Println("webSocket read error")
				return
			}
			message <- msg
		}

	}()

	go func() {
		defer ws.Close()

		for {
			mm, ok := <- message

			if !ok {
				//數據寫入失敗，關閉通道
				fmt.Printf("客戶端連接狀態：[%s]\n", websocket.CloseMessage)
				_ = ws.WriteMessage(websocket.CloseMessage, []byte{})
				//消息通道關閉後直接註銷客戶端
				return
			}

			err := ws.WriteMessage(websocket.TextMessage, mm)
			if err != nil {
				fmt.Println("webSocket write error")
				return
			}
		}
	}()

	//for {
	//
	//}

	 開始通信
	//for {
	//	fmt.Println("start transfer message")
	//	msgType, msg, err := ws.ReadMessage()
	//	if err != nil {
	//		fmt.Println("webSocket read error")
	//		return
	//	}
	//	err = ws.WriteMessage(msgType, msg)
	//	if err != nil {
	//		fmt.Println("webSocket write error")
	//		return
	//	}
	//}
}

// ws://127.0.0.1:9090/wsTest
func main() {
	r := gin.Default()

	// 方法一: 直接創建客戶端
	r.GET("/wsTest", NewConnection)
	// 方法二: 通過對象創建  客戶端連接
	r.GET("/wsTest1", WsClient)

	clientCount = 0
	r.Run("127.0.0.1:9090")
}