From 5f885a6369e660a4de512f79c5059613a621f4a7 Mon Sep 17 00:00:00 2001
From: TsMask <340112800@qq.com>
Date: Sat, 31 Aug 2024 14:22:44 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=8E=A5=E5=8F=A3=E5=8A=A0=E8=A7=A3?=
 =?UTF-8?q?=E5=AF=86=E4=B8=AD=E9=97=B4=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/framework/constants/result/result.go |   9 +-
 src/framework/errorcatch/errorcatch.go   |   6 +-
 src/framework/middleware/crypto_api.go   | 152 +++++++++++++++++++++++
 src/framework/middleware/report.go       |   7 +-
 src/framework/utils/crypto/aes.go        |  33 +++--
 5 files changed, 183 insertions(+), 24 deletions(-)
 create mode 100644 src/framework/middleware/crypto_api.go

diff --git a/src/framework/constants/result/result.go b/src/framework/constants/result/result.go
index dbd8630..19761ce 100644
--- a/src/framework/constants/result/result.go
+++ b/src/framework/constants/result/result.go
@@ -8,8 +8,13 @@ const (
 	// 响应-msg错误失败
 	MSG_ERROR = "error"
 
-	// 响应-msg正常成功
-	CODE_SUCCESS = 1
 	// 响应-code正常成功
+	CODE_SUCCESS = 1
+	// 响应-msg正常成功
 	MSG_SUCCESS = "success"
+
+	// 响应-code加密数据
+	CODE_ENCRYPT = 2
+	// 响应-msg加密数据
+	MSG_ENCRYPT = "encrypt"
 )
diff --git a/src/framework/errorcatch/errorcatch.go b/src/framework/errorcatch/errorcatch.go
index b196311..0a3583e 100644
--- a/src/framework/errorcatch/errorcatch.go
+++ b/src/framework/errorcatch/errorcatch.go
@@ -20,14 +20,14 @@ func ErrorCatch() gin.HandlerFunc {
 
 				// 返回错误响应给客户端
 				if config.Env() == "prod" {
-					c.JSON(500, result.ErrMsg("Internal Server Errors"))
+					c.JSON(500, result.CodeMsg(500, "Internal Server Errors"))
 				} else {
 					// 通过实现 error 接口的 Error() 方法自定义错误类型进行捕获
 					switch v := err.(type) {
 					case error:
-						c.JSON(500, result.ErrMsg(v.Error()))
+						c.JSON(500, result.CodeMsg(500, v.Error()))
 					default:
-						c.JSON(500, result.ErrMsg(fmt.Sprint(err)))
+						c.JSON(500, result.CodeMsg(500, fmt.Sprint(err)))
 					}
 				}
 
diff --git a/src/framework/middleware/crypto_api.go b/src/framework/middleware/crypto_api.go
new file mode 100644
index 0000000..b5b1862
--- /dev/null
+++ b/src/framework/middleware/crypto_api.go
@@ -0,0 +1,152 @@
+package middleware
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"strings"
+
+	"nms_cxy/src/framework/config"
+	constResult "nms_cxy/src/framework/constants/result"
+	"nms_cxy/src/framework/logger"
+	"nms_cxy/src/framework/utils/crypto"
+	"nms_cxy/src/framework/utils/parse"
+
+	"github.com/gin-gonic/gin"
+)
+
+// CryptoApi 接口加解密
+//
+// 示例参数：middleware.CryptoApi(true, true)
+//
+// 参数表示：对请求解密，对响应加密
+//
+// 请将中间件放在最前置，对请求优先处理
+func CryptoApi(requestDecrypt, responseEncrypt bool) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		// 请求解密时对请求data注入
+		if requestDecrypt {
+			method := c.Request.Method
+			contentType := c.ContentType()
+			contentDe := ""
+			// 请求参数解析
+			if method == "GET" {
+				contentDe = c.Query("data")
+			} else if contentType == gin.MIMEJSON {
+				var body struct {
+					Data string `json:"data" binding:"required"`
+				}
+				if err := c.ShouldBindJSON(&body); err == nil {
+					contentDe = body.Data
+				}
+			}
+
+			// 是否存在data字段数据
+			if contentDe == "" {
+				c.JSON(400, map[string]any{
+					"code": constResult.CODE_ERROR,
+					"msg":  "decrypt not found field data",
+				})
+				c.Abort() // 停止执行后续的处理函数
+				return
+			}
+
+			// 解密-原数据加密前含16位长度iv
+			apiKey := config.Get("aes.apiKey").(string)
+			dataBodyStr, err := crypto.AESDecryptBase64(contentDe, apiKey)
+			if err != nil {
+				logger.Errorf("CryptoApi decrypt err => %v", err)
+				c.JSON(400, map[string]any{
+					"code": constResult.CODE_ERROR,
+					"msg":  "decrypted data could not be parsed",
+				})
+				c.Abort() // 停止执行后续的处理函数
+				return
+			}
+
+			// 分配回请求体
+			if method == "GET" {
+				var urlParams map[string]any
+				json.Unmarshal([]byte(dataBodyStr), &urlParams)
+				rawQuery := []string{}
+				for k, v := range urlParams {
+					rawQuery = append(rawQuery, fmt.Sprintf("%s=%v", k, v))
+				}
+				c.Request.URL.RawQuery = strings.Join(rawQuery, "&")
+			} else if contentType == gin.MIMEJSON {
+				c.Request.Body = io.NopCloser(bytes.NewBuffer([]byte(dataBodyStr)))
+			}
+		}
+
+		// 响应加密时替换原有的响应体
+		var rbw *replaceBodyWriter
+		if responseEncrypt {
+			rbw = &replaceBodyWriter{
+				body:           &bytes.Buffer{},
+				ResponseWriter: c.Writer,
+			}
+			c.Writer = rbw
+		}
+
+		// 调用下一个处理程序
+		c.Next()
+
+		// 响应加密时对响应data数据进行加密
+		if responseEncrypt {
+			// 满足成功并带数据的响应进行加密
+			if c.Writer.Status() == 200 {
+				var resBody map[string]any
+				json.Unmarshal(rbw.body.Bytes(), &resBody)
+				codeV, codeOk := resBody["code"]
+				dataV, dataOk := resBody["data"]
+				if codeOk && dataOk {
+					if parse.Number(codeV) == constResult.CODE_SUCCESS {
+						byteBodyData, _ := json.Marshal(dataV)
+						// 加密-原数据头加入标记16位长度iv终止符
+						apiKey := config.Get("aes.apiKey").(string)
+						contentEn, err := crypto.AESEncryptBase64("=:)"+string(byteBodyData), apiKey)
+						if err != nil {
+							logger.Errorf("CryptoApi encrypt err => %v", err)
+							rbw.ReplaceWrite([]byte(fmt.Sprintf(`{"code":"%d","msg":"encrypt err"}`, constResult.CODE_ERROR)))
+						} else {
+							// 响应加密
+							byteBody, _ := json.Marshal(map[string]any{
+								"code": constResult.CODE_ENCRYPT,
+								"msg":  constResult.MSG_ENCRYPT,
+								"data": contentEn,
+							})
+							rbw.ReplaceWrite(byteBody)
+						}
+					}
+				} else {
+					rbw.ReplaceWrite(nil)
+				}
+			} else {
+				rbw.ReplaceWrite(nil)
+			}
+		}
+		//
+	}
+}
+
+// replaceBodyWriter 替换默认的响应体
+type replaceBodyWriter struct {
+	gin.ResponseWriter
+	body *bytes.Buffer
+}
+
+// Write 写入响应体
+func (r replaceBodyWriter) Write(b []byte) (int, error) {
+	return r.body.Write(b)
+}
+
+// ReplaceWrite 替换响应体
+func (r *replaceBodyWriter) ReplaceWrite(b []byte) (int, error) {
+	if b == nil {
+		return r.ResponseWriter.Write(r.body.Bytes())
+	}
+	r.body = &bytes.Buffer{}
+	r.body.Write(b)
+	return r.ResponseWriter.Write(r.body.Bytes())
+}
diff --git a/src/framework/middleware/report.go b/src/framework/middleware/report.go
index b700a5e..9dde09e 100644
--- a/src/framework/middleware/report.go
+++ b/src/framework/middleware/report.go
@@ -1,6 +1,7 @@
 package middleware
 
 import (
+	"runtime"
 	"time"
 
 	"nms_cxy/src/framework/logger"
@@ -18,6 +19,10 @@ func Report() gin.HandlerFunc {
 
 		// 计算请求处理时间，并打印日志
 		duration := time.Since(start)
-		logger.Infof("%s %s report end=> %v", c.Request.Method, c.Request.RequestURI, duration)
+		// logger.Infof("%s %s report end=> %v", c.Request.Method, c.Request.RequestURI, duration)
+		// 获取当前活跃的goroutine数量
+		num := runtime.NumGoroutine()
+		// logger.Infof("当前活跃的goroutine数量 %d\n", num)
+		logger.Infof("\n访问接口 %s %s\n总耗时 %v\n当前活跃的goroutine数量 %d\n", c.Request.Method, c.Request.RequestURI, duration, num)
 	}
 }
diff --git a/src/framework/utils/crypto/aes.go b/src/framework/utils/crypto/aes.go
index 290eb71..fd39516 100644
--- a/src/framework/utils/crypto/aes.go
+++ b/src/framework/utils/crypto/aes.go
@@ -10,12 +10,12 @@ import (
 	"io"
 )
 
-// StringEncryptByAES 字符串AES加密
-func StringEncryptByAES(text string) (string, error) {
+// AESEncryptBase64 AES加密转Base64字符串
+func AESEncryptBase64(text, key string) (string, error) {
 	if len(text) == 0 {
 		return "", nil
 	}
-	xpass, err := aesEncryptWithSalt([]byte(text))
+	xpass, err := AESEncrypt([]byte(text), []byte(key))
 	if err != nil {
 		return "", err
 	}
@@ -23,8 +23,8 @@ func StringEncryptByAES(text string) (string, error) {
 	return pass64, nil
 }
 
-// StringDecryptByAES 字符串AES解密
-func StringDecryptByAES(text string) (string, error) {
+// AESDecryptBase64 AES解密解Base64字符串
+func AESDecryptBase64(text, key string) (string, error) {
 	if len(text) == 0 {
 		return "", nil
 	}
@@ -32,21 +32,16 @@ func StringDecryptByAES(text string) (string, error) {
 	if err != nil {
 		return "", err
 	}
-
-	tpass, err := aesDecryptWithSalt(bytesPass)
+	tpass, err := AESDecrypt(bytesPass, []byte(key))
 	if err != nil {
 		return "", err
 	}
 	return string(tpass), nil
 }
 
-// aesKey 字符串AES加解密密钥
-const aesKey = "AGT66VfY4SMaiT97a7df0aef1704d5c5"
-
-// const aesKey = "AGT66VfY4SMaiT97"
-// aesEncryptWithSalt AES加密
-func aesEncryptWithSalt(plaintext []byte) ([]byte, error) {
-	block, err := aes.NewCipher([]byte(aesKey))
+// AESEncrypt AES加密
+func AESEncrypt(plaintext, aeskey []byte) ([]byte, error) {
+	block, err := aes.NewCipher(aeskey)
 	if err != nil {
 		return nil, err
 	}
@@ -68,8 +63,8 @@ func aesEncryptWithSalt(plaintext []byte) ([]byte, error) {
 	return ciphertext, nil
 }
 
-// aesDecryptWithSalt AES解密
-func aesDecryptWithSalt(ciphertext []byte) ([]byte, error) {
+// AESDecrypt AES解密
+func AESDecrypt(ciphertext, aeskey []byte) ([]byte, error) {
 	blockSize := aes.BlockSize
 	if len(ciphertext) < blockSize {
 		return nil, fmt.Errorf("ciphertext too short")
@@ -77,12 +72,14 @@ func aesDecryptWithSalt(ciphertext []byte) ([]byte, error) {
 
 	iv := ciphertext[:blockSize]
 	ciphertext = ciphertext[blockSize:]
+	block, err := aes.NewCipher([]byte(aeskey))
 
-	block, err := aes.NewCipher([]byte(aesKey))
 	if err != nil {
 		return nil, err
 	}
-
+	if len(ciphertext) == 0 {
+		return nil, fmt.Errorf("ciphertext is invalid")
+	}
 	if len(ciphertext)%blockSize != 0 {
 		return nil, fmt.Errorf("ciphertext is not a multiple of the block size")
 	}