long-agent
582ad7a069
- Add new test files for auth, service, and handler modules - Improve test organization and coverage - Refactor code for better maintainability - Add captcha, settings, stats, and theme handler tests - Add auth module tests (CAS, OAuth, password, SSO, state) - Add service layer tests for auth, export, permissions, roles - All Go tests pass (exit code 0) - All frontend tests pass (325 tests in 59 files)
136 lines
3.6 KiB
Go
136 lines
3.6 KiB
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()
|
||
}
|
||
}
|