pkg/middleware/check/check.go

package check

import (
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"net/http"
)

type (
	errorWriter struct {
		gin.ResponseWriter
		data []byte
	}
)

func (ew *errorWriter) WriteString(data string) (int, error) {
	ew.data = append(ew.data, []byte(data)...)
	return ew.ResponseWriter.WriteString(data)
}

func (ew *errorWriter) Write(data []byte) (int, error) {
	ew.data = append(ew.data, data...)

	return ew.ResponseWriter.Write(data)
}

func (ew *errorWriter) successful() bool {
	switch ew.Status() {
	case 200, 201, 202, 204:
		return true
	case 400, 401, 403, 404, 406, 410, 422, 500:
		return false
	default:
		panic(fmt.Sprintf("unsupported status code %d", ew.Status()))
	}
}

func CodeAnalysis() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		ew := &errorWriter{ResponseWriter: ctx.Writer}
		ctx.Writer = ew

		ctx.Next()

		switch ew.Header().Get("Content-Type") {
		case "application/json; charset=utf-8":
			rps := map[string]string{}

			err := json.Unmarshal(ew.data, &rps)
			if err != nil {
				panic("unsupported data format, only json now.")
			}

			if !ew.successful() {
				if len(rps) != 1 || rps["msg"] == "" {
					panic(fmt.Sprintf("you need return %s but, %s", `{"msg":"content can't empty'"}`, string(ew.data)))
				}
			} else {
				// 如果post请求的状态码不是201的话返回错误
				if ctx.Request.Method == "POST" && ew.Status() != http.StatusCreated {
					panic(fmt.Sprintf("you need set http status to %d", http.StatusCreated))
				}

				// 如果状态码是201但是请求方式不是post的话返回错误
				if ew.Status() == http.StatusCreated && ctx.Request.Method != "POST" {
					panic(fmt.Sprintf("the http status %d can set when method=POST", http.StatusCreated))
				}

				// 如果返回的数据为nil则报错
				if rps == nil {
					panic(fmt.Sprintf("if response body is null, please only set http status"))
				}

				// 如果返回状态码为204则body不能有任何数据
				if ew.Status() == http.StatusNoContent && len(ew.data) != 0 {
					panic(fmt.Sprintf("if response stauts is %d, can't have reponse body", http.StatusNoContent))
				}
			}
		default:
			if len(ew.data) == 0 && ew.Status() != http.StatusNoContent {
				panic(fmt.Sprintf("if response body is null, please only set http status %d", http.StatusNoContent))
			}

			if len(ew.data) != 0 {
				panic("unsupported data format, only json now.")
			}
		}
	}
}