Adds env features from original watchdog

Signed-off-by: Thomas Fan <[email protected]>

executor/serializing_fork_runner.go

@@ -1,6 +1,7 @@
 package executor
 
 import (
+	"fmt"
 	"io"
 	"io/ioutil"
 	"log"
@@ -17,13 +18,14 @@ type SerializingForkFunctionRunner struct {
 
 // Run run a fork for each invocation
 func (f *SerializingForkFunctionRunner) Run(req FunctionRequest, w http.ResponseWriter) error {
-	functionBytes, err := serializeFunction(req, f)
+	functionBytes, execDuration, err := serializeFunction(req, f)
 	if err != nil {
 		w.WriteHeader(500)
 		w.Write([]byte(err.Error()))
 		return err
 	}
 
+	w.Header().Set("X-Duration-Seconds", fmt.Sprintf("%f", execDuration))
 	w.WriteHeader(200)
 
 	if functionBytes != nil {
@@ -35,7 +37,7 @@ func (f *SerializingForkFunctionRunner) Run(req FunctionRequest, w http.Response
 	return err
 }
 
-func serializeFunction(req FunctionRequest, f *SerializingForkFunctionRunner) (*[]byte, error) {
+func serializeFunction(req FunctionRequest, f *SerializingForkFunctionRunner) (*[]byte, float64, error) {
 	log.Printf("Running %s", req.Process)
 
 	start := time.Now()
@@ -71,7 +73,7 @@ func serializeFunction(req FunctionRequest, f *SerializingForkFunctionRunner) (*
 		data, err = ioutil.ReadAll(limitReader)
 
 		if err != nil {
-			return nil, err
+			return nil, 0, err
 		}
 
 	}
@@ -81,24 +83,24 @@ func serializeFunction(req FunctionRequest, f *SerializingForkFunctionRunner) (*
 
 	err := cmd.Start()
 	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 
 	functionRes, errors := pipeToProcess(stdin, stdout, &data)
 
 	if len(errors) > 0 {
-		return nil, errors[0]
+		return nil, 0, errors[0]
 	}
 
 	waitErr := cmd.Wait()
 	if waitErr != nil {
-		return nil, err
+		return nil, 0, err
 	}
 
-	done := time.Since(start)
-	log.Printf("Took %f secs", done.Seconds())
+	done := time.Since(start).Seconds()
+	log.Printf("Took %f secs", done)
 
-	return functionRes, nil
+	return functionRes, done, nil
 }
 
 func pipeToProcess(stdin io.WriteCloser, stdout io.Reader, data *[]byte) (*[]byte, []error) {

main.go

@@ -262,6 +262,7 @@ func getEnvironment(r *http.Request) []string {
 		envs = append(envs, kv)
 	}
 	envs = append(envs, fmt.Sprintf("Http_Method=%s", r.Method))
+	envs = append(envs, fmt.Sprintf("Http_ContentLength=%d", r.ContentLength))
 
 	if len(r.URL.RawQuery) > 0 {
 		envs = append(envs, fmt.Sprintf("Http_Query=%s", r.URL.RawQuery))
@@ -271,6 +272,10 @@ func getEnvironment(r *http.Request) []string {
 		envs = append(envs, fmt.Sprintf("Http_Path=%s", r.URL.Path))
 	}
 
+	if len(r.Host) > 0 {
+		envs = append(envs, fmt.Sprintf("Http_Host=%s", r.Host))
+	}
+
 	return envs
 }
 

serializingforkrequesthandler_test.go

@@ -0,0 +1,215 @@
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/openfaas-incubator/of-watchdog/config"
+)
+
+func TestSerializingForkHandler_HasCustomHeaderInFunction_WithCGI(t *testing.T) {
+	rr := httptest.NewRecorder()
+
+	body := ""
+	req, err := http.NewRequest(http.MethodPost, "/", bytes.NewBufferString(body))
+	req.Header.Add("custom-header", "value")
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	config := config.WatchdogConfig{
+		FunctionProcess:  "env",
+		InjectCGIHeaders: true,
+	}
+	handler := makeSerializingForkRequestHandler(config)
+	handler(rr, req)
+
+	required := http.StatusOK
+
+	if status := rr.Code; status != required {
+		t.Errorf("handler returned wrong status code - got: %v, want: %v",
+			status, required)
+	}
+
+	read, _ := ioutil.ReadAll(rr.Body)
+	val := string(read)
+	if !strings.Contains(val, "Http_ContentLength=0") {
+		t.Errorf(config.FunctionProcess+" should print: Http_ContentLength=0, got: %s\n", val)
+	}
+	if !strings.Contains(val, "Http_Custom_Header=value") {
+		t.Errorf(config.FunctionProcess+" should print: Http_Custom_Header, got: %s\n", val)
+	}
+	seconds := rr.Header().Get("X-Duration-Seconds")
+	if len(seconds) == 0 {
+		t.Errorf(config.FunctionProcess + " should have given a duration as an X-Duration-Seconds header\n")
+	}
+}
+
+func TestSerializingForkHandler_HasCustomHeaderInFunction_WithBody_WithCGI(t *testing.T) {
+	rr := httptest.NewRecorder()
+
+	body := "test"
+	req, err := http.NewRequest(http.MethodPost, "/", bytes.NewBufferString(body))
+	req.Header.Add("custom-header", "value")
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	config := config.WatchdogConfig{
+		FunctionProcess:  "env",
+		InjectCGIHeaders: true,
+	}
+	handler := makeSerializingForkRequestHandler(config)
+	handler(rr, req)
+
+	required := http.StatusOK
+
+	if status := rr.Code; status != required {
+		t.Errorf("handler returned wrong status code - got: %v, want: %v",
+			status, required)
+	}
+
+	read, _ := ioutil.ReadAll(rr.Body)
+	val := string(read)
+	if !strings.Contains(val, fmt.Sprintf("Http_ContentLength=%d", len(body))) {
+		t.Errorf("'env' should printed: Http_ContentLength=0, got: %s\n", val)
+	}
+	if !strings.Contains(val, "Http_Custom_Header") {
+		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 TestSerializingForkHandler_HasHostHeaderWhenSet_WithCGI(t *testing.T) {
+	rr := httptest.NewRecorder()
+
+	body := "test"
+	req, err := http.NewRequest(http.MethodPost, "http://gateway/function", bytes.NewBufferString(body))
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	config := config.WatchdogConfig{
+		FunctionProcess:  "env",
+		InjectCGIHeaders: true,
+	}
+	handler := makeSerializingForkRequestHandler(config)
+	handler(rr, req)
+
+	required := http.StatusOK
+
+	if status := rr.Code; status != required {
+		t.Errorf("handler returned wrong status code - got: %v, want: %v",
+			status, required)
+	}
+
+	read, _ := ioutil.ReadAll(rr.Body)
+	val := string(read)
+	if !strings.Contains(val, fmt.Sprintf("Http_Host=%s", req.URL.Host)) {
+		t.Errorf("'env' should have printed: Http_Host=0, got: %s\n", val)
+	}
+}
+
+func TestSerializingForkHandler_HostHeader_Empty_WheNotSet_WithCGI(t *testing.T) {
+	rr := httptest.NewRecorder()
+
+	body := "test"
+	req, err := http.NewRequest(http.MethodPost, "/function", bytes.NewBufferString(body))
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	config := config.WatchdogConfig{
+		FunctionProcess:  "env",
+		InjectCGIHeaders: true,
+	}
+	handler := makeSerializingForkRequestHandler(config)
+	handler(rr, req)
+
+	required := http.StatusOK
+	if status := rr.Code; status != required {
+		t.Errorf("handler returned wrong status code - got: %v, want: %v",
+			status, required)
+	}
+
+	read, _ := ioutil.ReadAll(rr.Body)
+	val := string(read)
+	if strings.Contains(val, fmt.Sprintf("Http_Host=%s", req.URL.Host)) {
+		t.Errorf("Http_Host should not have been given, but was: %s\n", val)
+	}
+}
+
+func TestSerializingForkHandler_DoesntHaveCustomHeaderInFunction_WithoutCGI(t *testing.T) {
+	rr := httptest.NewRecorder()
+
+	body := ""
+	req, err := http.NewRequest(http.MethodPost, "/", bytes.NewBufferString(body))
+	req.Header.Add("custom-header", "value")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	config := config.WatchdogConfig{
+		FunctionProcess:  "env",
+		InjectCGIHeaders: false,
+	}
+	handler := makeSerializingForkRequestHandler(config)
+	handler(rr, req)
+
+	required := http.StatusOK
+	if status := rr.Code; status != required {
+		t.Errorf("handler returned wrong status code - got: %v, want: %v",
+			status, required)
+	}
+
+	read, _ := ioutil.ReadAll(rr.Body)
+	val := string(read)
+	if strings.Contains(val, "Http_Custom_Header") {
+		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 TestSerializingForkHandler_HasXDurationSecondsHeader(t *testing.T) {
+	rr := httptest.NewRecorder()
+
+	body := "hello"
+	req, err := http.NewRequest(http.MethodPost, "/", bytes.NewBufferString(body))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	config := config.WatchdogConfig{
+		FunctionProcess: "cat",
+	}
+	handler := makeSerializingForkRequestHandler(config)
+	handler(rr, req)
+
+	required := http.StatusOK
+	if status := rr.Code; status != required {
+		t.Errorf("handler returned wrong status code - got: %v, want: %v",
+			status, required)
+	}
+
+	seconds := rr.Header().Get("X-Duration-Seconds")
+	if len(seconds) == 0 {
+		t.Errorf("Exec of " + config.FunctionProcess + " should have given a duration as an X-Duration-Seconds header")
+	}
+}