Make meteredRoundTripper reconfigurable

8 files changed, 179 insertions, 74 deletions

diff --git a/internal/httprange/http_reader.go b/internal/httprange/http_reader.go
index 589351fa..ceb83ef8 100644
--- a/internal/httprange/http_reader.go
+++ b/internal/httprange/http_reader.go
@@ -54,7 +54,7 @@ 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(
 		"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..d04979a2
--- /dev/null
+++ b/internal/httptransport/metered_round_tripper.go
@@ -0,0 +1,107 @@
+package httptransport
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptrace"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	log "github.com/sirupsen/logrus"
+)
+
+// Options to configure a http.Transport
+type Options func(transport *http.Transport)
+
+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(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,
+	}
+}
+
+// WithFileProtocol option to be used while ReconfigureMeteredRoundTripper
+func WithFileProtocol(protocol string, rt http.RoundTripper) Options {
+	return func(transport *http.Transport) {
+		transport.RegisterProtocol(protocol, rt)
+	}
+}
+
+// ReconfigureMeteredRoundTripper clones meteredRoundTripper and applies options to the transport
+func ReconfigureMeteredRoundTripper(rt http.RoundTripper, opts ...Options) http.RoundTripper {
+	mrt, ok := rt.(*meteredRoundTripper)
+	if !ok {
+		return nil
+	}
+
+	t := clone(mrt.next.(*http.Transport))
+	for _, opt := range opts {
+		opt(t)
+	}
+
+	mrt.next = t
+
+	return mrt
+}
+
+// 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..c720fa83
--- /dev/null
+++ b/internal/httptransport/metered_round_tripper_test.go
@@ -0,0 +1,38 @@
+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)
+	mrt := NewMeteredRoundTripper(t.Name(), nil, histVec, counterVec, time.Millisecond)
+
+	rt := ReconfigureMeteredRoundTripper(mrt, WithFileProtocol("file", http.NewFileTransport(http.Dir("."))))
+
+	r := httptest.NewRequest("GET", "file:///testdata/", nil)
+
+	res, err := rt.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/index.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/index.html b/internal/httptransport/testdata/index.html
new file mode 100644
index 00000000..5fc555f4
--- /dev/null
+++ b/internal/httptransport/testdata/index.html
@@ -0,0 +1 @@
+httptransport/testdata/index.html
diff --git a/internal/httptransport/transport.go b/internal/httptransport/transport.go
index d8e6a3fe..3b0ec4db 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"
 )
 
@@ -33,15 +28,6 @@ var (
 	InternalTransport = newInternalTransport()
 )
 
-type meteredRoundTripper struct {
-	next        http.RoundTripper
-	name        string
-	tracer      *prometheus.HistogramVec
-	durations   *prometheus.HistogramVec
-	counter     *prometheus.CounterVec
-	ttfbTimeout time.Duration
-}
-
 func newInternalTransport() *http.Transport {
 	return &http.Transport{
 		DialTLS: func(network, addr string) (net.Conn, error) {
@@ -59,20 +45,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 +68,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_clone.go b/internal/httptransport/transport_clone.go
new file mode 100644
index 00000000..3b56250d
--- /dev/null
+++ b/internal/httptransport/transport_clone.go
@@ -0,0 +1,12 @@
+// +build go1.14
+
+package httptransport
+
+import (
+	"net/http"
+)
+
+// from go 1.14 onwards call Clone directly
+func clone(t *http.Transport) *http.Transport {
+	return t.Clone()
+}
diff --git a/internal/httptransport/transport_clone_go1_13.go b/internal/httptransport/transport_clone_go1_13.go
new file mode 100644
index 00000000..d3a91e03
--- /dev/null
+++ b/internal/httptransport/transport_clone_go1_13.go
@@ -0,0 +1,19 @@
+// +build !go1.14
+
+package httptransport
+
+import (
+    "crypto/tls"
+    "net/http"
+)
+
+func clone(t *http.Transport) *http.Transport {
+	// t.Clone() panics in Go 1.13 due to tls.Config being empty
+	// so we need to explicitly initialize it
+	// https://github.com/golang/go/issues/40565
+	if t.TLSClientConfig == nil {
+		t.TLSClientConfig = &tls.Config{}
+	}
+
+	return t.Clone()
+}
diff --git a/internal/source/gitlab/client/client.go b/internal/source/gitlab/client/client.go
index 98939641..f970c477 100644
--- a/internal/source/gitlab/client/client.go
+++ b/internal/source/gitlab/client/client.go
@@ -63,7 +63,7 @@ func NewClient(baseURL string, secretKey []byte, connectionTimeout, jwtTokenExpi
 		baseURL:   parsedURL,
 		httpClient: &http.Client{
 			Timeout: connectionTimeout,
-			Transport: httptransport.NewTransportWithMetrics(
+			Transport: httptransport.NewMeteredRoundTripper(
 				"gitlab_internal_api",
 				metrics.DomainsSourceAPITraceDuration,
 				metrics.DomainsSourceAPICallDuration,