zmkgC/third/ws/ws.go

// Package ws
// @Author 铁憨憨[cory] 2025/2/12 15:18:00
package ws

import (
	"encoding/json"
	"fmt"
	"github.com/gogf/gf/v2/os/gctx"
	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"github.com/tiger1103/gfast/v3/api/v1/system"
	"github.com/tiger1103/gfast/v3/internal/app/system/service"
	"golang.org/x/net/context"
	"log"
	"math/rand"
	"net/http"
	"strconv"
	"sync"
	"time"
)

// 存储所有连接的设备信息
var connectedDevices = make(map[string]*DeviceInfo)

// 存储每个 uuid 对应的通道
var responseChannels = make(map[string]chan CommonResponse)
var responseChannelsMutex sync.Mutex

// HandleWebSocket 处理WebSocket连接的函数
func HandleWebSocket(w http.ResponseWriter, r *http.Request) {
	ctx := gctx.New()
	// 将HTTP连接升级为WebSocket连接
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("WebSocket升级失败:", err)
		return
	}
	defer conn.Close()

	// 读取设备发送的第一条消息，从中提取 SN然后添加到设备列表中
	declare, err := addDevice(ctx, conn, r)
	if err != nil {
		log.Println("添加设备信息时出错:", err)
		return
	}

	// 持续读取从客户端（其他服务器）发送过来的数据
	for {
		// 读取消息类型和消息内容
		messageType, message, err := conn.ReadMessage()
		if err != nil {
			log.Println("读取消息时出错:", err)
			// 移除设备信息
			delete(connectedDevices, declare.SN)
			log.Printf("设备断开连接，设备信息: %+v", declare)
			err = service.BusAttendanceMachine().Change(ctx, &system.BusAttendanceMachineChangeReq{
				Sn:     declare.SN,
				Status: "0",
			})
			log.Println("修改考勤设备状态时出错:", err)
			break
		}

		// 先解析出 cmd 字段
		var genericMsg GenericMessage
		err = json.Unmarshal(message, &genericMsg)
		if err != nil {
			log.Println("解析 cmd 字段时出错:", err)
			continue
		}

		// 根据 cmd 字段进行不同处理
		switch genericMsg.CMD {
		case DECLARE:
			log.Println("设备在线：", declare)
		case PING:

			declareMessage := DeclareMessage{}
			if json.Unmarshal(message, &declareMessage) != nil {
				return
			}
			if err = handlePing(ctx, conn, r, declareMessage.SN); err != nil {
				log.Println("处理心跳回复:", err)
			}
		case ToClient:
			if err := requestResponse(message); err != nil {
				log.Println("处理响应:", err)
			}
		default:
			log.Printf("收到未知消息---> 类型: %d, 消息内容: %s", messageType, string(message))
		}
	}
	// 清理通道资源
	responseChannelsMutex.Lock()
	for key, ch := range responseChannels {
		if connectedDevices[declare.SN] != nil {
			delete(responseChannels, key)
			close(ch)
		}
	}
	responseChannelsMutex.Unlock()
}

// addDevice 将设备信息添加到设备列表中
func addDevice(ctx context.Context, conn *websocket.Conn, r *http.Request) (declareMessage DeclareMessage, err error) {
	// 读取设备发送的第一条消息，从中提取 SN
	_, message, err := conn.ReadMessage()
	if err != nil {
		return
	}
	declareMessage = DeclareMessage{}
	if json.Unmarshal(message, &declareMessage) != nil {
		return
	}
	// 假设设备信息可以从请求中获取，这里简单从 RemoteAddr 解析
	ip, port := parseRemoteAddr(r.RemoteAddr)
	deviceInfo := &DeviceInfo{
		IP:   ip,
		Port: port,
		Conn: conn,
	}
	if declareMessage.SN != "" {
		// 存储设备信息
		connectedDevices[declareMessage.SN] = deviceInfo
		log.Printf("新设备连接，设备信息: %+v", declareMessage)
		// 变更状态
		err = service.BusAttendanceMachine().Register(ctx, &system.BusAttendanceMachineRegisterReq{
			Sn: declareMessage.SN,
		})
		if err != nil {
			return
		}
	}
	return declareMessage, nil
}

// handlePing 处理 ping 消息
func handlePing(ctx context.Context, conn *websocket.Conn, r *http.Request, sn string) (err error) {
	pongMsg := PongMessageData{
		Cmd: "pong",
	}
	pongJSON, _ := json.Marshal(pongMsg)
	err = conn.WriteMessage(websocket.TextMessage, pongJSON)
	if err != nil {
		log.Println("发送 Pong 消息时出错:", err)
		return err
	}
	// 存储设备信息、变更状态
	connectedDevices[sn] = &DeviceInfo{
		Conn: conn,
	}
	err = service.BusAttendanceMachine().Register(ctx, &system.BusAttendanceMachineRegisterReq{
		Sn: sn,
	})
	return nil
}

// requestResponse 处理请求得到的响应
func requestResponse(message []byte) (err error) {
	log.Println("请求响应：", string(message))
	var common CommonResponse
	err = json.Unmarshal(message, &common)
	if err != nil {
		log.Println("解析 cmd 字段时出错:", err)
		return err
	}
	// 根据UUID查找对应的通道
	responseChannelsMutex.Lock()
	if ch, ok := responseChannels[common.To]; ok {
		ch <- common
		delete(responseChannels, common.To)
	}
	responseChannelsMutex.Unlock()

	return nil
}

// 根据 SN 发送消息给对应的设备
func sendMessageToDevice(sn string, uuid string, message interface{}) error {
	conn, exists := connectedDevices[sn]
	if !exists {
		return nil
	}
	msgJSON, err := json.Marshal(message)
	if err != nil {
		return err
	}

	err = conn.Conn.WriteMessage(websocket.TextMessage, msgJSON)
	if err != nil {
		responseChannelsMutex.Lock()
		delete(responseChannels, uuid)
		responseChannelsMutex.Unlock()
		return err
	}
	return nil
}

func SendRequestAndWaitResponse(sn string, sUuid string, payload interface{}) (CommonResponse, error) {
	responseChan := make(chan CommonResponse, 1)
	// 存入全局映射
	responseChannelsMutex.Lock()
	responseChannels[sUuid] = responseChan
	responseChannelsMutex.Unlock()

	// 发送请求
	err := sendMessageToDevice(sn, sUuid, payload)
	if err != nil {
		return CommonResponse{}, err
	}

	// 等待响应
	select {
	case resp := <-responseChan:
		fmt.Printf("收到响应: %+v\n", resp)
		return resp, nil
	case <-time.After(10 * time.Second):
		responseChannelsMutex.Lock()
		delete(responseChannels, sUuid)
		responseChannelsMutex.Unlock()
		return CommonResponse{}, fmt.Errorf("等待响应超时")
	}
}

/*
=========================================================业务无关=========================================================
=========================================================业务无关=========================================================
=========================================================业务无关=========================================================
*/

// 定义一个升级器，用于将HTTP连接升级为WebSocket连接
var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	// 允许跨域访问，这里设置为允许所有来源
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// 解析远程地址，获取 IP 和端口
func parseRemoteAddr(addr string) (string, string) {
	// 简单处理，假设地址格式为 IP:Port
	for i := len(addr) - 1; i >= 0; i-- {
		if addr[i] == ':' {
			return addr[:i], addr[i+1:]
		}
	}
	return "", ""
}

// GenerateUUIDWithSixRandomDigits 生成一个 UUID 并拼接 6 位随机数
func GenerateUUIDWithSixRandomDigits() string {
	// 生成 UUID
	uuidStr := uuid.New().String()
	// 初始化随机数种子
	rand.Seed(time.Now().UnixNano())
	// 生成 6 位随机数
	randomNum := rand.Intn(900000) + 100000
	// 将随机数转换为字符串
	randomNumStr := strconv.Itoa(randomNum)
	// 拼接 UUID 和 6 位随机数
	result := uuidStr + "-" + randomNumStr
	return result
}