From cb38271dc6155088962eb5c45d3c811238b8254c Mon Sep 17 00:00:00 2001
From: keks <keks@cryptoscope.co>
Date: Mon, 1 Oct 2018 16:45:26 +0200
Subject: [PATCH] implement more fine-grained setting of "Connection: close"
 header

---
 http/body.go            | 19 +++++++++++++++++++
 http/handler.go         | 26 ++++++++++++++++++--------
 http/responseemitter.go | 36 ++++++++++++++++++++++++++++++++++--
 3 files changed, 71 insertions(+), 10 deletions(-)
 create mode 100644 http/body.go

diff --git a/http/body.go b/http/body.go
new file mode 100644
index 00000000..01d2ab0c
--- /dev/null
+++ b/http/body.go
@@ -0,0 +1,19 @@
+package http
+
+import "io"
+
+// bodyWrapper wraps an io.Reader and calls onError whenever the Read function returns an error.
+// This was designed for wrapping the request body, so we can know whether it was closed.
+type bodyWrapper struct {
+	io.ReadCloser
+	onError func(err error)
+}
+
+func (bw bodyWrapper) Read(data []byte) (int, error) {
+	n, err := bw.ReadCloser.Read(data)
+	if err != nil {
+		bw.onError(err)
+	}
+
+	return n, err
+}
diff --git a/http/handler.go b/http/handler.go
index b0542290..3df03f4d 100644
--- a/http/handler.go
+++ b/http/handler.go
@@ -109,6 +109,23 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	// If we have a request body, make sure the preamble
+	// knows that it should close the body if it wants to
+	// write before completing reading.
+	// FIXME: https://github.com/ipfs/go-ipfs/issues/5168
+	// FIXME: https://github.com/golang/go/issues/15527
+	var bodyErrChan chan error
+	if r.Body != http.NoBody {
+		bodyErrChan = make(chan error)
+		bw := bodyWrapper{
+			ReadCloser: r.Body,
+			onError: func(err error) {
+				bodyErrChan <- err
+			},
+		}
+		r.Body = bw
+	}
+
 	req, err := parseRequest(ctx, r, h.root)
 	if err != nil {
 		if err == ErrNotFound {
@@ -157,14 +174,7 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	// If we have a request body, close the connection when we're done.
-	// FIXME: https://github.com/ipfs/go-ipfs/issues/5168
-	// FIXME: https://github.com/golang/go/issues/15527
-	if r.Body != http.NoBody {
-		w.Header().Set("Connection", "close")
-	}
-
-	re := NewResponseEmitter(w, r.Method, req)
+	re := NewResponseEmitter(w, r.Method, req, WithRequestBodyErrorChan(bodyErrChan))
 	h.root.Call(req, re, h.env)
 }
 
diff --git a/http/responseemitter.go b/http/responseemitter.go
index 4f754c07..192817d5 100644
--- a/http/responseemitter.go
+++ b/http/responseemitter.go
@@ -27,11 +27,10 @@ var (
 )
 
 // NewResponeEmitter returns a new ResponseEmitter.
-func NewResponseEmitter(w http.ResponseWriter, method string, req *cmds.Request) ResponseEmitter {
+func NewResponseEmitter(w http.ResponseWriter, method string, req *cmds.Request, opts ...ResponseEmitterOption) ResponseEmitter {
 	encType := cmds.GetEncoding(req)
 
 	var enc cmds.Encoder
-
 	if _, ok := cmds.Encoders[encType]; ok {
 		enc = cmds.Encoders[encType](req)(w)
 	}
@@ -43,9 +42,24 @@ func NewResponseEmitter(w http.ResponseWriter, method string, req *cmds.Request)
 		method:  method,
 		req:     req,
 	}
+
+	for _, opt := range opts {
+		opt(re)
+	}
+
 	return re
 }
 
+type ResponseEmitterOption func(*responseEmitter)
+
+func WithRequestBodyErrorChan(ch <-chan error) ResponseEmitterOption {
+	return func(re *responseEmitter) {
+		if ch != nil {
+			re.bodyErrChan = ch
+		}
+	}
+}
+
 type ResponseEmitter interface {
 	cmds.ResponseEmitter
 	http.Flusher
@@ -61,6 +75,8 @@ type responseEmitter struct {
 	length uint64
 	err    *cmdkit.Error
 
+	bodyErrChan <-chan error
+
 	streaming bool
 	once      sync.Once
 	method    string
@@ -179,6 +195,22 @@ func (re *responseEmitter) preamble(value interface{}) {
 		}
 	}
 
+	// If we have a request body, make sure we close the body
+	// if we want to write before completing reading.
+	// FIXME: https://github.com/ipfs/go-ipfs/issues/5168
+	// FIXME: https://github.com/golang/go/issues/15527
+	if re.bodyErrChan != nil {
+		select {
+		case <-re.bodyErrChan:
+			// all good, we received an error, so the body is read completely.
+			// we handle the error where it occurs, here we just want to know that we're done
+		default:
+			// set connection close header, because we want to write
+			// even though the body is not yet read completely.
+			h.Set("Connection", "close")
+		}
+	}
+
 	var mime string
 
 	switch v := value.(type) {