fusen-gateway/server/logic/gateway_response.go
2023-11-17 18:35:28 +08:00

101 lines
2.3 KiB
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
}