Fix afterburn for writing to response

Signed-off-by: Alex Ellis <[email protected]>

Author: alexellis

config/config.go

@@ -35,7 +35,7 @@ func New(env []string) (WatchdogConfig, error) {
 		FunctionProcess:  os.Getenv("fprocess"),
 		InjectCGIHeaders: true,
 		HardTimeout:      5 * time.Second,
-		OperationalMode:  ModeAfterBurn,
+		OperationalMode:  ModeStreaming,
 	}
 
 	envMap := mapEnv(env)

config/config_test.go

@@ -27,11 +27,12 @@ func Test_OperationalMode_AfterBurn(t *testing.T) {
 		"mode=afterburn",
 	}
 
-	defaults, err := New(env)
+	actual, err := New(env)
 	if err != nil {
 		t.Errorf("Expected no errors")
 	}
-	if defaults.OperationalMode != ModeStreaming {
-		t.Errorf("Want %s. got: %s", WatchdogMode(ModeStreaming), WatchdogMode(defaults.OperationalMode))
+
+	if actual.OperationalMode != ModeAfterBurn {
+		t.Errorf("Want %s. got: %s", WatchdogMode(ModeAfterBurn), WatchdogMode(actual.OperationalMode))
 	}
 }

functions/afterburn_runner.go

@@ -3,6 +3,7 @@ package functions
 import (
 	"bufio"
 	"io"
+	"io/ioutil"
 	"log"
 	"net/http"
 	"os/exec"
@@ -61,26 +62,43 @@ func (f *AfterBurnFunctionRunner) Start() error {
 
 // Run a function with a long-running process with a HTTP protocol for communication
 func (f *AfterBurnFunctionRunner) Run(req FunctionRequest, contentLength int64, r *http.Request, w http.ResponseWriter) error {
-	buffReader := bufio.NewReader(f.StdoutPipe)
 
+	// Submit body to function via stdin
 	writeErr := r.Write(f.StdinPipe)
+
 	if writeErr != nil {
 		return writeErr
 	}
 
-	processRes, err := http.ReadResponse(buffReader, r)
-	if err != nil {
-		return err
+	var processRes *http.Response
+
+	// Read response back from stdout
+	buffReader := bufio.NewReader(f.StdoutPipe)
+	var err1 error
+	processRes, err1 = http.ReadResponse(buffReader, r)
+	if err1 != nil {
+		return err1
 	}
 
-	if processRes.Body != nil {
-		defer processRes.Body.Close()
+	for h := range processRes.Header {
+		w.Header().Set(h, processRes.Header.Get(h))
 	}
 
 	w.WriteHeader(processRes.StatusCode)
-	processRes.Write(w)
+	if processRes.Body != nil {
+		defer processRes.Body.Close()
+		bodyBytes, bodyErr := ioutil.ReadAll(processRes.Body)
+		if bodyErr != nil {
+			log.Println("read body err", bodyErr)
+		}
 
-	log.Printf("%s %s - %s - ContentLength: %d\n", r.Method, r.RequestURI, processRes.Status, processRes.ContentLength)
+		// processRes.Write(w)
+		w.Write(bodyBytes)
+	}
+
+	if processRes != nil {
+		log.Printf("%s %s - %s - ContentLength: %d\n", r.Method, r.RequestURI, processRes.Status, processRes.ContentLength)
+	}
 
 	return nil
 }

main.go

@@ -69,13 +69,16 @@ func makeAfterBurnRequestHandler(watchdogConfig config.WatchdogConfig) func(http
 			InputReader:  r.Body,
 			OutputWriter: w,
 		}
+
 		functionInvoker.Mutex.Lock()
 
 		err := functionInvoker.Run(req, r.ContentLength, r, w)
+
 		if err != nil {
 			w.WriteHeader(500)
 			w.Write([]byte(err.Error()))
 		}
+
 		functionInvoker.Mutex.Unlock()
 	}
 }