From d0fc91b61b25db36a881162e6edb2b450bd59064 Mon Sep 17 00:00:00 2001
From: laodaming <11058467+laudamine@user.noreply.gitee.com>
Date: Tue, 25 Jul 2023 17:41:38 +0800
Subject: [PATCH] fix

---
 constants/websocket.go                        |  3 ++
 .../internal/handler/datatransferhandler.go   |  3 +-
 .../internal/handler/rendernotifyhandler.go   | 46 ++++++++++++++++++-
 server/websocket/internal/types/types.go      |  2 +
 server_api/websocket.api                      |  2 +
 5 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/constants/websocket.go b/constants/websocket.go
index 7fca9751..a9a6aa5d 100644
--- a/constants/websocket.go
+++ b/constants/websocket.go
@@ -13,3 +13,6 @@ const (
 	//图片渲染
 	WEBSOCKET_RENDER_IMAGE = "render-image"
 )
+
+// 云渲染通知需要的签名字符串
+const RENDER_NOTIFY_SIGN_KEY = "fusen-render-notify-%s-%d"
diff --git a/server/websocket/internal/handler/datatransferhandler.go b/server/websocket/internal/handler/datatransferhandler.go
index 9e0363eb..afcaf88d 100644
--- a/server/websocket/internal/handler/datatransferhandler.go
+++ b/server/websocket/internal/handler/datatransferhandler.go
@@ -203,7 +203,8 @@ func (w *wsConnectItem) dealwithReciveData(data []byte) {
 		return
 	}
 	switch parseInfo.MsgType {
-	case constants.WEBSOCKET_RENDER_IMAGE: //图片渲染
+	//图片渲染
+	case constants.WEBSOCKET_RENDER_IMAGE:
 		var renderImageData []types.RenderImageReqMsg
 		if err := json.Unmarshal([]byte(parseInfo.Message), &renderImageData); err != nil {
 			logx.Error("invalid format of websocket render image message", err)
diff --git a/server/websocket/internal/handler/rendernotifyhandler.go b/server/websocket/internal/handler/rendernotifyhandler.go
index 6dde0c3f..5fa2371a 100644
--- a/server/websocket/internal/handler/rendernotifyhandler.go
+++ b/server/websocket/internal/handler/rendernotifyhandler.go
@@ -1,11 +1,16 @@
 package handler
 
 import (
+	"crypto/sha256"
+	"encoding/hex"
 	"encoding/json"
+	"fmt"
+	"fusenapi/constants"
 	"fusenapi/utils/basic"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/rest/httpx"
 	"net/http"
+	"time"
 
 	"fusenapi/server/websocket/internal/svc"
 	"fusenapi/server/websocket/internal/types"
@@ -16,8 +21,45 @@ func RenderNotifyHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
 		var req types.RenderNotifyReq
 		_, err := basic.RequestParse(w, r, svcCtx, &req)
 		if err != nil {
+			httpx.OkJsonCtx(r.Context(), w, basic.Response{
+				Code:    basic.CodeRequestParamsErr.Code,
+				Message: "err param",
+				Data:    nil,
+			})
 			return
 		}
+		if len(req.NotifyList) == 0 {
+			httpx.OkJsonCtx(r.Context(), w, basic.Response{
+				Code:    basic.CodeRequestParamsErr.Code,
+				Message: "invalid param,notify list is empty",
+				Data:    nil,
+			})
+			return
+		}
+		if time.Now().Unix()-120 > req.Time {
+			httpx.OkJsonCtx(r.Context(), w, basic.Response{
+				Code:    basic.CodeRequestParamsErr.Code,
+				Message: "invalid param,time is expired",
+				Data:    nil,
+			})
+			return
+		}
+		//验证签名 sha256
+		notifyByte, _ := json.Marshal(req.NotifyList)
+		h := sha256.New()
+		h.Write([]byte(fmt.Sprintf(constants.RENDER_NOTIFY_SIGN_KEY, string(notifyByte), req.Time)))
+		signHex := h.Sum(nil)
+		sign := hex.EncodeToString(signHex)
+		fmt.Println(sign)
+		if req.Sign != sign {
+			httpx.OkJsonCtx(r.Context(), w, basic.Response{
+				Code:    basic.CodeRequestParamsErr.Code,
+				Message: "invalid sign",
+				Data:    nil,
+			})
+			return
+		}
+		//遍历链接
 		mapConnPool.Range(func(key, value any) bool {
 			ws, ok := value.(wsConnectItem)
 			if !ok {
@@ -40,7 +82,7 @@ func setOutRenderImage(req types.RenderNotifyReq, ws wsConnectItem) {
 	defer ws.mutex.Unlock()
 	for _, notifyItem := range req.NotifyList {
 		renderKey := ws.getRenderImageMapKey(notifyItem.ProductId, notifyItem.SizeId, notifyItem.TemplateId)
-		//加载并删除
+		//查询
 		_, ok := ws.renderImage[renderKey]
 		if !ok {
 			continue
@@ -58,5 +100,7 @@ func setOutRenderImage(req types.RenderNotifyReq, ws wsConnectItem) {
 		case ws.outChan <- b:
 			logx.Info("notify send render result to out chan")
 		}
+		//删掉已经处理的渲染任务
+		delete(ws.renderImage, renderKey)
 	}
 }
diff --git a/server/websocket/internal/types/types.go b/server/websocket/internal/types/types.go
index 35df171d..6082881a 100644
--- a/server/websocket/internal/types/types.go
+++ b/server/websocket/internal/types/types.go
@@ -29,6 +29,8 @@ type RenderImageRspMsg struct {
 }
 
 type RenderNotifyReq struct {
+	Sign       string       `json:"sign"`
+	Time       int64        `json:"time"`
 	NotifyList []NotifyItem `json:"notify_list"`
 }
 
diff --git a/server_api/websocket.api b/server_api/websocket.api
index ee90cdcc..4b5c1fed 100644
--- a/server_api/websocket.api
+++ b/server_api/websocket.api
@@ -39,6 +39,8 @@ type RenderImageRspMsg { //websocket发送渲染完的数据
 }
 //渲染完了通知接口
 type RenderNotifyReq {
+	Sign       string       `json:"sign"`
+	Time       int64        `json:"time"`
 	NotifyList []NotifyItem `json:"notify_list"`
 }
 type NotifyItem {