From 5bf701a22b967c134e87828033de42573f9baea3 Mon Sep 17 00:00:00 2001
From: Yun <995116474@qq.com>
Date: Wed, 26 Jun 2024 17:27:44 +0800
Subject: [PATCH] =?UTF-8?q?=E5=B0=81=E8=A3=85=E4=B8=AD=E9=97=B4=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 middleware.go           | 55 -----------------------
 middleware/ginParams.go | 99 +++++++++++++++++++++++++++++++++++++++++
 middleware/traceId.go   | 37 +++++++++++++++
 3 files changed, 136 insertions(+), 55 deletions(-)
 delete mode 100644 middleware.go
 create mode 100644 middleware/ginParams.go
 create mode 100644 middleware/traceId.go

diff --git a/middleware.go b/middleware.go
deleted file mode 100644
index b8f8582..0000000
--- a/middleware.go
+++ /dev/null
@@ -1,55 +0,0 @@
-package loggerx
-
-import (
-	"bytes"
-	"context"
-	"io"
-	"net/http"
-
-	"github.com/gin-gonic/gin"
-	uuid "github.com/satori/go.uuid"
-)
-
-// 可以被引入的中间件
-
-// 自动设置请求ID
-func SetTraceId(ctx context.Context) context.Context {
-	return context.WithValue(ctx, "trace_id", uuid.NewV4().String())
-}
-
-// 中间件 gin框架保存请求相关信息
-func MiddlewareGetGinParams(log *Logger) gin.HandlerFunc {
-	return func(ctx *gin.Context) {
-		// 使用bytes.Buffer来读取并记录请求体，同时避免改变ctx.Request.Body
-		buf := &bytes.Buffer{}
-		tea := io.TeeReader(ctx.Request.Body, buf)
-
-		// 读取body
-		body, err := io.ReadAll(tea)
-		if err != nil {
-			log.Infof(ctx, "Error reading request body: %v", err)
-			ctx.AbortWithStatus(http.StatusInternalServerError)
-			return
-		}
-
-		// 截取body的前1000个字符
-		bodyStr := string(body)
-		if len(bodyStr) > 1000 {
-			bodyStr = bodyStr[:1000]
-		}
-
-		// 使用NopCloser包裹buffer，仅为了确保在读取body之后body可以被关闭，但并不改变原始的Request.Body
-		ctx.Request.Body = io.NopCloser(buf)
-
-		m := map[string]interface{}{
-			"method": ctx.Request.Method,
-			"uri":    ctx.Request.RequestURI,
-			"body":   bodyStr,
-			"query":  ctx.Request.URL.Query(),
-			"header": ctx.Request.Header,
-		}
-		log.Infof(ctx, "request %+v", m)
-		ctx.Next()
-
-	}
-}
diff --git a/middleware/ginParams.go b/middleware/ginParams.go
new file mode 100644
index 0000000..2df668d
--- /dev/null
+++ b/middleware/ginParams.go
@@ -0,0 +1,99 @@
+package middleware
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+	"github.com/yuninks/loggerx"
+)
+
+// 可以被引入的中间件
+
+// 自动设置请求ID
+
+// 中间件 gin框架保存请求响应相关信息
+func SetGinParams(log *loggerx.Logger) gin.HandlerFunc {
+	return func(ctx *gin.Context) {
+
+		// 使用bytes.Buffer来读取并记录请求体，同时避免改变ctx.Request.Body
+		buf := &bytes.Buffer{}
+		tea := io.TeeReader(ctx.Request.Body, buf)
+
+		// 读取body
+		body, err := io.ReadAll(tea)
+		if err != nil {
+			log.Infof(ctx, "Error reading request body: %v", err)
+			ctx.AbortWithStatus(http.StatusInternalServerError)
+			return
+		}
+
+		// 截取body的前1000个字符
+		bodyStr := string(body)
+		if len(bodyStr) > 1000 {
+			bodyStr = bodyStr[:1000]
+		}
+
+		// 使用NopCloser包裹buffer，仅为了确保在读取body之后body可以被关闭，但并不改变原始的Request.Body
+		ctx.Request.Body = io.NopCloser(buf)
+
+		// 下面是响应参数缓存
+		blw := &bodyParamsWriter{
+			body:           bytes.NewBufferString(""),
+			ResponseWriter: ctx.Writer,
+		}
+		ctx.Writer = blw
+
+		req := struct {
+			Method   string              `json:"method"`
+			Path     string              `json:"path"`
+			ClientIP string              `json:"client_ip"`
+			Body     string              `json:"body"`
+			Header   map[string][]string `json:"header"`
+		}{
+			Method:   ctx.Request.Method,
+			Path:     ctx.FullPath(),
+			ClientIP: ctx.ClientIP(),
+			Body:     bodyStr,
+			Header:   ctx.Request.Header,
+		}
+
+		log.Info(ctx, "request", req)
+
+		ctx.Next()
+
+		// 判断返回的值是否json
+		contentType := ctx.Writer.Header().Get("Content-Type")
+
+		if !strings.Contains(contentType, "application/json") {
+			return
+		}
+
+		respData := blw.body.String()
+
+		resp := struct {
+			HttpCode int    `json:"http_code"`
+			Response string `json:"response"`
+		}{
+			HttpCode: ctx.Writer.Status(),
+			Response: respData,
+			// Header:   ctx.Writer.Header(),
+		}
+		log.Info(ctx, "response", resp)
+
+	}
+}
+
+type bodyParamsWriter struct {
+	gin.ResponseWriter
+	body *bytes.Buffer
+}
+
+// 真正输出的时候，它会调用这个输出方法
+func (w bodyParamsWriter) Write(b []byte) (int, error) {
+	// b = []byte(`{"code":0}`)
+	w.body.Write(b)
+	return w.ResponseWriter.Write(b)
+}
diff --git a/middleware/traceId.go b/middleware/traceId.go
new file mode 100644
index 0000000..deec9e3
--- /dev/null
+++ b/middleware/traceId.go
@@ -0,0 +1,37 @@
+package middleware
+
+import (
+	"context"
+
+	"github.com/gin-gonic/gin"
+	uuid "github.com/satori/go.uuid"
+)
+
+// 设置普通的traceId
+func SetTraceId(ctx context.Context, traceKey string) context.Context {
+	if traceKey == "" {
+		traceKey = "trace_id"
+	}
+
+	val := ctx.Value(traceKey)
+	if val == nil {
+		ctx = context.WithValue(ctx, traceKey, uuid.NewV4().String())
+	}
+	return ctx
+}
+
+// 设置Gin的traceId
+func SetGinTraceId(traceKey string) gin.HandlerFunc {
+
+	if traceKey == "" {
+		traceKey = "trace_id"
+	}
+
+	return func(ctx *gin.Context) {
+		traceId := ctx.Request.Header.Get(traceKey)
+		if traceId == "" {
+			traceId = uuid.NewV4().String()
+		}
+		ctx.Set(traceKey, traceId)
+	}
+}