From 93caa459e3f66792d887c1156c3258f089110bb8 Mon Sep 17 00:00:00 2001
From: Tim Gallant <me@timgallant.us>
Date: Sun, 9 Feb 2020 23:50:08 -0800
Subject: [PATCH] fs/config: add header-download and header-upload flags

---
 docs/content/flags.md                |  2 ++
 fs/config.go                         |  2 ++
 fs/config/configflags/configflags.go | 29 ++++++++++++++++++++++++++++
 3 files changed, 33 insertions(+)

diff --git a/docs/content/flags.md b/docs/content/flags.md
index 47001efa3..647829fcc 100755
--- a/docs/content/flags.md
+++ b/docs/content/flags.md
@@ -48,6 +48,8 @@ These flags are available for every command.
       --files-from stringArray               Read list of source-file names from file (use - to read from stdin).
   -f, --filter stringArray                   Add a file-filtering rule
       --filter-from stringArray              Read filtering patterns from a file (use - to read from stdin).
+      --header-download stringArray          Add HTTP header for download transactions
+      --header-upload stringArray            Add HTTP header for upload transactions
       --ignore-case                          Ignore case in filters (case insensitive)
       --ignore-case-sync                     Ignore case when synchronizing
       --ignore-checksum                      Skip post copy check of checksums.
diff --git a/fs/config.go b/fs/config.go
index 4e3cc6a60..555085c34 100644
--- a/fs/config.go
+++ b/fs/config.go
@@ -111,6 +111,8 @@ type ConfigInfo struct {
 	MultiThreadStreams     int
 	MultiThreadSet         bool   // whether MultiThreadStreams was set (set in fs/config/configflags)
 	OrderBy                string // instructions on how to order the transfer
+	UploadHeaders          []*HTTPOption
+	DownloadHeaders        []*HTTPOption
 }
 
 // NewConfig creates a new config with everything set to the default
diff --git a/fs/config/configflags/configflags.go b/fs/config/configflags/configflags.go
index d142f42c6..f06f83ec8 100644
--- a/fs/config/configflags/configflags.go
+++ b/fs/config/configflags/configflags.go
@@ -29,6 +29,8 @@ var (
 	deleteAfter     bool
 	bindAddr        string
 	disableFeatures string
+	uploadHeaders   []string
+	downloadHeaders []string
 )
 
 // AddFlags adds the non filing system specific flags to the command
@@ -112,6 +114,25 @@ func AddFlags(flagSet *pflag.FlagSet) {
 	flags.IntVarP(flagSet, &fs.Config.MultiThreadStreams, "multi-thread-streams", "", fs.Config.MultiThreadStreams, "Max number of streams to use for multi-thread downloads.")
 	flags.BoolVarP(flagSet, &fs.Config.UseJSONLog, "use-json-log", "", fs.Config.UseJSONLog, "Use json log format.")
 	flags.StringVarP(flagSet, &fs.Config.OrderBy, "order-by", "", fs.Config.OrderBy, "Instructions on how to order the transfers, eg 'size,descending'")
+	flags.StringArrayVarP(flagSet, &uploadHeaders, "header-upload", "", nil, "Set HTTP header for upload transactions")
+	flags.StringArrayVarP(flagSet, &downloadHeaders, "header-download", "", nil, "Set HTTP header for download transactions")
+}
+
+// ParseHeaders converts the strings passed in via the header flags into HTTPOptions
+func ParseHeaders(headers []string) []*fs.HTTPOption {
+	opts := []*fs.HTTPOption{}
+	for _, header := range headers {
+		parts := strings.SplitN(header, ":", 2)
+		if len(parts) == 1 {
+			log.Fatalf("Failed to parse '%s' as an HTTP header. Expecting a string like: 'Content-Encoding: gzip'", header)
+		}
+		option := &fs.HTTPOption{
+			Key:   strings.TrimSpace(parts[0]),
+			Value: strings.TrimSpace(parts[1]),
+		}
+		opts = append(opts, option)
+	}
+	return opts
 }
 
 // SetFlags converts any flags into config which weren't straight forward
@@ -212,6 +233,14 @@ func SetFlags() {
 		fs.Config.DisableFeatures = strings.Split(disableFeatures, ",")
 	}
 
+	if len(uploadHeaders) != 0 {
+		fs.Config.UploadHeaders = ParseHeaders(uploadHeaders)
+	}
+
+	if len(downloadHeaders) != 0 {
+		fs.Config.DownloadHeaders = ParseHeaders(downloadHeaders)
+	}
+
 	// Make the config file absolute
 	configPath, err := filepath.Abs(config.ConfigPath)
 	if err == nil {