From fff4bf98b0a6c0bea4f56fee035e3bf46577cea2 Mon Sep 17 00:00:00 2001
From: hayson <hayson97@qq.com>
Date: Fri, 15 Nov 2024 01:14:33 +0800
Subject: [PATCH] feat: add WithServer option (#25)

---
 graceful.go      | 10 ++++++++++
 graceful_test.go | 19 ++++++++++++++++++-
 options.go       | 23 +++++++++++++++++++++++
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/graceful.go b/graceful.go
index 6787e0f..77b1ea9 100644
--- a/graceful.go
+++ b/graceful.go
@@ -274,6 +274,16 @@ func (g *Graceful) appendHTTPServer() *http.Server {
 	return srv
 }
 
+// appendExistHTTPServer appends an existing HTTP server to the list of servers managed by the Graceful instance.
+// This allows for customization of the http.Server, and srv.Handler will be set to the current g.Engine.
+func (g *Graceful) appendExistHTTPServer(srv *http.Server) {
+	srv.Handler = g.Engine
+
+	g.lock.Lock()
+	defer g.lock.Unlock()
+	g.servers = append(g.servers, srv)
+}
+
 // ensureAtLeastDefaultServer ensures that there is at least one server running with the default address ":8080".
 // If no server is running, it creates a new server with the default address and starts it.
 // It returns an error if there was a problem creating or starting the server.
diff --git a/graceful_test.go b/graceful_test.go
index d26ca37..6894607 100644
--- a/graceful_test.go
+++ b/graceful_test.go
@@ -187,6 +187,20 @@ func TestWithListener(t *testing.T) {
 	}, fmt.Sprintf("http://localhost:%d/example", listener.Addr().(*net.TCPAddr).Port))
 }
 
+func TestWithServer(t *testing.T) {
+	cert, err := tls.LoadX509KeyPair("./testdata/certificate/cert.pem", "./testdata/certificate/key.pem")
+	assert.NoError(t, err)
+	tlsConfig := &tls.Config{
+		Certificates: []tls.Certificate{cert},
+	}
+	testRouterConstructor(t, func() (*Graceful, error) {
+		return Default(
+			WithServer(&http.Server{Addr: ":8811"}),
+			WithServer(&http.Server{Addr: ":9443", TLSConfig: tlsConfig}),
+		)
+	}, "http://localhost:8811/example", "https://localhost:9443/example")
+}
+
 func TestWithAll(t *testing.T) {
 	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
 	assert.NoError(t, err)
@@ -198,11 +212,14 @@ func TestWithAll(t *testing.T) {
 		return Default(WithAddr(":8080"),
 			WithTLS(":8443", "./testdata/certificate/cert.pem", "./testdata/certificate/key.pem"),
 			WithListener(listener),
+			WithServer(&http.Server{Addr: ":8811"}),
 		)
 	},
 		"http://localhost:8080/example",
 		"https://localhost:8443/example",
-		fmt.Sprintf("http://localhost:%d/example", listener.Addr().(*net.TCPAddr).Port))
+		fmt.Sprintf("http://localhost:%d/example", listener.Addr().(*net.TCPAddr).Port),
+		"http://localhost:8811/example",
+	)
 }
 
 func TestWithContext(t *testing.T) {
diff --git a/options.go b/options.go
index 054e46d..da49b71 100644
--- a/options.go
+++ b/options.go
@@ -1,8 +1,10 @@
 package graceful
 
 import (
+	"errors"
 	"fmt"
 	"net"
+	"net/http"
 	"os"
 )
 
@@ -48,6 +50,27 @@ func WithTLS(addr string, certFile string, keyFile string) Option {
 	})
 }
 
+// WithServer configure an existing http.Server to serve HTTP or HTTPS requests.
+// This allows for a more complete customization of the http.Server,
+// and srv Handler will be set to the current gin.Engine.
+// If srv contains TLSConfig, ListenAndServeTLS will be used;
+// otherwise, ListenAndServe will be used.
+func WithServer(srv *http.Server) Option {
+	return optionFunc(func(g *Graceful) (listenAndServe, cleanup, error) {
+		if srv == nil {
+			return nil, donothing, errors.New("nil http server")
+		}
+		return func() error {
+			g.appendExistHTTPServer(srv)
+			if srv.TLSConfig == nil {
+				return srv.ListenAndServe()
+			} else {
+				return srv.ListenAndServeTLS("", "")
+			}
+		}, donothing, nil
+	})
+}
+
 // WithUnix configure a http.Server to listen on the given unix socket file.
 func WithUnix(file string) Option {
 	return optionFunc(func(g *Graceful) (listenAndServe, cleanup, error) {