diff --git a/net/ghttp/ghttp_server_config.go b/net/ghttp/ghttp_server_config.go index 867543cb4..4e22165f5 100644 --- a/net/ghttp/ghttp_server_config.go +++ b/net/ghttp/ghttp_server_config.go @@ -9,8 +9,13 @@ package ghttp import ( "context" "crypto/tls" + "fmt" + "github.com/gogf/gf/v2/errors/gcode" + "github.com/gogf/gf/v2/errors/gerror" + "net" "net/http" "strconv" + "strings" "time" "github.com/gogf/gf/v2/internal/intlog" @@ -50,6 +55,9 @@ type ServerConfig struct { // HTTPSAddr specifies the HTTPS addresses, multiple addresses joined using char ','. HTTPSAddr string `json:"httpsAddr"` + // Listeners specifies the custom listeners. + Listeners []net.Listener `json:"listeners"` + // HTTPSCertPath specifies certification file path for HTTPS service. HTTPSCertPath string `json:"httpsCertPath"` @@ -254,6 +262,7 @@ func NewConfig() ServerConfig { Name: DefaultServerName, Address: ":0", HTTPSAddr: "", + Listeners: nil, Handler: nil, ReadTimeout: 60 * time.Second, WriteTimeout: 0, // No timeout. @@ -408,6 +417,25 @@ func (s *Server) SetHTTPSPort(port ...int) { } } +// SetListener set the custom listener for the server. +func (s *Server) SetListener(listeners ...net.Listener) error { + if listeners == nil { + return gerror.NewCodef(gcode.CodeInvalidParameter, "SetListener failed: listener can not be nil") + } + if len(listeners) > 0 { + ports := make([]string, len(listeners)) + for k, v := range listeners { + if v == nil { + return gerror.NewCodef(gcode.CodeInvalidParameter, "SetListener failed: listener can not be nil") + } + ports[k] = fmt.Sprintf(":%d", (v.Addr().(*net.TCPAddr)).Port) + } + s.config.Address = strings.Join(ports, ",") + s.config.Listeners = listeners + } + return nil +} + // EnableHTTPS enables HTTPS with given certification and key files for the server. // The optional parameter `tlsConfig` specifies custom TLS configuration. func (s *Server) EnableHTTPS(certFile, keyFile string, tlsConfig ...*tls.Config) { diff --git a/net/ghttp/ghttp_server_graceful.go b/net/ghttp/ghttp_server_graceful.go index a1a1dcf92..4f933e7b1 100644 --- a/net/ghttp/ghttp_server_graceful.go +++ b/net/ghttp/ghttp_server_graceful.go @@ -13,6 +13,7 @@ import ( "net" "net/http" "os" + "strconv" "sync" "github.com/gogf/gf/v2/errors/gerror" @@ -49,6 +50,18 @@ func (s *Server) newGracefulServer(address string, fd ...int) *gracefulServer { if len(fd) > 0 && fd[0] > 0 { gs.fd = uintptr(fd[0]) } + if s.config.Listeners != nil { + addrArray := gstr.SplitAndTrim(address, ":") + addrPort, err := strconv.Atoi(addrArray[len(addrArray)-1]) + if err == nil { + for _, v := range s.config.Listeners { + if listenerPort := (v.Addr().(*net.TCPAddr)).Port; listenerPort == addrPort { + gs.rawListener = v + break + } + } + } + } return gs } @@ -175,6 +188,9 @@ func (s *gracefulServer) doServe(ctx context.Context) error { // getNetListener retrieves and returns the wrapped net.Listener. func (s *gracefulServer) getNetListener() (net.Listener, error) { + if s.rawListener != nil { + return s.rawListener, nil + } var ( ln net.Listener err error diff --git a/net/ghttp/ghttp_z_unit_feature_config_test.go b/net/ghttp/ghttp_z_unit_feature_config_test.go index ff04f7e7e..a67cd0a16 100644 --- a/net/ghttp/ghttp_z_unit_feature_config_test.go +++ b/net/ghttp/ghttp_z_unit_feature_config_test.go @@ -8,6 +8,8 @@ package ghttp_test import ( "fmt" + "github.com/gogf/gf/v2/net/gtcp" + "net" "testing" "time" @@ -23,8 +25,15 @@ import ( func Test_ConfigFromMap(t *testing.T) { gtest.C(t, func(t *gtest.T) { + p, _ := gtcp.GetFreePort() + addr := fmt.Sprintf(":%d", p) + ln, err := net.Listen("tcp", addr) + t.AssertNil(err) + listeners := []net.Listener{ln} + m := g.Map{ - "address": ":8199", + "address": addr, + "listeners": listeners, "readTimeout": "60s", "indexFiles": g.Slice{"index.php", "main.php"}, "errorLogEnabled": true, @@ -38,6 +47,7 @@ func Test_ConfigFromMap(t *testing.T) { d1, _ := time.ParseDuration(gconv.String(m["readTimeout"])) d2, _ := time.ParseDuration(gconv.String(m["cookieMaxAge"])) t.Assert(config.Address, m["address"]) + t.Assert(config.Listeners, listeners) t.Assert(config.ReadTimeout, d1) t.Assert(config.CookieMaxAge, d2) t.Assert(config.IndexFiles, m["indexFiles"]) diff --git a/net/ghttp/ghttp_z_unit_feature_custom_listeners_test.go b/net/ghttp/ghttp_z_unit_feature_custom_listeners_test.go new file mode 100644 index 000000000..47c48c229 --- /dev/null +++ b/net/ghttp/ghttp_z_unit_feature_custom_listeners_test.go @@ -0,0 +1,101 @@ +// Copyright GoFrame Author(https://goframe.org). All Rights Reserved. +// +// This Source Code Form is subject to the terms of the MIT License. +// If a copy of the MIT was not distributed with this file, +// You can obtain one at https://github.com/gogf/gf. + +package ghttp_test + +import ( + "fmt" + "github.com/gogf/gf/v2/net/gtcp" + "github.com/gogf/gf/v2/test/gtest" + "github.com/gogf/gf/v2/text/gstr" + "github.com/gogf/gf/v2/util/guid" + "net" + "testing" + "time" + + "github.com/gogf/gf/v2/frame/g" + "github.com/gogf/gf/v2/net/ghttp" +) + +func Test_SetSingleCustomListener(t *testing.T) { + p, _ := gtcp.GetFreePort() + ln, _ := net.Listen("tcp", fmt.Sprintf(":%d", p)) + s := g.Server(guid.S()) + s.Group("/", func(group *ghttp.RouterGroup) { + group.GET("/test", func(r *ghttp.Request) { + r.Response.Write("test") + }) + }) + + s.SetListener(ln) + + s.Start() + defer s.Shutdown() + + time.Sleep(100 * time.Millisecond) + + gtest.C(t, func(t *gtest.T) { + c := g.Client() + c.SetPrefix(fmt.Sprintf("http://127.0.0.1:%d", s.GetListenedPort())) + + t.Assert( + gstr.Trim(c.GetContent(ctx, "/test")), + "test", + ) + }) +} + +func Test_SetMultipleCustomListeners(t *testing.T) { + p1, _ := gtcp.GetFreePort() + p2, _ := gtcp.GetFreePort() + + ln1, _ := net.Listen("tcp", fmt.Sprintf(":%d", p1)) + ln2, _ := net.Listen("tcp", fmt.Sprintf(":%d", p2)) + + s := g.Server(guid.S()) + s.Group("/", func(group *ghttp.RouterGroup) { + group.GET("/test", func(r *ghttp.Request) { + r.Response.Write("test") + }) + }) + + s.SetListener(ln1, ln2) + + s.Start() + defer s.Shutdown() + + time.Sleep(100 * time.Millisecond) + + gtest.C(t, func(t *gtest.T) { + c := g.Client() + c.SetPrefix(fmt.Sprintf("http://127.0.0.1:%d", p1)) + + t.Assert( + gstr.Trim(c.GetContent(ctx, "/test")), + "test", + ) + + c.SetPrefix(fmt.Sprintf("http://127.0.0.1:%d", p2)) + + t.Assert( + gstr.Trim(c.GetContent(ctx, "/test")), + "test", + ) + }) +} + +func Test_SetWrongCustomListeners(t *testing.T) { + gtest.C(t, func(t *gtest.T) { + s := g.Server(guid.S()) + s.Group("/", func(group *ghttp.RouterGroup) { + group.GET("/test", func(r *ghttp.Request) { + r.Response.Write("test") + }) + }) + err := s.SetListener(nil) + t.AssertNQ(err, nil) + }) +}