From e097bb0dbdc1e7f2b1220192db0cf43bc21ad5c0 Mon Sep 17 00:00:00 2001
From: Alexander Baryshnikov <dev@baryshnikov.net>
Date: Thu, 17 Sep 2020 18:48:42 +0800
Subject: [PATCH] allow set bad request code after execution for bin mode

---
 server/mode_bin.go   | 4 ++++
 worker/request_io.go | 4 ++++
 worker/worker.go     | 2 +-
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/server/mode_bin.go b/server/mode_bin.go
index ce467c2..00e5ad7 100644
--- a/server/mode_bin.go
+++ b/server/mode_bin.go
@@ -77,6 +77,10 @@ func (bh *binHandler) ServeHTTP(writer http.ResponseWriter, request *http.Reques
 	internal.SetBinFlags(cmd)
 	err := cmd.Run()
 
+	if codeReset, ok := writer.(interface{ Status(status int) }); ok && err != nil {
+		codeReset.Status(http.StatusBadGateway)
+	}
+
 	if marker.dataSent {
 		return
 	}
diff --git a/worker/request_io.go b/worker/request_io.go
index ad15acb..54c05e4 100644
--- a/worker/request_io.go
+++ b/worker/request_io.go
@@ -35,3 +35,7 @@ func (mo *responseStream) WriteHeader(statusCode int) {
 	mo.statusSent = true
 	mo.meta.Code = statusCode
 }
+
+func (mo *responseStream) Status(status int) {
+	mo.meta.Code = status
+}
diff --git a/worker/worker.go b/worker/worker.go
index fcc794c..09a6058 100644
--- a/worker/worker.go
+++ b/worker/worker.go
@@ -268,8 +268,8 @@ func (mgr *Worker) call(ctx context.Context, requestID string, info *meta.Reques
 	for k, v := range info.Headers {
 		req.Header[k] = v
 	}
-	req.Header.Set("X-Correlation-Id", requestID)
 
+	req.Header.Set("X-Correlation-Id", requestID)
 	req.Header.Set("X-Attempt-Id", attemptID)
 	req.Header.Set("X-Attempt", strconv.Itoa(len(info.Attempts)+1))