fusen-gateway/server/logic/gateway_response.go

package logic

import (
	"bytes"
	"context"
	"fmt"
	"fusen-basic/basic"
	"fusen-gateway/gen/go/service"
	"net/http"
	"regexp"

	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
)

type EmptyMarshaler struct {
	runtime.JSONPb
}

func (m *EmptyMarshaler) Marshal(v interface{}) ([]byte, error) {
	return nil, nil
}

var fsDefaultProtoJSON = &runtime.JSONPb{
	MarshalOptions: protojson.MarshalOptions{
		EmitUnpopulated: true,
		UseProtoNames:   true,
	},
	UnmarshalOptions: protojson.UnmarshalOptions{
		DiscardUnknown: true,
	},
}

var codeRE = regexp.MustCompile(`(?i)"code"\s*:\s*(\d+)`)

func WriteFusenResponse(fsHeader string, w http.ResponseWriter, resp proto.Message) error {
	w.WriteHeader(200)
	var buf = bytes.NewBufferString(fsHeader[:len(fsHeader)-1] + ",")
	var rdata []byte
	var err error

	if customResp, ok := resp.(*service.Response); ok {
		if customResp.GetData() != nil {
			rdata, err = fsDefaultProtoJSON.Marshal(customResp.GetData())
			if err != nil {
				return err
			}
		} else {
			rdata = []byte("null")
		}
	} else {
		rdata, err = fsDefaultProtoJSON.Marshal(resp)
		if err != nil {
			return err
		}
	}

	buf.WriteString(`"data":`)
	buf.Write(rdata)
	buf.WriteString(`}`)

	w.Write(buf.Bytes())
	return nil
}

func ResponseHeaderMatcher(ctx context.Context, w http.ResponseWriter, resp proto.Message) error {
	headers := w.Header()

	key := http.CanonicalHeaderKey(basic.GM_FusenResponse.GrpcMetadataKey())

	if grpcResp, ok := headers[key]; ok {
		fsHeader := grpcResp[0]
		code := codeRE.FindStringSubmatch(string(fsHeader))
		if len(code) > 0 && code[1] != "200" {
			return WriteFusenResponse(fsHeader, w, resp)
		}

		// 判断自定返回的内容, 用html的页面返回等...
		selfwrite := http.CanonicalHeaderKey(basic.GM_SelfWrite.GrpcMetadataKey())
		if sw, ok := headers[selfwrite]; ok {
			w.WriteHeader(200)
			w.Write([]byte(sw[0]))
			return nil
		}

		// 重定向
		redirect := http.CanonicalHeaderKey(basic.GM_Redirect.GrpcMetadataKey())
		if location, ok := headers[redirect]; ok {
			w.WriteHeader(http.StatusFound)
			w.Header().Set("Location", location[0])
			return nil
		}

		return WriteFusenResponse(fsHeader, w, resp)
	}

	w.Write([]byte(fmt.Sprintf("%s error", basic.GM_FusenResponse.String())))
	return nil
}