From 40143257786aa8e41f804f58b6844282b99926a4 Mon Sep 17 00:00:00 2001
From: dustinxie <dahuaxie@gmail.com>
Date: Tue, 19 Jul 2022 09:53:24 -0700
Subject: [PATCH] [httputil] add ReadHeaderTimeout (#3550)

---
 api/http.go                        | 16 ++++------
 pkg/probe/probe.go                 |  2 +-
 pkg/util/httputil/httputil.go      | 50 ++++++++++++++++++++++++------
 pkg/util/httputil/httputil_test.go | 13 ++++----
 server/itx/server.go               |  2 +-
 5 files changed, 55 insertions(+), 28 deletions(-)

diff --git a/api/http.go b/api/http.go
index 66579ba1bd..f652fa336c 100644
--- a/api/http.go
+++ b/api/http.go
@@ -11,6 +11,7 @@ import (
 
 	apitypes "github.com/iotexproject/iotex-core/api/types"
 	"github.com/iotexproject/iotex-core/pkg/log"
+	"github.com/iotexproject/iotex-core/pkg/util/httputil"
 )
 
 type (
@@ -26,22 +27,17 @@ type (
 )
 
 // NewHTTPServer creates a new http server
-// TODO: move timeout into config
 func NewHTTPServer(route string, port int, handler http.Handler) *HTTPServer {
 	if port == 0 {
 		return nil
 	}
-	svr := &HTTPServer{
-		svr: &http.Server{
-			Addr:              ":" + strconv.Itoa(port),
-			WriteTimeout:      30 * time.Second,
-			ReadHeaderTimeout: 10 * time.Second,
-		},
-	}
 	mux := http.NewServeMux()
 	mux.Handle("/"+route, handler)
-	svr.svr.Handler = mux
-	return svr
+
+	svr := httputil.NewServer(":"+strconv.Itoa(port), mux, httputil.ReadHeaderTimeout(10*time.Second))
+	return &HTTPServer{
+		svr: &svr,
+	}
 }
 
 // Start starts the http server
diff --git a/pkg/probe/probe.go b/pkg/probe/probe.go
index d88f9a514a..6484e174af 100644
--- a/pkg/probe/probe.go
+++ b/pkg/probe/probe.go
@@ -55,7 +55,7 @@ func New(port int, opts ...Option) *Server {
 	mux.HandleFunc("/health", readiness)
 	mux.Handle("/metrics", promhttp.Handler())
 
-	s.server = httputil.Server(fmt.Sprintf(":%d", port), mux)
+	s.server = httputil.NewServer(fmt.Sprintf(":%d", port), mux)
 	return s
 }
 
diff --git a/pkg/util/httputil/httputil.go b/pkg/util/httputil/httputil.go
index d5dbea833b..5d3aa63194 100644
--- a/pkg/util/httputil/httputil.go
+++ b/pkg/util/httputil/httputil.go
@@ -10,19 +10,49 @@ import (
 
 const (
 	_connectionCount = 400
-	_readTimeout     = 35 * time.Second
-	_writeTimeout    = 35 * time.Second
-	_idleTimeout     = 120 * time.Second
 )
 
-// Server creates a HTTP server with time out settings.
-func Server(addr string, handler http.Handler) http.Server {
+type (
+	// ServerOption is a server option
+	ServerOption func(*serverConfig)
+
+	serverConfig struct {
+		ReadHeaderTimeout time.Duration
+		ReadTimeout       time.Duration
+		WriteTimeout      time.Duration
+		IdleTimeout       time.Duration
+	}
+)
+
+// DefaultServerConfig is the default server config
+var DefaultServerConfig = serverConfig{
+	ReadHeaderTimeout: 5 * time.Second,
+	ReadTimeout:       30 * time.Second,
+	WriteTimeout:      30 * time.Second,
+	IdleTimeout:       120 * time.Second,
+}
+
+// ReadHeaderTimeout sets header timeout
+func ReadHeaderTimeout(h time.Duration) ServerOption {
+	return func(cfg *serverConfig) {
+		cfg.ReadHeaderTimeout = h
+	}
+}
+
+// NewServer creates a HTTP server with time out settings.
+func NewServer(addr string, handler http.Handler, opts ...ServerOption) http.Server {
+	cfg := DefaultServerConfig
+	for _, opt := range opts {
+		opt(&cfg)
+	}
+
 	return http.Server{
-		ReadTimeout:  _readTimeout,
-		WriteTimeout: _writeTimeout,
-		IdleTimeout:  _idleTimeout,
-		Addr:         addr,
-		Handler:      handler,
+		ReadHeaderTimeout: cfg.ReadHeaderTimeout,
+		ReadTimeout:       cfg.ReadTimeout,
+		WriteTimeout:      cfg.WriteTimeout,
+		IdleTimeout:       cfg.IdleTimeout,
+		Addr:              addr,
+		Handler:           handler,
 	}
 }
 
diff --git a/pkg/util/httputil/httputil_test.go b/pkg/util/httputil/httputil_test.go
index 29189ed983..2c98dc9b78 100644
--- a/pkg/util/httputil/httputil_test.go
+++ b/pkg/util/httputil/httputil_test.go
@@ -15,13 +15,14 @@ func TestServer(t *testing.T) {
 		addr := "myAddress"
 
 		expectValue := http.Server{
-			ReadTimeout:  35 * time.Second,
-			WriteTimeout: 35 * time.Second,
-			IdleTimeout:  120 * time.Second,
-			Addr:         addr,
-			Handler:      handler,
+			ReadHeaderTimeout: 2 * time.Second,
+			ReadTimeout:       DefaultServerConfig.ReadTimeout,
+			WriteTimeout:      DefaultServerConfig.WriteTimeout,
+			IdleTimeout:       DefaultServerConfig.IdleTimeout,
+			Addr:              addr,
+			Handler:           handler,
 		}
-		result := Server(addr, handler)
+		result := NewServer(addr, handler, ReadHeaderTimeout(2*time.Second))
 		require.Equal(t, expectValue, result)
 	})
 }
diff --git a/server/itx/server.go b/server/itx/server.go
index 1c550b6886..69633b6ad3 100644
--- a/server/itx/server.go
+++ b/server/itx/server.go
@@ -243,7 +243,7 @@ func StartServer(ctx context.Context, svr *Server, probeSvr *probe.Server, cfg c
 		mux.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
 
 		port := fmt.Sprintf(":%d", cfg.System.HTTPAdminPort)
-		adminserv = httputil.Server(port, mux)
+		adminserv = httputil.NewServer(port, mux)
 		defer func() {
 			if err := adminserv.Shutdown(ctx); err != nil {
 				log.L().Error("Error when serving metrics data.", zap.Error(err))