Эх сурвалжийг харах

defines the method to customize http server (#2171)

sniperwzq 2 жил өмнө
parent
commit
3e093bf34e
2 өөрчлөгдсөн 21 нэмэгдсэн , 9 устгасан
  1. 11 7
      rest/engine.go
  2. 10 2
      rest/server.go

+ 11 - 7
rest/engine.go

@@ -301,21 +301,25 @@ func (ng *engine) signatureVerifier(signature signatureSetting) (func(chain.Chai
 	}, nil
 }
 
-func (ng *engine) start(router httpx.Router) error {
+func (ng *engine) start(router httpx.Router, opts ...internal.StartOption) error {
 	if err := ng.bindRoutes(router); err != nil {
 		return err
 	}
 
+	opts = append(opts, ng.withTimeout())
+
 	if len(ng.conf.CertFile) == 0 && len(ng.conf.KeyFile) == 0 {
-		return internal.StartHttp(ng.conf.Host, ng.conf.Port, router, ng.withTimeout())
+		return internal.StartHttp(ng.conf.Host, ng.conf.Port, router, opts...)
 	}
 
+	opts = append(opts, func(svr *http.Server) {
+		if ng.tlsConfig != nil {
+			svr.TLSConfig = ng.tlsConfig
+		}
+	})
+
 	return internal.StartHttps(ng.conf.Host, ng.conf.Port, ng.conf.CertFile,
-		ng.conf.KeyFile, router, func(svr *http.Server) {
-			if ng.tlsConfig != nil {
-				svr.TLSConfig = ng.tlsConfig
-			}
-		}, ng.withTimeout())
+		ng.conf.KeyFile, router, opts...)
 }
 
 func (ng *engine) use(middleware Middleware) {

+ 10 - 2
rest/server.go

@@ -11,6 +11,7 @@ import (
 	"github.com/zeromicro/go-zero/rest/chain"
 	"github.com/zeromicro/go-zero/rest/handler"
 	"github.com/zeromicro/go-zero/rest/httpx"
+	"github.com/zeromicro/go-zero/rest/internal"
 	"github.com/zeromicro/go-zero/rest/internal/cors"
 	"github.com/zeromicro/go-zero/rest/router"
 )
@@ -19,6 +20,9 @@ type (
 	// RunOption defines the method to customize a Server.
 	RunOption func(*Server)
 
+	// StartOption defines the method to customize http server.
+	StartOption func(svr *http.Server)
+
 	// A Server is a http server.
 	Server struct {
 		ngin   *engine
@@ -112,8 +116,12 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 // Start starts the Server.
 // Graceful shutdown is enabled by default.
 // Use proc.SetTimeToForceQuit to customize the graceful shutdown period.
-func (s *Server) Start() {
-	handleError(s.ngin.start(s.router))
+func (s *Server) Start(opts ...StartOption) {
+	var startOption []internal.StartOption
+	for _, opt := range opts {
+		startOption = append(startOption, internal.StartOption(opt))
+	}
+	handleError(s.ngin.start(s.router, startOption...))
 }
 
 // Stop stops the Server.