appNewPlatform/controller/common/websocketApi.go

package common

import (
	"appNewPlatform/controller/common/websocketpack"
	"appNewPlatform/utils/formatoutput"
	"errors"
	"fmt"
	"net/http"
	"time"

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

// 定义WebSocket升级器
var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		//允许跨域请求
		return true
	},
}

/*
*
@ 作者: 秦东
@ 时间: 2025-02-12 10:28:36
@ 功能: 测试通讯
@ 参数

	#

@ 返回值

	#

@ 方法原型

	#
*/
func (a *ApiMethod) WsIndex(c *gin.Context) {
	//升级HTTP链接为WebSocket连接
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		fmt.Printf("WebSocket升级失败:%v\n\n", err)
		return
	}
	defer conn.Close()
	fmt.Printf("WebSocket升级成功:%v\n\n", err)
	for {
		messageType, p, err := conn.ReadMessage()

		fmt.Printf("messageType:%v\n\n", messageType)
		if err != nil {
			fmt.Printf("读取消息失败:%v\n\n", err)
			break
		}
		fmt.Printf("收到消息: %v\n\n", string(p))
		err = conn.WriteMessage(messageType, p)
		if err != nil {
			fmt.Printf("发送消息失败:%v\n\n", err)
			break
		}
	}
}

/*
*
@ 作者: 秦东
@ 时间: 2025-02-12 11:04:58
@ 功能: 根据用户发送信息
@ 参数

	#

@ 返回值

	#

@ 方法原型

	#
*/
func (a *ApiMethod) IdentityLink(c *gin.Context) {
	clientId := c.Query("user_id")
	if clientId == "" {
		formatoutput.Result(200, "未知人员使用WebSocket！请先表明身份！", c)
		return
	}
	clientId = fmt.Sprintf("user_%v", clientId)
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		formatoutput.Result(200, "WebSocket升级失败:", c)
		return
	}
	defer conn.Close()
	//将客户链接添加到存储中
	mutex.Lock()
	Clients[clientId] = conn
	mutex.Unlock()
	fmt.Printf("clientId: %v\n\n", clientId)
	//信息操作
	for {
		_, msgCont, err := conn.ReadMessage()
		if err != nil {
			//客户端断开链接后踢出
			mutex.Lock()
			delete(Clients, clientId)
			mutex.Unlock()
			break
		}
		fmt.Printf("收到消息: %v\n\n", string(msgCont))
		receiveMessages(clientId, msgCont)
	}
}

/*
*
@ 作者: 秦东
@ 时间: 2025-02-12 11:36:57
@ 功能: 接收消息处理
@ 参数

	#

@ 返回值

	#

@ 方法原型

	#
*/
func receiveMessages(clientId string, message []byte) {
	switch clientId {
	case "user_w1":
		SendMsgToClient("user_w2", []byte("你好；user_w2"))
		break
	case "user_w2":
		SendMsgToClient("user_w1", []byte("你好；user_w1"))
		break
	default:
		for i, _ := range Clients {
			SendMsgToClient(i, message)
		}
	}
}

/*
*
@ 作者: 秦东
@ 时间: 2025-02-12 11:25:44
@ 功能: 服务端向指定客户发送信息
@ 参数

	#

@ 返回值

	#

@ 方法原型

	#
*/
func SendMsgToClient(clientId string, message []byte) error {
	mutex.Lock()
	defer mutex.Unlock()
	if client, isOk := Clients[clientId]; isOk {
		err := client.WriteMessage(websocket.TextMessage, message)
		if err != nil { //消息发送失败，说明该链接已经是死链接
			client.Close()
			delete(Clients, clientId)
			return err
		}
	} else {
		errMsg := fmt.Sprintf("未找到客户端ID为 %v 的连接", clientId)
		return errors.New(errMsg)
	}
	return nil
}

/*
========================================================================================================
使用webstock
========================================================================================================
*/

// var (
// 	upgrader = websocket.Upgrader{
// 		// 允许跨域
// 		CheckOrigin: func(r *http.Request) bool {
// 			return true
// 		},
// 	}
// )

func wsHandler(w http.ResponseWriter, r *http.Request) {
	// w.Write([]byte("hello"))
	var (
		wsConn *websocket.Conn
		err    error
		conn   *websocketpack.Connection
		data   []byte
	)
	// 完成ws协议的握手操作
	// Upgrade:websocket
	if wsConn, err = upgrader.Upgrade(w, r, nil); err != nil {
		return
	}

	if conn, err = websocketpack.InitConnection(wsConn); err != nil {
		goto ERR
	}

	// 启动线程，不断发消息
	go func() {
		var (
			err error
		)
		for {
			if err = conn.WriteMessage([]byte("heartbeat")); err != nil {
				return
			}
			time.Sleep(1 * time.Second)
		}
	}()

	for {
		if data, err = conn.ReadMessage(); err != nil {
			goto ERR
		}
		if err = conn.WriteMessage(data); err != nil {
			goto ERR
		}
	}

ERR:
	conn.Close()

}