Golang Gin 引入Gorilla WebSocket 实现指定连接通信

Golang Gin 引入Gorilla WebSocket 实现指定连接通信

Gorilla引入：

github.com/gorilla/websocket

代码完全时个人逻辑，存在一些问题，目前没有考虑线程锁，测试效果仅供娱乐。

package service

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

//待优化：加入 缓存、线程锁、MQ

type message struct {
	msg         []byte
	massageType int
	msgSize     int
}

type session struct {
	Request    *http.Request
	conn       *websocket.Conn
	outputDone chan struct{}
	open       bool
	rwmutex    *sync.RWMutex
}

// SessionMap 保存WebSocket链接
var SessionMap = make(map[string]session)

var keyMap = make(map[*websocket.Conn]string)

func WsGorilla(ctx *gin.Context) {
	//http 协议升级
	upgraded := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	Conn, err := upgraded.Upgrade(ctx.Writer, ctx.Request, nil)
	addSession(Conn, ctx.Request, "小明")
	if err != nil {
		return
	}
	var msg []byte
	var messageType int
	//读取并发送消息
	for {
		messageType, msg, err = Conn.ReadMessage()
		if err != nil {
			_ = Conn.Close()
			//清除会话信息
			clearSession(Conn)
			return
		} else if msg != nil {
			err := SendMessage(messageType, msg, "小明")
			if err != nil {
				return
			}
		}
	}
}

// 添加会话
func addSession(coon *websocket.Conn, request *http.Request, key string) {
	fmt.Println("addSession")
	SessionMap[key] = session{
		conn:    coon,
		Request: request,
	}
	keyMap[coon] = key
}

// 清除指定会话
func clearSession(coon *websocket.Conn) {
	delete(SessionMap, keyMap[coon])
	delete(keyMap, coon)
}

func SendMessage(messageType int, msg []byte, key string) error {
	conn := SessionMap[key].conn
	fmt.Println("SendMessage")
	err := conn.WriteMessage(messageType, msg)
	if msg = nil; err != nil {
		return err
	}
	return nil
}