lone-services/bff/internal/response/middleware.go

package response

import (
	"bytes"
	"encoding/base64"
	"lone-services/pkg/utils"
	"net/http"
	"strings"

	jsoniter "github.com/json-iterator/go"
	"github.com/zeromicro/go-zero/core/logx"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type bodyWriter struct {
	http.ResponseWriter
	status int
	buf    bytes.Buffer
}

type responseDataBack struct {
	ReturnCode string `json:"return_code"`
	ReturnMsg  string `json:"return_msg"`
}

type responseDataDown struct {
	Data string `json:"data"` // excel二进制
	Name string `json:"name"`
	Msg  string `json:"msg"`
}

type responseData struct {
	Code int32  `json:"code"`
	Data any    `json:"data"`
	Msg  string `json:"msg"`
}

func (w *bodyWriter) WriteHeader(statusCode int) {
	w.status = statusCode
}

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

func Wrap(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		bw := &bodyWriter{ResponseWriter: w, status: http.StatusOK}
		next(bw, r)
		ctrl := BaseController{}
		if bw.status != http.StatusOK {
			ctrl.Error(w, mapHTTPError(bw.status, bw.buf.String()))
			return
		}

		raw := bytes.TrimSpace(bw.buf.Bytes())

		res := make(map[string]interface{})
		if len(raw) > utils.NumberZero {
			if err := jsoniter.Unmarshal(raw, &res); err != nil {
				ctrl.Fail(w)
				return
			}
		}
		if _, ok := res["code"]; ok {
			var data responseData
			if err := jsoniter.Unmarshal(raw, &data); err != nil {
				ctrl.Fail(w)
				return
			}
			ctrl.OutPut(w, data.Code, unwrapJSONData(data.Data), data.Msg)
		} else if _, ok := res["return_code"]; ok {
			var data responseDataBack
			if err := jsoniter.Unmarshal(raw, &data); err != nil {
				ctrl.Fail(w)
				return
			}
			ctrl.OutPutBack(w, data)
		} else if _, ok := res["name"]; ok {

			var data responseDataDown
			if err := jsoniter.Unmarshal(raw, &data); err != nil {
				ctrl.Fail(w)
				return
			}
			excelBin, err := base64.StdEncoding.DecodeString(data.Data)

			if err != nil {
				logx.Error("excel base64 decode failed", err, data.Data)
				ctrl.Fail(w)
				return
			}
			data.Data = string(excelBin)
			ctrl.OutPutDown(w, data)
		}
	}
}

func unwrapJSONData(v any) any {
	s, ok := v.(string)
	if !ok {
		return v
	}
	s = strings.TrimSpace(s)
	if s == "" {
		return nil
	}
	var parsed any
	if err := jsoniter.Unmarshal([]byte(s), &parsed); err != nil {
		return v
	}
	return parsed
}

func mapHTTPError(httpStatus int, body string) Error {
	msg := strings.TrimSpace(body)
	msg = strings.TrimPrefix(msg, "rpc error: ")
	if i := strings.Index(msg, "desc = "); i >= 0 {
		msg = strings.TrimSpace(msg[i+len("desc = "):])
	}
	if msg == "" {
		msg = Fail.GetMsg()
	}

	switch httpStatus {
	case http.StatusBadRequest:
		return Define(ErrorMissingParams, msg)
	case http.StatusUnauthorized:
		return Define(ErrorNoLogin, msg)
	case http.StatusForbidden:
		return Define(ErrorAuthority, msg)
	case http.StatusNotFound:
		return Define(ErrorDataNotExist, msg)
	default:
		return Define(Fail, msg)
	}
}

func FromGRPC(err error) Error {
	if err == nil {
		return OK
	}
	st, ok := status.FromError(err)
	if !ok {
		return Define(Fail, err.Error())
	}
	msg := st.Message()
	if msg == "" {
		msg = st.Code().String()
	}
	switch st.Code() {
	case codes.InvalidArgument:
		return Define(ErrorMissingParams, msg)
	case codes.Unauthenticated:
		return Define(ErrorNoLogin, msg)
	case codes.PermissionDenied:
		return Define(ErrorAuthority, msg)
	case codes.NotFound:
		return Define(ErrorDataNotExist, msg)
	case codes.AlreadyExists:
		return Define(ErrorDataIsExist, msg)
	default:
		return Define(Fail, msg)
	}
}