From 93584c6ca67f833145cdd0ee1630d5622e143326 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Sat, 27 Aug 2022 11:39:42 +0800
Subject: [PATCH] chore: refactor gateway (#2303)

---
 gateway/internal/eventhandler.go | 47 ++++++++++++++++++++++++++++++++
 gateway/server.go                |  6 +---
 2 files changed, 48 insertions(+), 5 deletions(-)
 create mode 100644 gateway/internal/eventhandler.go

diff --git a/gateway/internal/eventhandler.go b/gateway/internal/eventhandler.go
new file mode 100644
index 00000000..a840930f
--- /dev/null
+++ b/gateway/internal/eventhandler.go
@@ -0,0 +1,47 @@
+package internal
+
+import (
+	"io"
+
+	"github.com/golang/protobuf/jsonpb"
+	"github.com/golang/protobuf/proto"
+	"github.com/jhump/protoreflect/desc"
+	"github.com/zeromicro/go-zero/core/logx"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+)
+
+type EventHandler struct {
+	Status    *status.Status
+	writer    io.Writer
+	marshaler jsonpb.Marshaler
+}
+
+func NewEventHandler(writer io.Writer, resolver jsonpb.AnyResolver) *EventHandler {
+	return &EventHandler{
+		writer: writer,
+		marshaler: jsonpb.Marshaler{
+			EmitDefaults: true,
+			AnyResolver:  resolver,
+		},
+	}
+}
+
+func (h *EventHandler) OnResolveMethod(_ *desc.MethodDescriptor) {
+}
+
+func (h *EventHandler) OnSendHeaders(_ metadata.MD) {
+}
+
+func (h *EventHandler) OnReceiveHeaders(_ metadata.MD) {
+}
+
+func (h *EventHandler) OnReceiveResponse(message proto.Message) {
+	if err := h.marshaler.Marshal(h.writer, message); err != nil {
+		logx.Error(err)
+	}
+}
+
+func (h *EventHandler) OnReceiveTrailers(status *status.Status, _ metadata.MD) {
+	h.Status = status
+}
diff --git a/gateway/server.go b/gateway/server.go
index ede9e537..80309629 100644
--- a/gateway/server.go
+++ b/gateway/server.go
@@ -120,11 +120,6 @@ func (s *Server) build() error {
 func (s *Server) buildHandler(source grpcurl.DescriptorSource, resolver jsonpb.AnyResolver,
 	cli zrpc.Client, rpcPath string) func(http.ResponseWriter, *http.Request) {
 	return func(w http.ResponseWriter, r *http.Request) {
-		handler := &grpcurl.DefaultEventHandler{
-			Out: w,
-			Formatter: grpcurl.NewJSONFormatter(true,
-				grpcurl.AnyResolverFromDescriptorSource(source)),
-		}
 		parser, err := internal.NewRequestParser(r, resolver)
 		if err != nil {
 			httpx.Error(w, err)
@@ -136,6 +131,7 @@ func (s *Server) buildHandler(source grpcurl.DescriptorSource, resolver jsonpb.A
 		defer can()
 
 		w.Header().Set(httpx.ContentType, httpx.JsonContentType)
+		handler := internal.NewEventHandler(w, resolver)
 		if err := grpcurl.InvokeRPC(ctx, source, cli.Conn(), rpcPath, s.prepareMetadata(r.Header),
 			handler, parser.Next); err != nil {
 			httpx.Error(w, err)