From f958e99bdc1b5c6157e6582db6fa030fbd29e151 Mon Sep 17 00:00:00 2001
From: Alex Ellis <alexellis2@gmail.com>
Date: Mon, 10 Apr 2017 22:18:37 +0100
Subject: [PATCH] Allow override of contentType response from watchdog. (#46)

---
 watchdog/README.md      |  4 ++--
 watchdog/config_test.go | 13 +++++++++++++
 watchdog/main.go        | 12 +++++++++---
 watchdog/readconfig.go  |  5 +++++
 4 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/watchdog/README.md b/watchdog/README.md
index ec299038..87126e09 100644
--- a/watchdog/README.md
+++ b/watchdog/README.md
@@ -53,6 +53,7 @@ A number of environmental overrides can be added for additional flexibility and
 |------------------------|--------------|
 | `fprocess`             | The process to invoke for each function call. This must be a UNIX binary and accept input via STDIN and output via STDOUT.  |
 | `marshal_requests`     | Instead of re-directing the raw HTTP body into your fprocess, it will first be marshalled into JSON. Use this if you need to work with HTTP headers |
+| `content_type`         | Force a specific Content-Type response for all responses.    |
 | `write_timeout`        | HTTP timeout for writing a response body from your function  |
 | `read_timeout`         | HTTP timeout for reading the payload from the client caller  |
 | `suppress_lock`        | The watchdog will attempt to write a lockfile to /tmp/ for swarm healthchecks - set this to true to disable behaviour. |
@@ -66,8 +67,7 @@ By default the watchdog will match the response of your function to the "Content
 * If your client sends a JSON post with a Content-Type of `application/json` this will be matched automatically in the response.
 * If your client sends a JSON post with a Content-Type of `text/plain` this will be matched automatically in the response too
 
-See open issues for custom override of response ContentType.
-
+To override the Content-Type of all your responses set the `content_type` environmental variable.
 
 **Tuning auto-scaling**
 
diff --git a/watchdog/config_test.go b/watchdog/config_test.go
index 71a9e02c..af373068 100644
--- a/watchdog/config_test.go
+++ b/watchdog/config_test.go
@@ -73,6 +73,19 @@ func TestRead_SuppressLockConfig(t *testing.T) {
 	}
 }
 
+func TestRead_ContentTypeConfig(t *testing.T) {
+	defaults := NewEnvBucket()
+	readConfig := ReadConfig{}
+	defaults.Setenv("content_type", "application/json")
+
+	config := readConfig.Read(defaults)
+
+	if config.contentType != "application/json" {
+		t.Logf("content_type envVariable incorrect, got: %s.\n", config.contentType)
+		t.Fail()
+	}
+}
+
 func TestRead_FprocessConfig(t *testing.T) {
 	defaults := NewEnvBucket()
 	readConfig := ReadConfig{}
diff --git a/watchdog/main.go b/watchdog/main.go
index 076d61f3..e05ba788 100644
--- a/watchdog/main.go
+++ b/watchdog/main.go
@@ -89,9 +89,15 @@ func pipeRequest(config *WatchdogConfig, w http.ResponseWriter, r *http.Request)
 		os.Stdout.Write(out)
 	}
 
-	clientContentType := r.Header.Get("Content-Type")
-	if len(clientContentType) > 0 {
-		w.Header().Set("Content-Type", "application/json")
+	if len(config.contentType) > 0 {
+		w.Header().Set("Content-Type", config.contentType)
+	} else {
+
+		// Match content-type of caller if no override specified.
+		clientContentType := r.Header.Get("Content-Type")
+		if len(clientContentType) > 0 {
+			w.Header().Set("Content-Type", clientContentType)
+		}
 	}
 
 	w.WriteHeader(200)
diff --git a/watchdog/readconfig.go b/watchdog/readconfig.go
index de47485a..f188e284 100644
--- a/watchdog/readconfig.go
+++ b/watchdog/readconfig.go
@@ -64,6 +64,8 @@ func (ReadConfig) Read(hasEnv HasEnv) WatchdogConfig {
 
 	cfg.suppressLock = parseBoolValue(hasEnv.Getenv("suppress_lock"))
 
+	cfg.contentType = hasEnv.Getenv("content_type")
+
 	return cfg
 }
 
@@ -85,4 +87,7 @@ type WatchdogConfig struct {
 
 	// Don't write a lock file to /tmp/
 	suppressLock bool
+
+	// contentType forces a specific pre-defined value for all responses
+	contentType string
 }