From 0ed10416a70c06f01bcd498db856edd0280311f2 Mon Sep 17 00:00:00 2001 From: Joe Elliott Date: Tue, 21 Jul 2020 13:27:46 -0400 Subject: [PATCH] Allow for Custom Transport on Ruler s3 Client (#2891) * stuff Signed-off-by: Joe Elliott * Added Middleware test Signed-off-by: Joe Elliott * Better naming. Split New methods to not break existing calls Signed-off-by: Joe Elliott * Added top level rule storage module Signed-off-by: Joe Elliott * Fixed test Signed-off-by: Joe Elliott * Don't start ruler or rulerstorage if not configured Signed-off-by: Joe Elliott * Made injectRequestMiddleware a configuration option Signed-off-by: Joe Elliott --- aws/s3_storage_client.go | 42 ++++++++++++------- aws/s3_storage_client_test.go | 77 +++++++++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 15 deletions(-) create mode 100644 aws/s3_storage_client_test.go diff --git a/aws/s3_storage_client.go b/aws/s3_storage_client.go index dd0e72bf8291c..eb35948f8abb6 100644 --- a/aws/s3_storage_client.go +++ b/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 diff --git a/aws/s3_storage_client_test.go b/aws/s3_storage_client_test.go new file mode 100644 index 0000000000000..e5bfd5a9a4a99 --- /dev/null +++ b/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")) + }) + } +}