Close target connection when request is canceled

ackport changes added to the net/http/httputil ReverseProxy for go1.15 .
See release notes here https://golang.org/doc/go1.15#net/http/httputil

Author: gernest

gateway/reverse_proxy.go

@@ -42,6 +42,7 @@ import (
 	"github.com/opentracing/opentracing-go/ext"
 	cache "github.com/pmylund/go-cache"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/net/http/httpguts"
 	"golang.org/x/net/http2"
 )
 
@@ -1108,6 +1109,13 @@ func (p *ReverseProxy) copyBuffer(dst io.Writer, src io.Reader) (int64, error) {
 	}
 }
 
+func upgradeType(h http.Header) string {
+	if !httpguts.HeaderValuesContainsToken(h["Connection"], "Upgrade") {
+		return ""
+	}
+	return strings.ToLower(h.Get("Upgrade"))
+}
+
 func (p *ReverseProxy) handleUpgradeResponse(rw http.ResponseWriter, req *http.Request, res *http.Response) error {
 	copyHeader(res.Header, rw.Header())
 
@@ -1119,7 +1127,18 @@ func (p *ReverseProxy) handleUpgradeResponse(rw http.ResponseWriter, req *http.R
 	if !ok {
 		return fmt.Errorf("internal error: 101 switching protocols response with non-writable body")
 	}
-	defer backConn.Close()
+	backConnCloseCh := make(chan bool)
+	go func() {
+		// Ensure that the cancelation of a request closes the backend.
+		// See issue https://golang.org/issue/35559.
+		select {
+		case <-req.Context().Done():
+		case <-backConnCloseCh:
+		}
+		backConn.Close()
+	}()
+
+	defer close(backConnCloseCh)
 	conn, brw, err := hj.Hijack()
 	if err != nil {
 		return fmt.Errorf("Hijack failed on protocol switch: %v", err)

gateway/reverse_proxy_test.go

@@ -1,9 +1,12 @@
 package gateway
 
 import (
+	"bufio"
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
@@ -781,3 +784,129 @@ func TestEnsureTransport(t *testing.T) {
 		})
 	}
 }
+
+func TestReverseProxyWebSocketCancelation(t *testing.T) {
+	c := config.Global()
+	c.HttpServerOptions.EnableWebSockets = true
+	config.SetGlobal(c)
+	n := 5
+	triggerCancelCh := make(chan bool, n)
+	nthResponse := func(i int) string {
+		return fmt.Sprintf("backend response #%d\n", i)
+	}
+	terminalMsg := "final message"
+
+	cst := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if g, ws := upgradeType(r.Header), "websocket"; g != ws {
+			t.Errorf("Unexpected upgrade type %q, want %q", g, ws)
+			http.Error(w, "Unexpected request", 400)
+			return
+		}
+		conn, bufrw, err := w.(http.Hijacker).Hijack()
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		defer conn.Close()
+
+		upgradeMsg := "HTTP/1.1 101 Switching Protocols\r\nConnection: upgrade\r\nUpgrade: WebSocket\r\n\r\n"
+		if _, err := io.WriteString(conn, upgradeMsg); err != nil {
+			t.Error(err)
+			return
+		}
+		if _, _, err := bufrw.ReadLine(); err != nil {
+			t.Errorf("Failed to read line from client: %v", err)
+			return
+		}
+
+		for i := 0; i < n; i++ {
+			if _, err := bufrw.WriteString(nthResponse(i)); err != nil {
+				select {
+				case <-triggerCancelCh:
+				default:
+					t.Errorf("Writing response #%d failed: %v", i, err)
+				}
+				return
+			}
+			bufrw.Flush()
+			time.Sleep(time.Second)
+		}
+		if _, err := bufrw.WriteString(terminalMsg); err != nil {
+			select {
+			case <-triggerCancelCh:
+			default:
+				t.Errorf("Failed to write terminal message: %v", err)
+			}
+		}
+		bufrw.Flush()
+	}))
+	defer cst.Close()
+
+	backendURL, _ := url.Parse(cst.URL)
+	spec := &APISpec{APIDefinition: &apidef.APIDefinition{}}
+	rproxy := TykNewSingleHostReverseProxy(backendURL, spec, nil)
+
+	handler := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("X-Header", "X-Value")
+		ctx, cancel := context.WithCancel(req.Context())
+		go func() {
+			<-triggerCancelCh
+			cancel()
+		}()
+		rproxy.ServeHTTP(rw, req.WithContext(ctx))
+	})
+
+	frontendProxy := httptest.NewServer(handler)
+	defer frontendProxy.Close()
+
+	req, _ := http.NewRequest("GET", frontendProxy.URL, nil)
+	req.Header.Set("Connection", "Upgrade")
+	req.Header.Set("Upgrade", "websocket")
+
+	res, err := frontendProxy.Client().Do(req)
+	if err != nil {
+		t.Fatalf("Dialing to frontend proxy: %v", err)
+	}
+	defer res.Body.Close()
+	if g, w := res.StatusCode, 101; g != w {
+		t.Fatalf("Switching protocols failed, got: %d, want: %d", g, w)
+	}
+
+	if g, w := res.Header.Get("X-Header"), "X-Value"; g != w {
+		t.Errorf("X-Header mismatch\n\tgot:  %q\n\twant: %q", g, w)
+	}
+
+	if g, w := upgradeType(res.Header), "websocket"; g != w {
+		t.Fatalf("Upgrade header mismatch\n\tgot:  %q\n\twant: %q", g, w)
+	}
+
+	rwc, ok := res.Body.(io.ReadWriteCloser)
+	if !ok {
+		t.Fatalf("Response body type mismatch, got %T, want io.ReadWriteCloser", res.Body)
+	}
+
+	if _, err := io.WriteString(rwc, "Hello\n"); err != nil {
+		t.Fatalf("Failed to write first message: %v", err)
+	}
+
+	// Read loop.
+
+	br := bufio.NewReader(rwc)
+	for {
+		line, err := br.ReadString('\n')
+		switch {
+		case line == terminalMsg: // this case before "err == io.EOF"
+			t.Fatalf("The websocket request was not canceled, unfortunately!")
+
+		case err == io.EOF:
+			return
+
+		case err != nil:
+			t.Fatalf("Unexpected error: %v", err)
+
+		case line == nthResponse(0): // We've gotten the first response back
+			// Let's trigger a cancel.
+			close(triggerCancelCh)
+		}
+	}
+}