net/http/httputil: ReverseProxy request cancellation

If an inbound connection is closed, cancel the outbound http request.

This is particularly useful if the outbound request may consume resources
unnecessarily until it is cancelled.

Fixes #8406

Change-Id: I738c4489186ce342f7e21d0ea3f529722c5b443a
Signed-off-by: Peter Waller <p@pwaller.net>
Reviewed-on: https://go-review.googlesource.com/2320
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

1 files changed, 46 insertions, 0 deletions

diff --git a/src/net/http/httputil/reverseproxy.go b/src/net/http/httputil/reverseproxy.go
index ab46370180..5a0c1edfe1 100644
--- a/src/net/http/httputil/reverseproxy.go
+++ b/src/net/http/httputil/reverseproxy.go
@@ -100,6 +100,24 @@ var hopHeaders = []string{
 	"Upgrade",
 }
 
+type requestCanceler interface {
+	CancelRequest(*http.Request)
+}
+
+type runOnFirstRead struct {
+	io.Reader
+
+	fn func() // Run before first Read, then set to nil
+}
+
+func (c *runOnFirstRead) Read(bs []byte) (int, error) {
+	if c.fn != nil {
+		c.fn()
+		c.fn = nil
+	}
+	return c.Reader.Read(bs)
+}
+
 func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	transport := p.Transport
 	if transport == nil {
@@ -109,6 +127,34 @@ func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	outreq := new(http.Request)
 	*outreq = *req // includes shallow copies of maps, but okay
 
+	if closeNotifier, ok := rw.(http.CloseNotifier); ok {
+		if requestCanceler, ok := transport.(requestCanceler); ok {
+			reqDone := make(chan struct{})
+			defer close(reqDone)
+
+			clientGone := closeNotifier.CloseNotify()
+
+			outreq.Body = struct {
+				io.Reader
+				io.Closer
+			}{
+				Reader: &runOnFirstRead{
+					Reader: outreq.Body,
+					fn: func() {
+						go func() {
+							select {
+							case <-clientGone:
+								requestCanceler.CancelRequest(outreq)
+							case <-reqDone:
+							}
+						}()
+					},
+				},
+				Closer: outreq.Body,
+			}
+		}
+	}
+
 	p.Director(outreq)
 	outreq.Proto = "HTTP/1.1"
 	outreq.ProtoMajor = 1