reverseproxy: simplify WebSocket header handling (#8365)

.golangci.yml

@@ -52,7 +52,6 @@ linters:
       # Match net/http's canonical MIME map keys to avoid re-canonicalisation.
       exclusions:
         - Etag
-        - Sec-Websocket-Key
         - Te
         - Www-Authenticate
     staticcheck:

modules/caddyhttp/reverseproxy/extended_connect_test.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"strconv"
 	"strings"
 	"sync"
 	"testing"
@@ -265,6 +266,9 @@ func TestServeHTTPRewritesExtendedConnectWebsocketSecWebSocketKeyMatrix(t *testi
 					var keyValues []string
 					for k, v := range captured.headers {
 						if strings.EqualFold(k, "Sec-WebSocket-Key") {
+							if k != "Sec-WebSocket-Key" {
+								t.Fatalf("upstream WebSocket key spelling = %q, want Sec-WebSocket-Key", k)
+							}
 							keyValues = append(keyValues, v...)
 						}
 					}
@@ -302,3 +306,60 @@ func TestServeHTTPRewritesExtendedConnectWebsocketSecWebSocketKeyMatrix(t *testi
 		}
 	}
 }
+
+func TestUpgradeResponseWebsocketHeaders(t *testing.T) {
+	for _, proto := range []int{1, 2, 3} {
+		t.Run("h"+strconv.Itoa(proto), func(t *testing.T) {
+			h := &Handler{logger: zap.NewNop()}
+			req := prepareTestRequest(httptest.NewRequest(http.MethodGet, "http://example.test/upgrade", nil))
+			req.ProtoMajor = proto
+			req.Header.Set("Connection", "Upgrade")
+			req.Header.Set("Upgrade", "websocket")
+			wantStatus := http.StatusSwitchingProtocols
+			if proto > 1 {
+				caddyhttp.SetVar(req.Context(), "extended_connect_websocket_body", http.NoBody)
+				wantStatus = http.StatusOK
+			}
+			res := &http.Response{
+				StatusCode: http.StatusSwitchingProtocols,
+				Header:     req.Header.Clone(),
+				Body:       fakeRWC{strings.NewReader("")},
+			}
+			res.Header.Set("Sec-WebSocket-Accept", "backend-accept")
+			res.Header.Set("Sec-WebSocket-Protocol", "chat")
+			rw := upgradeHeaderRecorder{httptest.NewRecorder()}
+			h.handleUpgradeResponse(h.logger, rw, req, res, "test-upstream")
+			response := rw.Result()
+			defer response.Body.Close()
+			if response.StatusCode != wantStatus {
+				t.Fatalf("status = %d, want %d", response.StatusCode, wantStatus)
+			}
+			var accept, protocol string
+			for key, values := range response.Header {
+				if strings.EqualFold(key, "Sec-WebSocket-Accept") {
+					if proto > 1 {
+						t.Fatalf("extended CONNECT response contains %s", key)
+					}
+					if key != "Sec-WebSocket-Accept" {
+						t.Fatalf("WebSocket accept spelling = %q, want Sec-WebSocket-Accept", key)
+					}
+					accept = strings.Join(values, ",")
+				}
+				if strings.EqualFold(key, "Sec-WebSocket-Protocol") {
+					protocol = strings.Join(values, ",")
+				}
+			}
+			if proto == 1 && accept != "backend-accept" {
+				t.Fatalf("accept = %q, want backend-accept", accept)
+			}
+			if protocol != "chat" {
+				t.Fatalf("subprotocol = %q, want chat", protocol)
+			}
+		})
+	}
+}
+
+type upgradeHeaderRecorder struct{ *httptest.ResponseRecorder }
+
+// Stop after the headers are committed, before the tunnel starts.
+func (upgradeHeaderRecorder) FlushError() error { return io.ErrClosedPipe }

modules/caddyhttp/reverseproxy/reverseproxy.go

@@ -655,8 +655,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyht
 		if randErr != nil {
 			return randErr
 		}
-		clonedReq.Header.Del("Sec-WebSocket-Key")                                                //nolint:canonicalheader // RFC 6455 casing, see normalizeWebsocketHeaders
-		clonedReq.Header["Sec-WebSocket-Key"] = []string{base64.StdEncoding.EncodeToString(key)} //nolint:canonicalheader // RFC 6455 casing, see normalizeWebsocketHeaders
+		clonedReq.Header.Set("Sec-WebSocket-Key", base64.StdEncoding.EncodeToString(key))
 	}
 
 	// we will need the original headers and Host value if

modules/caddyhttp/reverseproxy/streaming.go

@@ -103,7 +103,6 @@ func (h *Handler) handleUpgradeResponse(logger *zap.Logger, rw http.ResponseWrit
 	// write header first, response headers should not be counted in size
 	// like the rest of handler chain.
 	copyHeader(rw.Header(), res.Header)
-	normalizeWebsocketHeaders(rw.Header())
 
 	// Capture all h fields needed by the tunnel now, so that the Handler (h)
 	// is not referenced after this function returns (for HTTP/1.1 hijacked
@@ -127,7 +126,7 @@ func (h *Handler) handleUpgradeResponse(logger *zap.Logger, rw http.ResponseWrit
 		req.Body = body
 		rw.Header().Del("Upgrade")
 		rw.Header().Del("Connection")
-		delete(rw.Header(), "Sec-WebSocket-Accept") //nolint:canonicalheader // set by normalizeWebsocketHeaders above
+		rw.Header().Del("Sec-WebSocket-Accept")
 		rw.WriteHeader(http.StatusOK)
 
 		if c := logger.Check(zap.DebugLevel, "upgrading connection"); c != nil {
@@ -146,6 +145,7 @@ func (h *Handler) handleUpgradeResponse(logger *zap.Logger, rw http.ResponseWrit
 		// bufio is not needed, use minimal buffer
 		brw = bufio.NewReadWriter(bufio.NewReaderSize(conn, 1), bufio.NewWriterSize(conn, 1))
 	} else {
+		normalizeWebsocketHeaders(rw.Header())
 		if h.StreamDetached {
 			detached = caddyhttp.DetachResponseWriterAfterHijack(rw, req, true)
 			if !detached {