From 453fa309b1ae0f62308b08d9751c93128c4f2abb Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Sat, 16 Jul 2022 22:59:25 +0800
Subject: [PATCH] feat: export gateway.Server to let users add middlewares
 (#2157)

---
 gateway/config.go        |  1 +
 gateway/requestparser.go | 23 +++++++++++------------
 gateway/server.go        | 10 +++++-----
 3 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/gateway/config.go b/gateway/config.go
index aac835aa..684daa59 100644
--- a/gateway/config.go
+++ b/gateway/config.go
@@ -24,6 +24,7 @@ type (
 		// Rpc is the gRPC rpc method, with format of package.service/method
 		Rpc string
 	}
+
 	// upstream is the configuration for upstream.
 	upstream struct {
 		// Grpc is the target of upstream.
diff --git a/gateway/requestparser.go b/gateway/requestparser.go
index ed806319..fe495b32 100644
--- a/gateway/requestparser.go
+++ b/gateway/requestparser.go
@@ -10,6 +10,15 @@ import (
 	"github.com/zeromicro/go-zero/rest/pathvar"
 )
 
+func buildJsonRequestParser(v interface{}, resolver jsonpb.AnyResolver) (grpcurl.RequestParser, error) {
+	var buf bytes.Buffer
+	if err := json.NewEncoder(&buf).Encode(v); err != nil {
+		return nil, err
+	}
+
+	return grpcurl.NewJSONRequestParser(&buf, resolver), nil
+}
+
 func newRequestParser(r *http.Request, resolver jsonpb.AnyResolver) (grpcurl.RequestParser, error) {
 	vars := pathvar.Vars(r)
 	if len(vars) == 0 {
@@ -17,12 +26,7 @@ func newRequestParser(r *http.Request, resolver jsonpb.AnyResolver) (grpcurl.Req
 	}
 
 	if r.ContentLength == 0 {
-		var buf bytes.Buffer
-		if err := json.NewEncoder(&buf).Encode(vars); err != nil {
-			return nil, err
-		}
-
-		return grpcurl.NewJSONRequestParser(&buf, resolver), nil
+		return buildJsonRequestParser(vars, resolver)
 	}
 
 	m := make(map[string]interface{})
@@ -34,10 +38,5 @@ func newRequestParser(r *http.Request, resolver jsonpb.AnyResolver) (grpcurl.Req
 		m[k] = v
 	}
 
-	var buf bytes.Buffer
-	if err := json.NewEncoder(&buf).Encode(m); err != nil {
-		return nil, err
-	}
-
-	return grpcurl.NewJSONRequestParser(&buf, resolver), nil
+	return buildJsonRequestParser(m, resolver)
 }
diff --git a/gateway/server.go b/gateway/server.go
index 02c26d00..c95cdd17 100644
--- a/gateway/server.go
+++ b/gateway/server.go
@@ -19,7 +19,7 @@ import (
 
 // Server is a gateway server.
 type Server struct {
-	svr       *rest.Server
+	*rest.Server
 	upstreams []upstream
 	timeout   time.Duration
 }
@@ -27,7 +27,7 @@ type Server struct {
 // MustNewServer creates a new gateway server.
 func MustNewServer(c GatewayConf) *Server {
 	return &Server{
-		svr:       rest.MustNewServer(c.RestConf),
+		Server:    rest.MustNewServer(c.RestConf),
 		upstreams: c.Upstreams,
 		timeout:   c.Timeout,
 	}
@@ -36,12 +36,12 @@ func MustNewServer(c GatewayConf) *Server {
 // Start starts the gateway server.
 func (s *Server) Start() {
 	logx.Must(s.build())
-	s.svr.Start()
+	s.Server.Start()
 }
 
 // Stop stops the gateway server.
 func (s *Server) Stop() {
-	s.svr.Stop()
+	s.Server.Stop()
 }
 
 func (s *Server) build() error {
@@ -69,7 +69,7 @@ func (s *Server) build() error {
 	}, func(pipe <-chan interface{}, cancel func(error)) {
 		for item := range pipe {
 			route := item.(rest.Route)
-			s.svr.AddRoute(route)
+			s.Server.AddRoute(route)
 		}
 	})
 }