Allow for Custom Transport on Ruler s3 Client (grafana#2891)

* stuff

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Added Middleware test

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Better naming.  Split New methods to not break existing calls

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Added top level rule storage module

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Fixed test

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Don't start ruler or rulerstorage if not configured

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Made injectRequestMiddleware a configuration option

Signed-off-by: Joe Elliott <number101010@gmail.com>

aws/s3_storage_client.go

@@ -35,6 +35,10 @@ var (
 	}, []string{"operation", "status_code"}))
 )
 
+// InjectRequestMiddleware gives users of this client the ability to make arbitrary
+// changes to outgoing requests.
+type InjectRequestMiddleware func(next http.RoundTripper) http.RoundTripper
+
 func init() {
 	s3RequestDuration.Register()
 }
@@ -52,6 +56,8 @@ type S3Config struct {
 	Insecure        bool       `yaml:"insecure"`
 	SSEEncryption   bool       `yaml:"sse_encryption"`
 	HTTPConfig      HTTPConfig `yaml:"http_config"`
+
+	Inject InjectRequestMiddleware `yaml:"-"`
 }
 
 // HTTPConfig stores the http.Transport configuration
@@ -165,22 +171,28 @@ func buildS3Config(cfg S3Config) (*aws.Config, []string, error) {
 	// to maintain backwards compatibility with previous versions of Cortex while providing
 	// more flexible configuration of the http client
 	// https://github.com/weaveworks/common/blob/4b1847531bc94f54ce5cf210a771b2a86cd34118/aws/config.go#L23
+	transport := http.RoundTripper(&http.Transport{
+		Proxy: http.ProxyFromEnvironment,
+		DialContext: (&net.Dialer{
+			Timeout:   30 * time.Second,
+			KeepAlive: 30 * time.Second,
+			DualStack: true,
+		}).DialContext,
+		MaxIdleConns:          100,
+		IdleConnTimeout:       cfg.HTTPConfig.IdleConnTimeout,
+		MaxIdleConnsPerHost:   100,
+		TLSHandshakeTimeout:   3 * time.Second,
+		ExpectContinueTimeout: 1 * time.Second,
+		ResponseHeaderTimeout: time.Duration(cfg.HTTPConfig.ResponseHeaderTimeout),
+		TLSClientConfig:       &tls.Config{InsecureSkipVerify: cfg.HTTPConfig.InsecureSkipVerify},
+	})
+
+	if cfg.Inject != nil {
+		transport = cfg.Inject(transport)
+	}
+
 	s3Config = s3Config.WithHTTPClient(&http.Client{
-		Transport: &http.Transport{
-			Proxy: http.ProxyFromEnvironment,
-			DialContext: (&net.Dialer{
-				Timeout:   30 * time.Second,
-				KeepAlive: 30 * time.Second,
-				DualStack: true,
-			}).DialContext,
-			MaxIdleConns:          100,
-			IdleConnTimeout:       cfg.HTTPConfig.IdleConnTimeout,
-			MaxIdleConnsPerHost:   100,
-			TLSHandshakeTimeout:   3 * time.Second,
-			ExpectContinueTimeout: 1 * time.Second,
-			ResponseHeaderTimeout: time.Duration(cfg.HTTPConfig.ResponseHeaderTimeout),
-			TLSClientConfig:       &tls.Config{InsecureSkipVerify: cfg.HTTPConfig.InsecureSkipVerify},
-		},
+		Transport: transport,
 	})
 
 	// bucketnames

aws/s3_storage_client_test.go

@@ -0,0 +1,77 @@
+package aws
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type RoundTripperFunc func(*http.Request) (*http.Response, error)
+
+func (f RoundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
+	return f(req)
+}
+
+func TestRequestMiddleware(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintln(w, r.Header.Get("echo-me"))
+	}))
+	defer ts.Close()
+
+	cfg := S3Config{
+		Endpoint:         ts.URL,
+		BucketNames:      "buck-o",
+		S3ForcePathStyle: true,
+		Insecure:         true,
+		AccessKeyID:      "key",
+		SecretAccessKey:  "secret",
+	}
+
+	tests := []struct {
+		name     string
+		fn       InjectRequestMiddleware
+		expected string
+	}{
+		{
+			name:     "Test Nil",
+			fn:       nil,
+			expected: "",
+		},
+		{
+			name: "Test Header Injection",
+			fn: func(next http.RoundTripper) http.RoundTripper {
+				return RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
+					req.Header["echo-me"] = []string{"blerg"}
+					return next.RoundTrip(req)
+				})
+			},
+			expected: "blerg",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cfg.Inject = tt.fn
+			client, err := NewS3ObjectClient(cfg, "/")
+			require.NoError(t, err)
+
+			readCloser, err := client.GetObject(context.Background(), "key")
+			require.NoError(t, err)
+
+			buffer := make([]byte, 100)
+			_, err = readCloser.Read(buffer)
+			if err != io.EOF {
+				require.NoError(t, err)
+			}
+
+			assert.Equal(t, tt.expected, strings.Trim(string(buffer), "\n\x00"))
+		})
+	}
+}