ez-api/internal/middleware/response_envelope.go

package middleware

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

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

const businessCodeKey = "response_business_code"

type responseEnvelope struct {
	Code    string          `json:"code"`
	Data    json.RawMessage `json:"data"`
	Message string          `json:"message"`
}

type envelopeWriter struct {
	gin.ResponseWriter
	body   bytes.Buffer
	status int
	size   int
	wrote  bool
}

func (w *envelopeWriter) WriteHeader(code int) {
	w.status = code
	w.wrote = true
}

func (w *envelopeWriter) WriteHeaderNow() {
	if !w.wrote {
		w.WriteHeader(http.StatusOK)
	}
}

func (w *envelopeWriter) Write(data []byte) (int, error) {
	if !w.wrote {
		w.WriteHeader(http.StatusOK)
	}
	n, err := w.body.Write(data)
	w.size += n
	return n, err
}

func (w *envelopeWriter) WriteString(s string) (int, error) {
	if !w.wrote {
		w.WriteHeader(http.StatusOK)
	}
	n, err := w.body.WriteString(s)
	w.size += n
	return n, err
}

func (w *envelopeWriter) Status() int {
	if w.status == 0 {
		return http.StatusOK
	}
	return w.status
}

func (w *envelopeWriter) Size() int {
	return w.size
}

func (w *envelopeWriter) Written() bool {
	return w.wrote
}

// SetBusinessCode sets an explicit business code for the response envelope.
func SetBusinessCode(c *gin.Context, code string) {
	if c == nil {
		return
	}
	code = strings.TrimSpace(code)
	if code == "" {
		return
	}
	c.Set(businessCodeKey, code)
}

func ResponseEnvelope() gin.HandlerFunc {
	return func(c *gin.Context) {
		originalWriter := c.Writer
		writer := &envelopeWriter{ResponseWriter: originalWriter}
		c.Writer = writer

		c.Next()

		status := writer.Status()
		body := writer.body.Bytes()
		if !bodyAllowedForStatus(status) || len(body) == 0 {
			writeStatusOnly(originalWriter, status)
			return
		}

		contentType := originalWriter.Header().Get("Content-Type")
		if !isJSONContentType(contentType) {
			writeThrough(originalWriter, status, body)
			return
		}

		obj, objOK := parseObject(body)
		if objOK && isEnvelopeObject(obj) {
			writeThrough(originalWriter, status, body)
			return
		}

		code := businessCodeFromContext(c)
		if code == "" {
			code = defaultBusinessCode(status)
		}

		message := ""
		if status >= http.StatusBadRequest && objOK {
			if raw, ok := obj["error"]; ok {
				var msg string
				if err := json.Unmarshal(raw, &msg); err == nil {
					message = msg
				}
			}
		}

		envelope := responseEnvelope{
			Code:    code,
			Data:    json.RawMessage(body),
			Message: message,
		}
		payload, err := json.Marshal(envelope)
		if err != nil {
			writeThrough(originalWriter, status, body)
			return
		}

		originalWriter.Header().Set("Content-Type", "application/json; charset=utf-8")
		originalWriter.Header().Del("Content-Length")
		writeThrough(originalWriter, status, payload)
	}
}

func businessCodeFromContext(c *gin.Context) string {
	if c == nil {
		return ""
	}
	value, ok := c.Get(businessCodeKey)
	if !ok {
		return ""
	}
	code, ok := value.(string)
	if !ok {
		return ""
	}
	return strings.TrimSpace(code)
}

func defaultBusinessCode(status int) string {
	switch {
	case status >= http.StatusOK && status < http.StatusMultipleChoices:
		return "ok"
	case status == http.StatusBadRequest:
		return "invalid_request"
	case status == http.StatusUnauthorized:
		return "unauthorized"
	case status == http.StatusForbidden:
		return "forbidden"
	case status == http.StatusNotFound:
		return "not_found"
	case status == http.StatusConflict:
		return "conflict"
	case status == http.StatusTooManyRequests:
		return "rate_limited"
	case status >= http.StatusBadRequest && status < http.StatusInternalServerError:
		return "request_error"
	case status >= http.StatusInternalServerError:
		return "internal_error"
	default:
		return "ok"
	}
}

func bodyAllowedForStatus(status int) bool {
	switch {
	case status >= 100 && status <= 199:
		return false
	case status == http.StatusNoContent:
		return false
	case status == http.StatusNotModified:
		return false
	}
	return true
}

func isJSONContentType(contentType string) bool {
	if contentType == "" {
		return false
	}
	return strings.Contains(strings.ToLower(contentType), "application/json")
}

func parseObject(body []byte) (map[string]json.RawMessage, bool) {
	var obj map[string]json.RawMessage
	if err := json.Unmarshal(body, &obj); err != nil {
		return nil, false
	}
	return obj, true
}

func isEnvelopeObject(obj map[string]json.RawMessage) bool {
	if obj == nil {
		return false
	}
	if _, ok := obj["code"]; !ok {
		return false
	}
	if _, ok := obj["data"]; !ok {
		return false
	}
	if _, ok := obj["message"]; !ok {
		return false
	}
	return true
}

func writeStatusOnly(w gin.ResponseWriter, status int) {
	w.WriteHeader(status)
}

func writeThrough(w gin.ResponseWriter, status int, body []byte) {
	w.WriteHeader(status)
	_, _ = w.Write(body)
}