From 089f833ce0a14a841fbcbff33fc157cb835edf20 Mon Sep 17 00:00:00 2001
From: jiuker <2818723467@qq.com>
Date: Fri, 14 Feb 2025 02:18:53 +0800
Subject: [PATCH] feat: support --custom-header flags (#5127)

---
 cmd/client.go  | 23 ++++++++++++++++++++++-
 cmd/flags.go   |  4 ++++
 cmd/globals.go | 22 ++++++++++++++++++++++
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/cmd/client.go b/cmd/client.go
index 2bb55988a5..cb81f57348 100644
--- a/cmd/client.go
+++ b/cmd/client.go
@@ -372,7 +372,16 @@ func (config *Config) initTransport(withS3v2 bool) {
 			// 	return nil, probe.NewError(e)
 			// }
 		}
-		transport = tr
+
+		if len(globalCustomHeader) > 0 {
+			transport = &headerTransport{
+				RoundTripper: tr,
+				customHeader: globalCustomHeader.Clone(),
+			}
+		} else {
+			transport = tr
+		}
+
 	}
 
 	transport = limiter.New(config.UploadLimit, config.DownloadLimit, transport)
@@ -399,3 +408,15 @@ type SelectObjectOpts struct {
 	OutputSerOpts   map[string]map[string]string
 	CompressionType minio.SelectCompressionType
 }
+
+type headerTransport struct {
+	http.RoundTripper
+	customHeader http.Header
+}
+
+func (h *headerTransport) RoundTrip(request *http.Request) (*http.Response, error) {
+	for k, v := range h.customHeader {
+		request.Header[k] = v
+	}
+	return h.RoundTripper.RoundTrip(request)
+}
diff --git a/cmd/flags.go b/cmd/flags.go
index df1b1fd04c..6f88b259aa 100644
--- a/cmd/flags.go
+++ b/cmd/flags.go
@@ -96,6 +96,10 @@ var globalFlags = []cli.Flag{
 		Hidden: true,
 		Value:  10 * time.Minute,
 	},
+	cli.StringSliceFlag{
+		Name:  "custom-header,H",
+		Usage: "add custom HTTP header to the request. 'key:value' format.",
+	},
 }
 
 // bundled encryption flags
diff --git a/cmd/globals.go b/cmd/globals.go
index 319830db1a..cae9279ce3 100644
--- a/cmd/globals.go
+++ b/cmd/globals.go
@@ -23,6 +23,7 @@ import (
 	"crypto/x509"
 	"fmt"
 	"net"
+	"net/http"
 	"net/netip"
 	"net/url"
 	"os"
@@ -36,6 +37,7 @@ import (
 	"github.com/minio/madmin-go/v3"
 	"github.com/minio/pkg/v3/console"
 	"github.com/muesli/termenv"
+	"golang.org/x/net/http/httpguts"
 )
 
 const (
@@ -90,6 +92,8 @@ var (
 	globalLimitDownload uint64
 
 	globalContext, globalCancel = context.WithCancel(context.Background())
+
+	globalCustomHeader http.Header
 )
 
 var (
@@ -200,5 +204,23 @@ func setGlobalsFromContext(ctx *cli.Context) error {
 			globalResolvers[host] = addr
 		}
 	}
+
+	customHeaders := ctx.StringSlice("custom-header")
+	if len(customHeaders) > 0 {
+		globalCustomHeader = make(http.Header)
+		for _, header := range customHeaders {
+			i := strings.IndexByte(header, ':')
+			if i <= 0 {
+				return fmt.Errorf("invalid custom header entry %s", header)
+			}
+			h := strings.TrimSpace(header[:i])
+			hv := strings.TrimSpace(header[i+1:])
+			if !httpguts.ValidHeaderFieldName(h) || !httpguts.ValidHeaderFieldValue(hv) {
+				return fmt.Errorf("invalid custom header entry %s", header)
+			}
+			globalCustomHeader.Add(h, hv)
+		}
+	}
+
 	return nil
 }