From cc0d351e696bd01aebc76a54ea0546f05bda449e Mon Sep 17 00:00:00 2001
From: Alex <alexellis2@gmail.com>
Date: Tue, 9 May 2017 09:13:42 +0100
Subject: [PATCH] Watchdog - enable forwarding of HTTP headers via CGI-like env
 vars.

---
 watchdog/main.go                | 10 +++++
 watchdog/readconfig.go          |  5 +++
 watchdog/requesthandler_test.go | 73 +++++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+)

diff --git a/watchdog/main.go b/watchdog/main.go
index 387b00df..66dcd8a7 100644
--- a/watchdog/main.go
+++ b/watchdog/main.go
@@ -51,6 +51,16 @@ func pipeRequest(config *WatchdogConfig, w http.ResponseWriter, r *http.Request)
 	}
 
 	targetCmd := exec.Command(parts[0], parts[1:]...)
+
+	if config.cgiHeaders {
+		envs := os.Environ()
+		for k, v := range r.Header {
+			kv := fmt.Sprintf("Http_%s=%s", k, v[0])
+			envs = append(envs, kv)
+		}
+		targetCmd.Env = envs
+	}
+
 	writer, _ := targetCmd.StdinPipe()
 
 	var out []byte
diff --git a/watchdog/readconfig.go b/watchdog/readconfig.go
index f188e284..2c999709 100644
--- a/watchdog/readconfig.go
+++ b/watchdog/readconfig.go
@@ -59,6 +59,8 @@ func (ReadConfig) Read(hasEnv HasEnv) WatchdogConfig {
 		cfg.writeDebug = parseBoolValue(hasEnv.Getenv("write_debug"))
 	}
 
+	cfg.cgiHeaders = parseBoolValue(hasEnv.Getenv("cgi_headers"))
+
 	cfg.marshalRequest = parseBoolValue(hasEnv.Getenv("marshal_request"))
 	cfg.debugHeaders = parseBoolValue(hasEnv.Getenv("debug_headers"))
 
@@ -82,6 +84,9 @@ type WatchdogConfig struct {
 
 	marshalRequest bool
 
+	// cgiHeaders will make environmental variables available with all the HTTP headers.
+	cgiHeaders bool
+
 	// prints out all incoming and out-going HTTP headers
 	debugHeaders bool
 
diff --git a/watchdog/requesthandler_test.go b/watchdog/requesthandler_test.go
index 87adebaa..c2a439f2 100644
--- a/watchdog/requesthandler_test.go
+++ b/watchdog/requesthandler_test.go
@@ -8,6 +8,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 )
 
@@ -20,6 +21,78 @@ func TestHandler_make(t *testing.T) {
 	}
 }
 
+func TestHandler_HasCustomHeaderInFunction_WithCgi_Mode(t *testing.T) {
+	rr := httptest.NewRecorder()
+
+	body := ""
+	req, err := http.NewRequest("POST", "/", bytes.NewBufferString(body))
+	req.Header.Add("custom-header", "value")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	config := WatchdogConfig{
+		faasProcess: "env",
+		cgiHeaders:  true,
+	}
+	handler := makeRequestHandler(&config)
+	handler(rr, req)
+
+	required := http.StatusOK
+	if status := rr.Code; status != required {
+		t.Errorf("handler returned wrong status code: got %v, but wanted %v",
+			status, required)
+	}
+
+	read, _ := ioutil.ReadAll(rr.Body)
+	val := string(read)
+	if strings.Index(val, "Http_Custom-Header") == -1 {
+		t.Errorf("'env' should printed: Http_Custom-Header, got: %s\n", val)
+
+	}
+
+	seconds := rr.Header().Get("X-Duration-Seconds")
+	if len(seconds) == 0 {
+		t.Errorf("Exec of cat should have given a duration as an X-Duration-Seconds header\n")
+	}
+}
+
+func TestHandler_DoesntHaveCustomHeaderInFunction_WithoutCgi_Mode(t *testing.T) {
+	rr := httptest.NewRecorder()
+
+	body := ""
+	req, err := http.NewRequest("POST", "/", bytes.NewBufferString(body))
+	req.Header.Add("custom-header", "value")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	config := WatchdogConfig{
+		faasProcess: "env",
+		cgiHeaders:  false,
+	}
+	handler := makeRequestHandler(&config)
+	handler(rr, req)
+
+	required := http.StatusOK
+	if status := rr.Code; status != required {
+		t.Errorf("handler returned wrong status code: got %v, but wanted %v",
+			status, required)
+	}
+
+	read, _ := ioutil.ReadAll(rr.Body)
+	val := string(read)
+	if strings.Index(val, "Http_Custom-Header") != -1 {
+		t.Errorf("'env' should not have printed: Http_Custom-Header, got: %s\n", val)
+
+	}
+
+	seconds := rr.Header().Get("X-Duration-Seconds")
+	if len(seconds) == 0 {
+		t.Errorf("Exec of cat should have given a duration as an X-Duration-Seconds header\n")
+	}
+}
+
 func TestHandler_HasXDurationSecondsHeader(t *testing.T) {
 	rr := httptest.NewRecorder()