spark-server/internal/transport/http/middleware.go

package httpx

import (
	"bytes"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"go.uber.org/zap"

	"spark-server/internal/infrastructure/logger"
	"spark-server/internal/pkg/errs"
	"spark-server/internal/pkg/i18n"
	"spark-server/internal/pkg/reply"
	"spark-server/internal/pkg/reqctx"
)

const (
	headerRequestID = "X-Request-ID"
	headerTraceID   = "X-Trace-ID"
)

func Trace(base *zap.Logger) gin.HandlerFunc {
	return func(c *gin.Context) {
		reqID := c.GetHeader(headerRequestID)
		if reqID == "" {
			reqID = reqctx.NewID()
		}
		traceID := c.GetHeader(headerTraceID)
		if traceID == "" {
			traceID = reqID
		}
		rc := &reqctx.RequestContext{
			TraceID:    traceID,
			RequestID:  reqID,
			RemoteAddr: c.Request.RemoteAddr,
			RealIP: reqctx.ParseHTTPRealIP(
				c.GetHeader(headerRealIP),
				c.GetHeader(headerForwardedFor),
				c.Request.RemoteAddr,
			),
			AppVersion: c.GetHeader(headerAppVersion),
			Language:   requestLanguage(c),
		}
		ctx := reqctx.With(c.Request.Context(), rc)
		c.Request = c.Request.WithContext(logger.Attach(ctx, base))
		c.Header(headerRequestID, reqID)
		c.Header(headerTraceID, traceID)
		c.Next()
	}
}

func Recovery(_ *zap.Logger) gin.HandlerFunc {
	return gin.CustomRecovery(func(c *gin.Context, recovered any) {
		logger.From(c.Request.Context()).Error("panic", zap.Any("panic", recovered), zap.String("path", c.Request.URL.Path))
		c.Abort()
		c.JSON(http.StatusOK, reply.FromContext(c.Request.Context(), errs.ErrInner.WithInnerMsg("panic")))
	})
}

func AccessLog() gin.HandlerFunc {
	return func(c *gin.Context) {
		if skipAccessLog(c.Request.URL.Path) {
			c.Next()
			return
		}
		reqBody := readRequestBody(c)
		bw := &bodyWriter{ResponseWriter: c.Writer}
		c.Writer = bw
		start := time.Now()
		c.Next()
		logger.From(c.Request.Context()).Info("http_access",
			zap.String("method", c.Request.Method),
			zap.String("path", c.Request.URL.Path),
			zap.Int("status", c.Writer.Status()),
			zap.Int64("latency_ms", time.Since(start).Milliseconds()),
			zap.String("request_body", logger.ClipBody(reqBody)),
			zap.String("response_body", logger.ClipBody(bw.buf.Bytes())),
		)
	}
}

func skipAccessLog(path string) bool {
	switch path {
	case "/healthz", "/readyz", "/openapi.yaml", "/docs", "/docs/":
		return true
	default:
		return strings.HasPrefix(path, "/docs/")
	}
}

func readRequestBody(c *gin.Context) []byte {
	if c.Request.Body == nil {
		return nil
	}
	if ct := c.ContentType(); strings.HasPrefix(ct, "multipart/") {
		return []byte("<multipart omitted>")
	}
	raw, err := io.ReadAll(io.LimitReader(c.Request.Body, int64(logger.MaxBody)+1))
	if err != nil {
		return nil
	}
	_ = c.Request.Body.Close()
	c.Request.Body = io.NopCloser(bytes.NewReader(raw))
	return raw
}

type bodyWriter struct {
	gin.ResponseWriter
	buf bytes.Buffer
}

func (w *bodyWriter) Write(b []byte) (int, error) {
	w.capture(b)
	return w.ResponseWriter.Write(b)
}

func (w *bodyWriter) WriteString(s string) (int, error) {
	if w.buf.Len() < logger.MaxBody {
		remain := logger.MaxBody - w.buf.Len()
		if len(s) < remain {
			_, _ = w.buf.WriteString(s)
		} else {
			_, _ = w.buf.WriteString(s[:remain])
		}
	}
	return w.ResponseWriter.WriteString(s)
}

func (w *bodyWriter) capture(b []byte) {
	if w.buf.Len() >= logger.MaxBody {
		return
	}
	remain := logger.MaxBody - w.buf.Len()
	if len(b) < remain {
		_, _ = w.buf.Write(b)
		return
	}
	_, _ = w.buf.Write(b[:remain])
}

func requestLanguage(c *gin.Context) string {
	raw := c.GetHeader(headerXLanguage)
	if raw == "" {
		raw = c.GetHeader(headerLanguage)
	}
	return i18n.Parse(raw).String()
}