xy_sh-20260727101126/xy_sh/internal/middleware/middleware.go

package middleware

import (
	"encoding/json"
	"fmt"
	"xy_sh/tmpl/errcode"

	"xy_sh/internal/entities"
	"xy_sh/pkg/crypto"

	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/log"
)

func SignVerifyMiddleware() fiber.Handler {
	return func(c *fiber.Ctx) error {
		c.Locals("decryptedData", "")
		timestamp := c.Get("timestamp")
		sign := c.Get("sign")

		if timestamp == "" || sign == "" {
			return fmt.Errorf("缺少timestamp或sign请求头")
		}

		body := c.Body()

		var req entities.EncryptedRequest
		if err := json.Unmarshal(body, &req); err != nil {
			return fmt.Errorf("请求体格式错误")
		}

		if req.EncryptedData == "" {
			return fmt.Errorf("缺少encryptedData参数")
		}

		if !crypto.VerifySign(timestamp, req.EncryptedData, sign) {
			return fmt.Errorf("签名验证失败")
		}

		decryptedData, err := crypto.SM4CBCDecrypt(req.EncryptedData)
		if err != nil {
			return fmt.Errorf("数据解密失败: " + err.Error())
		}

		c.Locals("decryptedData", string(decryptedData))
		c.Locals("timestamp", timestamp)

		return c.Next()
	}
}

func SignVerifyMiddleware2() fiber.Handler {
	return func(c *fiber.Ctx) error {
		timestamp := c.Get("timestamp")
		sign := c.Get("sign")

		if timestamp == "" || sign == "" {
			return c.Status(fiber.StatusUnauthorized).JSON(entities.CommonResponse{
				Code: entities.CodeFailed,
				Msg:  "缺少timestamp或sign请求头",
				Data: nil,
			})
		}

		body := c.Body()

		var req entities.EncryptedRequest
		if err := json.Unmarshal(body, &req); err != nil {
			return c.Status(fiber.StatusBadRequest).JSON(entities.CommonResponse{
				Code: entities.CodeFailed,
				Msg:  "请求体格式错误",
				Data: nil,
			})
		}

		if req.EncryptedData == "" {
			return c.Status(fiber.StatusBadRequest).JSON(entities.CommonResponse{
				Code: entities.CodeFailed,
				Msg:  "缺少encryptedData参数",
				Data: nil,
			})
		}

		if !crypto.VerifySign(timestamp, req.EncryptedData, sign) {
			return c.Status(fiber.StatusUnauthorized).JSON(entities.CommonResponse{
				Code: entities.CodeFailed,
				Msg:  "签名验证失败",
				Data: nil,
			})
		}

		decryptedData, err := crypto.SM4CBCDecrypt(req.EncryptedData)
		if err != nil {
			return c.Status(fiber.StatusBadRequest).JSON(entities.CommonResponse{
				Code: entities.CodeFailed,
				Msg:  "数据解密失败: " + err.Error(),
				Data: nil,
			})
		}

		c.Locals("decryptedData", decryptedData)
		c.Locals("timestamp", timestamp)

		return c.Next()
	}
}

type Res struct {
	Code int         `json:"code"`
	Data interface{} `json:"data"`
	Msg  string      `json:"msg"`
}

func EncryptResponseMiddleware() fiber.Handler {
	return func(c *fiber.Ctx) error {
		err := c.Next()
		if c.Path() == "/api/v1/callback/notify" {
			return nil
		}
		var resp *Res
		// 如果有错误发生
		if err != nil {
			// 返回自定义错误响应
			resp = &Res{
				Msg:  err.Error(),
				Code: entities.CodeFailed,
				Data: nil,
			}
		} else {
			body := c.Response().Body()
			if len(body) == 0 {
				return nil
			}
			resp = &Res{
				Code: entities.CodeSuccess,
				Data: body,
				Msg:  errcode.Success.Error(),
			}
		}

		//var resp entities.CommonResponse
		//if err := json.Unmarshal(res, &resp); err != nil {
		//	return nil
		//}
		var newBody []byte
		if resp.Data != nil {
			dataBytes, err := json.Marshal(resp.Data)
			if err != nil {
				return c.Status(fiber.StatusInternalServerError).JSON(entities.CommonResponse{
					Code: entities.CodeFailed,
					Msg:  "响应数据序列化失败",
					Data: nil,
				})
			}

			encryptedData, err := crypto.SM4CBCEncrypt(dataBytes)
			if err != nil {
				return c.Status(fiber.StatusInternalServerError).JSON(entities.CommonResponse{
					Code: entities.CodeFailed,
					Msg:  "响应数据加密失败",
					Data: nil,
				})
			}

			resp.Data = encryptedData

			newBody, err = json.Marshal(resp)
			if err != nil {
				return c.Status(fiber.StatusInternalServerError).JSON(entities.CommonResponse{
					Code: entities.CodeFailed,
					Msg:  "响应序列化失败",
					Data: nil,
				})
			}

		}

		return c.JSON(newBody)
	}
}

func SetLogReq() fiber.Handler {
	return func(c *fiber.Ctx) error {
		err := c.Next()
		if err != nil {
			return err
		}

		log.Info("Request: ", string(c.Request().Header.Header()), " Body: ", string(c.Request().Body()))
		return nil
	}
}