Merge branch '485-make-transport-reconfigurable' into 'master'

Make meteredRoundTripper take a custom transport

See merge request gitlab-org/gitlab-pages!427

9 files changed, 143 insertions, 86 deletions

diff --git a/internal/artifact/artifact.go b/internal/artifact/artifact.go
index 922aeeee..64156589 100644
--- a/internal/artifact/artifact.go
+++ b/internal/artifact/artifact.go
@@ -52,7 +52,7 @@ func New(server string, timeoutSeconds int, pagesDomain string) *Artifact {
 		suffix: "." + strings.ToLower(pagesDomain),
 		client: &http.Client{
 			Timeout:   time.Second * time.Duration(timeoutSeconds),
-			Transport: httptransport.InternalTransport,
+			Transport: httptransport.DefaultTransport,
 		},
 	}
 }
diff --git a/internal/auth/auth.go b/internal/auth/auth.go
index 215290ba..cbbc720e 100644
--- a/internal/auth/auth.go
+++ b/internal/auth/auth.go
@@ -661,7 +661,7 @@ func New(pagesDomain string, storeSecret string, clientID string, clientSecret s
 		gitLabServer: strings.TrimRight(gitLabServer, "/"),
 		apiClient: &http.Client{
 			Timeout:   5 * time.Second,
-			Transport: httptransport.InternalTransport,
+			Transport: httptransport.DefaultTransport,
 		},
 		store:         sessions.NewCookieStore(keys[0], keys[1]),
 		authSecret:    storeSecret,
diff --git a/internal/httprange/http_reader.go b/internal/httprange/http_reader.go
index 589351fa..d1523177 100644
--- a/internal/httprange/http_reader.go
+++ b/internal/httprange/http_reader.go
@@ -54,7 +54,9 @@ var _ vfs.SeekableFile = &Reader{}
 var httpClient = &http.Client{
 	// The longest time the request can be executed
 	Timeout: 30 * time.Minute,
-	Transport: httptransport.NewTransportWithMetrics(
+	Transport: httptransport.NewMeteredRoundTripper(
+		// TODO: register file protocol https://gitlab.com/gitlab-org/gitlab-pages/-/issues/485
+		nil,
 		"httprange_client",
 		metrics.HTTPRangeTraceDuration,
 		metrics.HTTPRangeRequestDuration,
diff --git a/internal/httptransport/metered_round_tripper.go b/internal/httptransport/metered_round_tripper.go
new file mode 100644
index 00000000..fc652086
--- /dev/null
+++ b/internal/httptransport/metered_round_tripper.go
@@ -0,0 +1,84 @@
+package httptransport
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptrace"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	log "github.com/sirupsen/logrus"
+)
+
+type meteredRoundTripper struct {
+	next        http.RoundTripper
+	name        string
+	tracer      *prometheus.HistogramVec
+	durations   *prometheus.HistogramVec
+	counter     *prometheus.CounterVec
+	ttfbTimeout time.Duration
+}
+
+// NewMeteredRoundTripper will create a custom http.RoundTripper that can be used with an http.Client.
+// The RoundTripper will report metrics based on the collectors passed.
+func NewMeteredRoundTripper(transport *http.Transport, name string, tracerVec, durationsVec *prometheus.
+	HistogramVec, counterVec *prometheus.CounterVec, ttfbTimeout time.Duration) http.RoundTripper {
+	if transport == nil {
+		transport = DefaultTransport
+	}
+
+	return &meteredRoundTripper{
+		next:        transport,
+		name:        name,
+		tracer:      tracerVec,
+		durations:   durationsVec,
+		counter:     counterVec,
+		ttfbTimeout: ttfbTimeout,
+	}
+}
+
+// RoundTripper wraps the original http.Transport into a meteredRoundTripper which
+// reports metrics on request duration, tracing and request count
+func (mrt *meteredRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	start := time.Now()
+
+	ctx := httptrace.WithClientTrace(r.Context(), mrt.newTracer(start))
+	ctx, cancel := context.WithCancel(ctx)
+
+	timer := time.AfterFunc(mrt.ttfbTimeout, cancel)
+	defer timer.Stop()
+
+	r = r.WithContext(ctx)
+
+	resp, err := mrt.next.RoundTrip(r)
+	if err != nil {
+		mrt.counter.WithLabelValues("error").Inc()
+		return nil, err
+	}
+
+	mrt.logResponse(r, resp)
+
+	statusCode := strconv.Itoa(resp.StatusCode)
+	mrt.durations.WithLabelValues(statusCode).Observe(time.Since(start).Seconds())
+	mrt.counter.WithLabelValues(statusCode).Inc()
+
+	return resp, nil
+}
+
+func (mrt *meteredRoundTripper) logResponse(req *http.Request, resp *http.Response) {
+	if log.GetLevel() == log.TraceLevel {
+		l := log.WithFields(log.Fields{
+			"client_name":     mrt.name,
+			"req_url":         req.URL.String(),
+			"res_status_code": resp.StatusCode,
+		})
+
+		for header, value := range resp.Header {
+			l = l.WithField(strings.ToLower(header), strings.Join(value, ";"))
+		}
+
+		l.Traceln("response")
+	}
+}
diff --git a/internal/httptransport/metered_round_tripper_test.go b/internal/httptransport/metered_round_tripper_test.go
new file mode 100644
index 00000000..218ed4a7
--- /dev/null
+++ b/internal/httptransport/metered_round_tripper_test.go
@@ -0,0 +1,39 @@
+package httptransport
+
+import (
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus/testutil"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestReconfigureMeteredRoundTripper(t *testing.T) {
+	histVec, counterVec := newTestMetrics(t)
+	transport := NewTransport()
+	transport.RegisterProtocol("file", http.NewFileTransport(http.Dir(".")))
+
+	mrt := NewMeteredRoundTripper(transport, t.Name(), nil, histVec, counterVec, time.Millisecond)
+
+	r := httptest.NewRequest("GET", "file:///testdata/file.html", nil)
+
+	res, err := mrt.RoundTrip(r)
+	require.NoError(t, err)
+	defer res.Body.Close()
+
+	require.Equal(t, http.StatusOK, res.StatusCode)
+	body, err := ioutil.ReadAll(res.Body)
+	require.NoError(t, err)
+
+	require.Equal(t, "httptransport/testdata/file.html\n", string(body))
+
+	// make sure counter still works
+	statusCode := strconv.Itoa(res.StatusCode)
+	counterCount := testutil.ToFloat64(counterVec.WithLabelValues(statusCode))
+	require.Equal(t, float64(1), counterCount, statusCode)
+}
diff --git a/internal/httptransport/testdata/file.html b/internal/httptransport/testdata/file.html
new file mode 100644
index 00000000..d4ddc151
--- /dev/null
+++ b/internal/httptransport/testdata/file.html
@@ -0,0 +1 @@
+httptransport/testdata/file.html
diff --git a/internal/httptransport/transport.go b/internal/httptransport/transport.go
index d8e6a3fe..fcadc5fe 100644
--- a/internal/httptransport/transport.go
+++ b/internal/httptransport/transport.go
@@ -1,18 +1,13 @@
 package httptransport
 
 import (
-	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"net"
 	"net/http"
-	"net/http/httptrace"
-	"strconv"
-	"strings"
 	"sync"
 	"time"
 
-	"github.com/prometheus/client_golang/prometheus"
 	log "github.com/sirupsen/logrus"
 )
 
@@ -29,20 +24,13 @@ var (
 
 	// only overridden by transport_darwin.go
 	loadExtraCerts = func() {}
-	// InternalTransport can be used with http.Client with TLS and certificates
-	InternalTransport = newInternalTransport()
+	// DefaultTransport can be used with http.Client with TLS and certificates
+	DefaultTransport = NewTransport()
 )
 
-type meteredRoundTripper struct {
-	next        http.RoundTripper
-	name        string
-	tracer      *prometheus.HistogramVec
-	durations   *prometheus.HistogramVec
-	counter     *prometheus.CounterVec
-	ttfbTimeout time.Duration
-}
-
-func newInternalTransport() *http.Transport {
+// NewTransport initializes an http.Transport with a custom dialer that includes TLS Root CAs.
+// It sets default connection values such as timeouts and max idle connections.
+func NewTransport() *http.Transport {
 	return &http.Transport{
 		DialTLS: func(network, addr string) (net.Conn, error) {
 			return tls.Dial(network, addr, &tls.Config{RootCAs: pool()})
@@ -59,20 +47,6 @@ func newInternalTransport() *http.Transport {
 	}
 }
 
-// NewTransportWithMetrics will create a custom http.RoundTripper that can be used with an http.Client.
-// The RoundTripper will report metrics based on the collectors passed.
-func NewTransportWithMetrics(name string, tracerVec, durationsVec *prometheus.
-	HistogramVec, counterVec *prometheus.CounterVec, ttfbTimeout time.Duration) http.RoundTripper {
-	return &meteredRoundTripper{
-		next:        InternalTransport,
-		name:        name,
-		tracer:      tracerVec,
-		durations:   durationsVec,
-		counter:     counterVec,
-		ttfbTimeout: ttfbTimeout,
-	}
-}
-
 // This is here because macOS does not support the SSL_CERT_FILE and
 // SSL_CERT_DIR environment variables. We have arranged things to read
 // SSL_CERT_FILE and SSL_CERT_DIR  as late as possible to avoid conflicts
@@ -96,47 +70,3 @@ func loadPool() {
 	// load them manually in OSX. See https://golang.org/src/crypto/x509/root_unix.go
 	loadExtraCerts()
 }
-
-// withRoundTripper takes an original RoundTripper, reports metrics based on the
-// gauge and counter collectors passed
-func (mrt *meteredRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
-	start := time.Now()
-
-	ctx := httptrace.WithClientTrace(r.Context(), mrt.newTracer(start))
-	ctx, cancel := context.WithCancel(ctx)
-
-	timer := time.AfterFunc(mrt.ttfbTimeout, cancel)
-	defer timer.Stop()
-
-	r = r.WithContext(ctx)
-
-	resp, err := mrt.next.RoundTrip(r)
-	if err != nil {
-		mrt.counter.WithLabelValues("error").Inc()
-		return nil, err
-	}
-
-	mrt.logResponse(r, resp)
-
-	statusCode := strconv.Itoa(resp.StatusCode)
-	mrt.durations.WithLabelValues(statusCode).Observe(time.Since(start).Seconds())
-	mrt.counter.WithLabelValues(statusCode).Inc()
-
-	return resp, nil
-}
-
-func (mrt *meteredRoundTripper) logResponse(req *http.Request, resp *http.Response) {
-	if log.GetLevel() == log.TraceLevel {
-		l := log.WithFields(log.Fields{
-			"client_name":     mrt.name,
-			"req_url":         req.URL.String(),
-			"res_status_code": resp.StatusCode,
-		})
-
-		for header, value := range resp.Header {
-			l = l.WithField(strings.ToLower(header), strings.Join(value, ";"))
-		}
-
-		l.Traceln("response")
-	}
-}
diff --git a/internal/httptransport/transport_test.go b/internal/httptransport/transport_test.go
index 9059ea15..feaf63b6 100644
--- a/internal/httptransport/transport_test.go
+++ b/internal/httptransport/transport_test.go
@@ -125,11 +125,11 @@ func (mrt *mockRoundTripper) RoundTrip(r *http.Request) (*http.Response, error)
 }
 
 func TestInternalTransportShouldHaveCustomConnectionPoolSettings(t *testing.T) {
-	require.EqualValues(t, 100, InternalTransport.MaxIdleConns)
-	require.EqualValues(t, 100, InternalTransport.MaxIdleConnsPerHost)
-	require.EqualValues(t, 0, InternalTransport.MaxConnsPerHost)
-	require.EqualValues(t, 90*time.Second, InternalTransport.IdleConnTimeout)
-	require.EqualValues(t, 10*time.Second, InternalTransport.TLSHandshakeTimeout)
-	require.EqualValues(t, 15*time.Second, InternalTransport.ResponseHeaderTimeout)
-	require.EqualValues(t, 15*time.Second, InternalTransport.ExpectContinueTimeout)
+	require.EqualValues(t, 100, DefaultTransport.MaxIdleConns)
+	require.EqualValues(t, 100, DefaultTransport.MaxIdleConnsPerHost)
+	require.EqualValues(t, 0, DefaultTransport.MaxConnsPerHost)
+	require.EqualValues(t, 90*time.Second, DefaultTransport.IdleConnTimeout)
+	require.EqualValues(t, 10*time.Second, DefaultTransport.TLSHandshakeTimeout)
+	require.EqualValues(t, 15*time.Second, DefaultTransport.ResponseHeaderTimeout)
+	require.EqualValues(t, 15*time.Second, DefaultTransport.ExpectContinueTimeout)
 }
diff --git a/internal/source/gitlab/client/client.go b/internal/source/gitlab/client/client.go
index 98939641..c5a37e17 100644
--- a/internal/source/gitlab/client/client.go
+++ b/internal/source/gitlab/client/client.go
@@ -63,7 +63,8 @@ func NewClient(baseURL string, secretKey []byte, connectionTimeout, jwtTokenExpi
 		baseURL:   parsedURL,
 		httpClient: &http.Client{
 			Timeout: connectionTimeout,
-			Transport: httptransport.NewTransportWithMetrics(
+			Transport: httptransport.NewMeteredRoundTripper(
+				httptransport.DefaultTransport,
 				"gitlab_internal_api",
 				metrics.DomainsSourceAPITraceDuration,
 				metrics.DomainsSourceAPICallDuration,