From 5851c39776c6be082eb1efe179f456b5c645f491 Mon Sep 17 00:00:00 2001
From: Marcin Wolny <marcin@0xproject.com>
Date: Fri, 15 Dec 2023 22:29:28 +0100
Subject: [PATCH] Hello http.TimeoutHandler

Prefer http.TimeoutHandler instead ResponseHeaderTimeout. This way, we
can simplify the interface and stop use UpstreamTimeut upon ReverseProxy
creation.
---
 internal/proxy/proxy.go        | 16 +++++++++++++---
 internal/proxy/proxy_test.go   |  3 ++-
 internal/proxy/reverseproxy.go |  3 +--
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go
index 892191e..82b99b0 100644
--- a/internal/proxy/proxy.go
+++ b/internal/proxy/proxy.go
@@ -84,7 +84,7 @@ func NewProxy(proxyConfig Config, healthCheckManager *HealthcheckManager) *Proxy
 }
 
 func (p *Proxy) AddTarget(target TargetConfig) error {
-	proxy, err := NewReverseProxy(target, p.config)
+	proxy, err := NewReverseProxy(target)
 	if err != nil {
 		return err
 	}
@@ -113,6 +113,16 @@ func (p *Proxy) copyHeaders(dst http.ResponseWriter, src http.ResponseWriter) {
 	}
 }
 
+func (p *Proxy) timeoutHandler(next http.Handler) http.Handler {
+	fn := func(w http.ResponseWriter, r *http.Request) {
+		http.TimeoutHandler(next,
+			p.config.Proxy.UpstreamTimeout,
+			http.StatusText(http.StatusGatewayTimeout)).ServeHTTP(w, r)
+	}
+
+	return http.HandlerFunc(fn)
+}
+
 func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	body := &bytes.Buffer{}
 
@@ -127,9 +137,9 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		r.Body = io.NopCloser(bytes.NewBuffer(body.Bytes()))
 
 		if !target.Config.Connection.HTTP.Compression && strings.Contains(r.Header.Get(headers.ContentEncoding), "gzip") {
-			middleware.Gunzip(target.Proxy).ServeHTTP(pw, r)
+			p.timeoutHandler(middleware.Gunzip(target.Proxy)).ServeHTTP(pw, r)
 		} else {
-			target.Proxy.ServeHTTP(pw, r)
+			p.timeoutHandler(target.Proxy).ServeHTTP(pw, r)
 		}
 
 		if p.HasNodeProviderFailed(pw.statusCode) {
diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go
index 840a008..16c292b 100644
--- a/internal/proxy/proxy_test.go
+++ b/internal/proxy/proxy_test.go
@@ -8,6 +8,7 @@ import (
 	"net/http/httptest"
 	"strconv"
 	"testing"
+	"time"
 
 	"github.com/go-http-utils/headers"
 	"github.com/prometheus/client_golang/prometheus"
@@ -17,7 +18,7 @@ import (
 func createConfig() Config {
 	return Config{
 		Proxy: ProxyConfig{
-			UpstreamTimeout: 0,
+			UpstreamTimeout: time.Second * 3,
 		},
 		HealthChecks: HealthCheckConfig{
 			Interval:         0,
diff --git a/internal/proxy/reverseproxy.go b/internal/proxy/reverseproxy.go
index 52d4101..2b768cc 100644
--- a/internal/proxy/reverseproxy.go
+++ b/internal/proxy/reverseproxy.go
@@ -11,7 +11,7 @@ import (
 	"github.com/pkg/errors"
 )
 
-func NewReverseProxy(targetConfig TargetConfig, config Config) (*httputil.ReverseProxy, error) {
+func NewReverseProxy(targetConfig TargetConfig) (*httputil.ReverseProxy, error) {
 	target, err := url.Parse(targetConfig.Connection.HTTP.URL)
 	if err != nil {
 		return nil, errors.Wrap(err, "cannot parse url")
@@ -42,7 +42,6 @@ func NewReverseProxy(targetConfig TargetConfig, config Config) (*httputil.Revers
 		IdleConnTimeout:       30 * time.Second,
 		TLSHandshakeTimeout:   10 * time.Second,
 		ExpectContinueTimeout: 1 * time.Second,
-		ResponseHeaderTimeout: config.Proxy.UpstreamTimeout,
 	}
 
 	conntrack.PreRegisterDialerMetrics(targetConfig.Name)