user-system/internal/api/middleware/response_wrapper.go

package middleware

import (
	"bytes"
	"encoding/json"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
)

// responseWrapper 捕获 handler 输出的中间件
// 将所有裸 JSON 响应自动包装为 {code: 0, message: "success", data: ...} 格式
type responseWrapper struct {
	gin.ResponseWriter
	body       *bytes.Buffer
	statusCode int
}

func (w *responseWrapper) Write(b []byte) (int, error) {
	w.body.Write(b)
	// 不再同时写到原始 writer，让 body 完全缓冲
	return len(b), nil
}

func (w *responseWrapper) WriteString(s string) (int, error) {
	w.body.WriteString(s)
	return len(s), nil
}

func (w *responseWrapper) WriteHeader(code int) {
	w.statusCode = code
	// 不实际写入，让 gin 的最终写入处理
}

// ResponseWrapper 返回包装响应格式的中间件
func ResponseWrapper() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 跳过非 JSON 响应（如文件下载、流式响应）
		contentType := c.GetHeader("Content-Type")
		if strings.Contains(contentType, "text/event-stream") ||
			contentType == "application/octet-stream" ||
			strings.HasPrefix(c.Request.URL.Path, "/swagger/") {
			c.Next()
			return
		}

		// 包装 response writer 以捕获输出
		wrapper := &responseWrapper{
			ResponseWriter: c.Writer,
			body:          bytes.NewBuffer(nil),
			statusCode:    http.StatusOK,
		}
		c.Writer = wrapper

		c.Next()

		// 检查是否已标记为已包装
		if _, exists := c.Get("response_wrapped"); exists {
			// 直接把捕获的内容写回到底层 writer
			wrapper.ResponseWriter.WriteHeader(wrapper.statusCode)
			wrapper.ResponseWriter.Write(wrapper.body.Bytes())
			return
		}

		// 只处理成功响应（2xx）
		if wrapper.statusCode < 200 || wrapper.statusCode >= 300 {
			// 非成功状态，直接把捕获的内容写回
			wrapper.ResponseWriter.WriteHeader(wrapper.statusCode)
			wrapper.ResponseWriter.Write(wrapper.body.Bytes())
			return
		}

		// 解析捕获的 body
		if wrapper.body.Len() == 0 {
			wrapper.ResponseWriter.WriteHeader(wrapper.statusCode)
			return
		}

		bodyBytes := wrapper.body.Bytes()

		// 尝试解析为 JSON 对象
		var raw json.RawMessage
		if err := json.Unmarshal(bodyBytes, &raw); err != nil {
			// 不是有效 JSON，不包装
			wrapper.ResponseWriter.WriteHeader(wrapper.statusCode)
			wrapper.ResponseWriter.Write(bodyBytes)
			return
		}

		// 检查是否已经是标准格式（有 code 字段）
		var checkMap map[string]interface{}
		if err := json.Unmarshal(bodyBytes, &checkMap); err == nil {
			if _, hasCode := checkMap["code"]; hasCode {
				// 已经是标准格式，不重复包装
				wrapper.ResponseWriter.WriteHeader(wrapper.statusCode)
				wrapper.ResponseWriter.Write(bodyBytes)
				return
			}
		}

		// 包装为标准格式
		wrapped := map[string]interface{}{
			"code":    0,
			"message": "success",
			"data":    raw,
		}

		wrappedBytes, err := json.Marshal(wrapped)
		if err != nil {
			wrapper.ResponseWriter.WriteHeader(wrapper.statusCode)
			wrapper.ResponseWriter.Write(bodyBytes)
			return
		}

		// 设置响应头并写入包装后的内容
		wrapper.ResponseWriter.Header().Set("Content-Type", "application/json")
		wrapper.ResponseWriter.WriteHeader(wrapper.statusCode)
		wrapper.ResponseWriter.Write(wrappedBytes)
	}
}

// WrapResponse 标记响应为已包装，防止重复包装
// handler 中使用 response.Success() 等方法后调用此函数
func WrapResponse(c *gin.Context) {
	c.Set("response_wrapped", true)
}

// NoWrapper 跳过包装的中间件处理器
func NoWrapper() gin.HandlerFunc {
	return func(c *gin.Context) {
		WrapResponse(c)
		c.Next()
	}
}