Repository: istio/fortio (go)

Issue:
add ?header= to echosrv
@mandarjog request

Pull request:
Allow custom headers to be set on query for the echo server replies
Fixes #191

You can specify multiple headers to be set/returned by the echo server
(including multiple for the same header key)

Patch:
diff --git a/README.md b/README.md
index 5dab36d..694d11b 100644
--- a/README.md
+++ b/README.md
@@ -282,6 +282,7 @@ For instance `curl -d abcdef http://localhost:8080/` returns `abcdef` back. It s
 | status    | http status to return instead of 200. Can be a single value or a coma separated list of probabilities, e.g `status=404:10,503:5,429:1` for 10% of chance of a 404 status, 5% of a 503 status and 1% of a 429 status |
 | size      | size of the payload to reply instead of echoing input. Also works as probabilities list. `size=1024:10,512:5` 10% of response will be 1k and 5% will be 512 bytes payload and the rest defaults to echoing back. |
 | close     | close the socket after answering e.g `close=true` |
+| header    | header(s) to add to the reply e.g. `&header=Foo:Bar&header=X:Y` |
 - `/debug` will echo back the request in plain text for human debugging.
 - `/fortio/` A UI to
   - Run/Trigger tests and graph the results.
diff --git a/fhttp/http_server.go b/fhttp/http_server.go
index a9e71cf..571205f 100644
--- a/fhttp/http_server.go
+++ b/fhttp/http_server.go
@@ -82,6 +82,19 @@ func EchoHandler(w http.ResponseWriter, r *http.Request) {
 		log.Debugf("Adding Connection:close / will close socket")
 		w.Header().Set("Connection", "close")
 	}
+	// process header(s) args, must be before size to compose properly
+	for _, hdr := range r.Form["header"] {
+		log.LogVf("Adding requested header %s", hdr)
+		if len(hdr) == 0 {
+			continue
+		}
+		s := strings.SplitN(hdr, ":", 2)
+		if len(s) != 2 {
+			log.Errf("invalid extra header '%s', expecting Key: Value", hdr)
+			continue
+		}
+		w.Header().Add(s[0], s[1])
+	}
 	size := generateSize(r.FormValue("size"))
 	if size >= 0 {
 		log.LogVf("Writing %d size with %d status", size, status)
diff --git a/fhttp/http_utils.go b/fhttp/http_utils.go
index df5d30b..774d41a 100644
--- a/fhttp/http_utils.go
+++ b/fhttp/http_utils.go
@@ -41,7 +41,7 @@ func toUpper(b byte) byte {
 }
 
 // ASCIIToUpper returns a byte array equal to the input string but in lowercase.
-// Only wotks for ASCII, not meant for unicode.
+// Only works for ASCII, not meant for unicode.
 func ASCIIToUpper(str string) []byte {
 	numChars := utf8.RuneCountInString(str)
 	if numChars != len(str) && log.LogVerbose() {
@@ -266,7 +266,7 @@ func generateStatus(status string) int {
 // generateSize from string, format: "size=512" for 100% 512 bytes body replies,
 // size="512:20,16384:10" for 20% 512 bytes, 10% 16k, 70% default echo back.
 // returns -1 for the default case, so one can specify 0 and force no payload
-// even if it's a post request with a payload (to test asymetric large inbound
+// even if it's a post request with a payload (to test asymmetric large inbound
 // small outbound).
 // TODO: refactor similarities with status and delay
 func generateSize(sizeInput string) (size int) {
diff --git a/fhttp/httprunner.go b/fhttp/httprunner.go
index 4b7f2c3..0da80c1 100644
--- a/fhttp/httprunner.go
+++ b/fhttp/httprunner.go
@@ -60,8 +60,8 @@ func (httpstate *HTTPRunnerResults) Run(t int) {
 	httpstate.sizes.Record(float64(size))
 	httpstate.headerSizes.Record(float64(headerSize))
 	if httpstate.AbortOn == code {
-		log.Infof("Aborting run because of code %d", code)
 		httpstate.aborter.Abort()
+		log.Infof("Aborted run because of code %d - data %s", code, DebugSummary(body, 1024))
 	}
 }
 


Test patch:
diff --git a/fhttp/http_test.go b/fhttp/http_test.go
index dde57d5..35bebba 100644
--- a/fhttp/http_test.go
+++ b/fhttp/http_test.go
@@ -728,6 +728,53 @@ func TestDebugHandlerSortedHeaders(t *testing.T) {
 	}
 }
 
+func TestEchoHeaders(t *testing.T) {
+	_, a := Serve("0", "")
+	var headers = []struct {
+		key   string
+		value string
+	}{
+		{"Foo", "Bar1"},
+		{"Foo", "Bar2"}, // Test multiple same header
+		{"X", "Y"},
+		{"Z", "abc def:xyz"},
+	}
+	v := url.Values{}
+	for _, pair := range headers {
+		v.Add("header", pair.key+":"+pair.value)
+	}
+	// minimal manual encoding (only escape the space) + errors for coverage sake
+	var urls []string
+	urls = append(urls,
+		fmt.Sprintf("http://localhost:%d/echo?size=10&header=Foo:Bar1&header=Foo:Bar2&header=X:Y&header=Z:abc+def:xyz&header=&header=Foo",
+			a.Port))
+	// proper encoding
+	urls = append(urls, fmt.Sprintf("http://localhost:%d/echo?%s", a.Port, v.Encode()))
+	for _, url := range urls {
+		resp, err := http.Get(url)
+		if err != nil {
+			t.Fatalf("Failed get for %s : %v", url, err)
+		}
+		t.Logf("TestEchoHeaders url = %s : status %s", url, resp.Status)
+		if resp.StatusCode != http.StatusOK {
+			t.Errorf("Got %d instead of 200", resp.StatusCode)
+		}
+		for _, pair := range headers {
+			got := resp.Header[pair.key]
+			found := false
+			for _, v := range got {
+				if v == pair.value {
+					found = true
+					break // found == good
+				}
+			}
+			if !found {
+				t.Errorf("Mismatch: got %+v and didn't find \"%s\" for header %s (url %s)", got, pair.value, pair.key, url)
+			}
+		}
+	}
+}
+
 func TestPPROF(t *testing.T) {
 	mux, addr := HTTPServer("test pprof", "0")
 	url := fmt.Sprintf("localhost:%d/debug/pprof/heap?debug=1", addr.Port)
@@ -872,7 +919,29 @@ func TestDefaultHeadersAndOptionsInit(t *testing.T) {
 	}
 }
 
-// --- for bench mark/comparaison
+func TestChangeMaxPayloadSize(t *testing.T) {
+	var tests = []struct {
+		input    int
+		expected int
+	}{
+		// negative test cases
+		{-1, 0},
+		// lesser than current default
+		{0, 0},
+		{64, 64},
+		// Greater than current default
+		{987 * 1024, 987 * 1024},
+	}
+	for _, tst := range tests {
+		ChangeMaxPayloadSize(tst.input)
+		actual := len(payload)
+		if len(payload) != tst.expected {
+			t.Errorf("Got %d, expected %d for ChangeMaxPayloadSize(%d)", actual, tst.expected, tst.input)
+		}
+	}
+}
+
+// --- for bench mark/comparison
 
 func asciiFold0(str string) []byte {
 	return []byte(strings.ToUpper(str))
@@ -948,24 +1017,4 @@ func BenchmarkFoldFind(b *testing.B) {
 	}
 }
 
-func TestChangeMaxPayloadSize(t *testing.T) {
-	var tests = []struct {
-		input    int
-		expected int
-	}{
-		// negative test cases
-		{-1, 0},
-		// lesser than current default
-		{0, 0},
-		{64, 64},
-		// Greater than current default
-		{987 * 1024, 987 * 1024},
-	}
-	for _, tst := range tests {
-		ChangeMaxPayloadSize(tst.input)
-		actual := len(payload)
-		if len(payload) != tst.expected {
-			t.Errorf("Got %d, expected %d for ChangeMaxPayloadSize(%d)", actual, tst.expected, tst.input)
-		}
-	}
-}
+// -- end of benchmark tests / end of this file