From cbe75d373642ef9ff64ca76dc7175dec050556e3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Giedrius=20Statkevi=C4=8Dius?= <giedriuswork@gmail.com>
Date: Thu, 17 Oct 2024 15:57:51 +0300
Subject: [PATCH] Wrapping transport

---
 client/factory.go  | 14 ++++++++++++--
 providers/s3/s3.go |  8 +++-----
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/client/factory.go b/client/factory.go
index bd345024..c7baf0be 100644
--- a/client/factory.go
+++ b/client/factory.go
@@ -48,9 +48,19 @@ type BucketConfig struct {
 	Prefix string      `yaml:"prefix" default:""`
 }
 
+/*
+func(rt http.RoundTripper) http.RoundTripper {
+  if rt == nil {
+    rt = http.DefaultTransport
+  }
+
+  return wrapHedgedRoundtripper(rt)
+}
+*/
+
 // NewBucket initializes and returns new object storage clients.
 // NOTE: confContentYaml can contain secrets.
-func NewBucket(logger log.Logger, confContentYaml []byte, component string, rt http.RoundTripper) (objstore.Bucket, error) {
+func NewBucket(logger log.Logger, confContentYaml []byte, component string, wrapRoundtripper func(http.RoundTripper) http.RoundTripper) (objstore.Bucket, error) {
 	level.Info(logger).Log("msg", "loading bucket configuration")
 	bucketConf := &BucketConfig{}
 	if err := yaml.UnmarshalStrict(confContentYaml, bucketConf); err != nil {
@@ -67,7 +77,7 @@ func NewBucket(logger log.Logger, confContentYaml []byte, component string, rt h
 	case string(GCS):
 		bucket, err = gcs.NewBucket(context.Background(), logger, config, component, rt)
 	case string(S3):
-		bucket, err = s3.NewBucket(logger, config, component, rt)
+		bucket, err = s3.NewBucket(logger, config, component, wrapRoundtripper)
 	case string(AZURE):
 		bucket, err = azure.NewBucket(logger, config, component, rt)
 	case string(SWIFT):
diff --git a/providers/s3/s3.go b/providers/s3/s3.go
index eac8191a..afe6fe91 100644
--- a/providers/s3/s3.go
+++ b/providers/s3/s3.go
@@ -176,7 +176,7 @@ func parseConfig(conf []byte) (Config, error) {
 }
 
 // NewBucket returns a new Bucket using the provided s3 config values.
-func NewBucket(logger log.Logger, conf []byte, component string, rt http.RoundTripper) (*Bucket, error) {
+func NewBucket(logger log.Logger, conf []byte, component string, wrapRoundtripper func(http.RoundTripper) http.RoundTripper) (*Bucket, error) {
 	config, err := parseConfig(conf)
 	if err != nil {
 		return nil, err
@@ -202,7 +202,7 @@ func (s *overrideSignerType) Retrieve() (credentials.Value, error) {
 }
 
 // NewBucketWithConfig returns a new Bucket using the provided s3 config values.
-func NewBucketWithConfig(logger log.Logger, config Config, component string, rt http.RoundTripper) (*Bucket, error) {
+func NewBucketWithConfig(logger log.Logger, config Config, component string, wrapRoundtripper func(http.RoundTripper) http.RoundTripper) (*Bucket, error) {
 	var chain []credentials.Provider
 
 	// TODO(bwplotka): Don't do flags as they won't scale, use actual params like v2, v4 instead
@@ -242,9 +242,6 @@ func NewBucketWithConfig(logger log.Logger, config Config, component string, rt
 			}),
 		}
 	}
-	if rt != nil {
-		config.HTTPConfig.Transport = rt
-	}
 	// Check if a roundtripper has been set in the config
 	// otherwise build the default transport.
 	var tpt http.RoundTripper
@@ -255,6 +252,7 @@ func NewBucketWithConfig(logger log.Logger, config Config, component string, rt
 	if config.HTTPConfig.Transport != nil {
 		tpt = config.HTTPConfig.Transport
 	}
+	tpt = wrapRoundtripper(tpt)
 
 	client, err := minio.New(config.Endpoint, &minio.Options{
 		Creds:        credentials.NewChainCredentials(chain),