From a83f96e46ff97d9fd9a945cb39e0a3a521c2250a Mon Sep 17 00:00:00 2001
From: gyakkun <xiaowufeixia@gmail.com>
Date: Sat, 31 Aug 2024 16:19:08 +0800
Subject: [PATCH] Try to support multiple custom headers.

---
 cmd/cassowary/cli.go | 15 +++++++++------
 pkg/client/load.go   |  9 +++++++--
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/cmd/cassowary/cli.go b/cmd/cassowary/cli.go
index 9bf1b44..f341832 100644
--- a/cmd/cassowary/cli.go
+++ b/cmd/cassowary/cli.go
@@ -132,11 +132,14 @@ func validateCLI(c *cli.Context) error {
 		prometheusEnabled = true
 	}
 
-	if c.String("header") != "" {
-		length := 0
-		length, header = client.SplitHeader(c.String("header"))
-		if length != 2 {
-			return errNotValidHeader
+	if len(c.StringSlice("header")) != 0 {
+		allHeaders := c.StringSlice("header")
+		for _, hdr := range allHeaders {
+			thisLen, thisHdr := client.SplitHeader(hdr)
+			if thisLen != 2 {
+				return errNotValidHeader
+			}
+			header = append(header, thisHdr...)
 		}
 	}
 
@@ -279,7 +282,7 @@ func runCLI(args []string) {
 					Aliases: []string{"prompushgwurl"},
 					Usage:   "specify prometheus push gateway url to send metrics (optional)",
 				},
-				&cli.StringFlag{
+				&cli.StringSliceFlag{
 					Name:    "H",
 					Aliases: []string{"header"},
 					Usage:   "add arbitrary header, eg. 'Host: www.example.com'",
diff --git a/pkg/client/load.go b/pkg/client/load.go
index 1a2f23f..c4c71d8 100644
--- a/pkg/client/load.go
+++ b/pkg/client/load.go
@@ -68,8 +68,13 @@ func (c *Cassowary) runLoadTest(outPutChan chan<- durationMetrics, workerChan ch
 			}
 		}
 
-		if len(c.RequestHeader) == 2 {
-			request.Header.Add(c.RequestHeader[0], c.RequestHeader[1])
+		if len(c.RequestHeader)%2 == 0 {
+			for idx := range c.RequestHeader {
+				if idx%2 == 1 {
+					continue
+				}
+				request.Header.Add(c.RequestHeader[idx], c.RequestHeader[idx+1])
+			}
 		}
 
 		var t0, t1, t2, t3, t4, t5, t6 time.Time