encode: write status immediate for success response for CONNECT requests (#6738)

* encode: write status immediate for success response for CONNECT requests

* fix compile

* fix test

* fix lint

* treat first write and flush for encode response writer to CONNECT request as success if status is not set explicitly

Author: WeidiDeng

modules/caddyhttp/encode/encode.go

@@ -156,7 +156,7 @@ func (enc *Encode) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyh
 			if _, ok := enc.writerPools[encName]; !ok {
 				continue // encoding not offered
 			}
-			w = enc.openResponseWriter(encName, w)
+			w = enc.openResponseWriter(encName, w, r.Method == http.MethodConnect)
 			defer w.(*responseWriter).Close()
 
 			// to comply with RFC 9110 section 8.8.3(.3), we modify the Etag when encoding
@@ -201,21 +201,22 @@ func (enc *Encode) addEncoding(e Encoding) error {
 // openResponseWriter creates a new response writer that may (or may not)
 // encode the response with encodingName. The returned response writer MUST
 // be closed after the handler completes.
-func (enc *Encode) openResponseWriter(encodingName string, w http.ResponseWriter) *responseWriter {
+func (enc *Encode) openResponseWriter(encodingName string, w http.ResponseWriter, isConnect bool) *responseWriter {
 	var rw responseWriter
-	return enc.initResponseWriter(&rw, encodingName, w)
+	return enc.initResponseWriter(&rw, encodingName, w, isConnect)
 }
 
 // initResponseWriter initializes the responseWriter instance
 // allocated in openResponseWriter, enabling mid-stack inlining.
-func (enc *Encode) initResponseWriter(rw *responseWriter, encodingName string, wrappedRW http.ResponseWriter) *responseWriter {
+func (enc *Encode) initResponseWriter(rw *responseWriter, encodingName string, wrappedRW http.ResponseWriter, isConnect bool) *responseWriter {
 	if rww, ok := wrappedRW.(*caddyhttp.ResponseWriterWrapper); ok {
 		rw.ResponseWriter = rww
 	} else {
 		rw.ResponseWriter = &caddyhttp.ResponseWriterWrapper{ResponseWriter: wrappedRW}
 	}
 	rw.encodingName = encodingName
 	rw.config = enc
+	rw.isConnect = isConnect
 
 	return rw
 }
@@ -230,6 +231,7 @@ type responseWriter struct {
 	config       *Encode
 	statusCode   int
 	wroteHeader  bool
+	isConnect    bool
 }
 
 // WriteHeader stores the status to write when the time comes
@@ -245,6 +247,14 @@ func (rw *responseWriter) WriteHeader(status int) {
 		rw.Header().Add("Vary", "Accept-Encoding")
 	}
 
+	// write status immediately if status is 2xx and the request is CONNECT
+	// since it means the response is successful.
+	// see: https://github.com/caddyserver/caddy/issues/6733#issuecomment-2525058845
+	if rw.isConnect && 200 <= status && status <= 299 {
+		rw.ResponseWriter.WriteHeader(status)
+		rw.wroteHeader = true
+	}
+
 	// write status immediately when status code is informational
 	// see: https://caddy.community/t/disappear-103-early-hints-response-with-encode-enable-caddy-v2-7-6/23081/5
 	if 100 <= status && status <= 199 {
@@ -260,6 +270,12 @@ func (enc *Encode) Match(rw *responseWriter) bool {
 // FlushError is an alternative Flush returning an error. It delays the actual Flush of the underlying
 // ResponseWriterWrapper until headers were written.
 func (rw *responseWriter) FlushError() error {
+	// WriteHeader wasn't called and is a CONNECT request, treat it as a success.
+	// otherwise, wait until header is written.
+	if rw.isConnect && !rw.wroteHeader && rw.statusCode == 0 {
+		rw.WriteHeader(http.StatusOK)
+	}
+
 	if !rw.wroteHeader {
 		// flushing the underlying ResponseWriter will write header and status code,
 		// but we need to delay that until we can determine if we must encode and
@@ -288,6 +304,12 @@ func (rw *responseWriter) Write(p []byte) (int, error) {
 		return 0, nil
 	}
 
+	// WriteHeader wasn't called and is a CONNECT request, treat it as a success.
+	// otherwise, determine if the response should be compressed.
+	if rw.isConnect && !rw.wroteHeader && rw.statusCode == 0 {
+		rw.WriteHeader(http.StatusOK)
+	}
+
 	// sniff content-type and determine content-length
 	if !rw.wroteHeader && rw.config.MinLength > 0 {
 		var gtMinLength bool

modules/caddyhttp/encode/encode_test.go

@@ -9,7 +9,7 @@ import (
 func BenchmarkOpenResponseWriter(b *testing.B) {
 	enc := new(Encode)
 	for n := 0; n < b.N; n++ {
-		enc.openResponseWriter("test", nil)
+		enc.openResponseWriter("test", nil, false)
 	}
 }