Fix cloudflared parity gaps

Author: nekohasekai

protocol/cloudflare/connection_drain_test.go

@@ -10,8 +10,9 @@ import (
 	"testing"
 	"time"
 
-	"github.com/google/uuid"
 	"github.com/sagernet/quic-go"
+
+	"github.com/google/uuid"
 )
 
 type stubNetConn struct {
@@ -43,6 +44,7 @@ func (c *stubQUICConn) OpenStream() (*quic.Stream, error) { return nil, errors.N
 func (c *stubQUICConn) AcceptStream(context.Context) (*quic.Stream, error) {
 	return nil, errors.New("unused")
 }
+
 func (c *stubQUICConn) ReceiveDatagram(context.Context) ([]byte, error) {
 	return nil, errors.New("unused")
 }

protocol/cloudflare/connection_http2.go

@@ -9,6 +9,7 @@ import (
 	"math"
 	"net"
 	"net/http"
+	"runtime/debug"
 	"strconv"
 	"strings"
 	"sync"
@@ -27,8 +28,17 @@ const (
 	h2EdgeSNI                        = "h2.cftunnel.com"
 	h2ResponseMetaCloudflared        = `{"src":"cloudflared"}`
 	h2ResponseMetaCloudflaredLimited = `{"src":"cloudflared","flow_rate_limited":true}`
+	contentTypeHeader                = "content-type"
+	contentLengthHeader              = "content-length"
+	transferEncodingHeader           = "transfer-encoding"
+	chunkTransferEncoding            = "chunked"
+	sseContentType                   = "text/event-stream"
+	grpcContentType                  = "application/grpc"
+	ndjsonContentType                = "application/x-ndjson"
 )
 
+var flushableContentTypes = []string{sseContentType, grpcContentType, ndjsonContentType}
+
 // HTTP2Connection manages a single HTTP/2 connection to the Cloudflare edge.
 // Uses role reversal: we dial the edge as a TLS client but serve HTTP/2 as server.
 type HTTP2Connection struct {
@@ -191,7 +201,7 @@ func (c *HTTP2Connection) handleControlStream(ctx context.Context, r *http.Reque
 		return
 	}
 	c.registrationResult = result
-	c.inbound.notifyConnected(c.connIndex)
+	c.inbound.notifyConnected(c.connIndex, "http2")
 
 	c.logger.Info("connected to ", result.Location,
 		" (connection ", result.ConnectionID, ")")
@@ -246,14 +256,18 @@ func (c *HTTP2Connection) handleH2DataStream(ctx context.Context, r *http.Reques
 		}
 	}
 
+	flushState := &http2FlushState{shouldFlush: connectionType != ConnectionTypeHTTP}
 	stream := &http2DataStream{
 		reader:  r.Body,
 		writer:  w,
 		flusher: flusher,
+		state:   flushState,
+		logger:  c.logger,
 	}
 	respWriter := &http2ResponseWriter{
-		writer:  w,
-		flusher: flusher,
+		writer:     w,
+		flusher:    flusher,
+		flushState: flushState,
 	}
 
 	c.inbound.dispatchRequest(ctx, stream, respWriter, request)
@@ -386,15 +400,26 @@ type http2DataStream struct {
 	reader  io.ReadCloser
 	writer  http.ResponseWriter
 	flusher http.Flusher
+	state   *http2FlushState
+	logger  log.ContextLogger
 }
 
 func (s *http2DataStream) Read(p []byte) (int, error) {
 	return s.reader.Read(p)
 }
 
-func (s *http2DataStream) Write(p []byte) (int, error) {
-	n, err := s.writer.Write(p)
-	if err == nil {
+func (s *http2DataStream) Write(p []byte) (n int, err error) {
+	defer func() {
+		if recovered := recover(); recovered != nil {
+			if s.logger != nil {
+				s.logger.Debug("recovered from HTTP/2 data stream panic: ", recovered, "\n", string(debug.Stack()))
+			}
+			n = 0
+			err = io.ErrClosedPipe
+		}
+	}()
+	n, err = s.writer.Write(p)
+	if err == nil && s.state != nil && s.state.shouldFlush {
 		s.flusher.Flush()
 	}
 	return n, err
@@ -409,6 +434,7 @@ type http2ResponseWriter struct {
 	writer      http.ResponseWriter
 	flusher     http.Flusher
 	headersSent bool
+	flushState  *http2FlushState
 }
 
 func (w *http2ResponseWriter) AddTrailer(name, value string) {
@@ -462,12 +488,37 @@ func (w *http2ResponseWriter) WriteResponse(responseError error, metadata []Meta
 
 	w.writer.Header().Set(h2HeaderResponseUser, SerializeHeaders(userHeaders))
 	w.writer.Header().Set(h2HeaderResponseMeta, h2ResponseMetaOrigin)
+	if w.flushState != nil && shouldFlushHTTPHeaders(userHeaders) {
+		w.flushState.shouldFlush = true
+	}
 
 	if statusCode == http.StatusSwitchingProtocols {
 		statusCode = http.StatusOK
 	}
 
 	w.writer.WriteHeader(statusCode)
-	w.flusher.Flush()
+	if w.flushState != nil && w.flushState.shouldFlush {
+		w.flusher.Flush()
+	}
 	return nil
 }
+
+type http2FlushState struct {
+	shouldFlush bool
+}
+
+func shouldFlushHTTPHeaders(headers http.Header) bool {
+	if headers.Get(contentLengthHeader) == "" {
+		return true
+	}
+	if transferEncoding := strings.ToLower(headers.Get(transferEncodingHeader)); transferEncoding != "" && strings.Contains(transferEncoding, chunkTransferEncoding) {
+		return true
+	}
+	contentType := strings.ToLower(headers.Get(contentTypeHeader))
+	for _, flushable := range flushableContentTypes {
+		if strings.HasPrefix(contentType, flushable) {
+			return true
+		}
+	}
+	return false
+}

protocol/cloudflare/connection_http2_behavior_test.go

@@ -0,0 +1,167 @@
+//go:build with_cloudflared
+
+package cloudflare
+
+import (
+	"io"
+	"net/http"
+	"testing"
+
+	"github.com/sagernet/sing-box/log"
+)
+
+type captureHTTP2Writer struct {
+	header     http.Header
+	flushCount int
+	statusCode int
+	body       []byte
+	panicWrite bool
+}
+
+func (w *captureHTTP2Writer) Header() http.Header {
+	if w.header == nil {
+		w.header = make(http.Header)
+	}
+	return w.header
+}
+
+func (w *captureHTTP2Writer) WriteHeader(statusCode int) {
+	w.statusCode = statusCode
+}
+
+func (w *captureHTTP2Writer) Write(p []byte) (int, error) {
+	if w.panicWrite {
+		panic("write after close")
+	}
+	w.body = append(w.body, p...)
+	return len(p), nil
+}
+
+func (w *captureHTTP2Writer) Flush() {
+	w.flushCount++
+}
+
+func TestHTTP2NonStreamingResponseDoesNotFlush(t *testing.T) {
+	writer := &captureHTTP2Writer{}
+	flushState := &http2FlushState{}
+	respWriter := &http2ResponseWriter{
+		writer:     writer,
+		flusher:    writer,
+		flushState: flushState,
+	}
+
+	err := respWriter.WriteResponse(nil, encodeResponseHeaders(http.StatusOK, http.Header{
+		"Content-Type":   []string{"application/json"},
+		"Content-Length": []string{"2"},
+	}))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if writer.flushCount != 0 {
+		t.Fatalf("expected no header flush for non-streaming response, got %d", writer.flushCount)
+	}
+
+	stream := &http2DataStream{
+		writer:  writer,
+		flusher: writer,
+		state:   flushState,
+		logger:  log.NewNOPFactory().NewLogger("test"),
+	}
+	if _, err := stream.Write([]byte("ok")); err != nil {
+		t.Fatal(err)
+	}
+	if writer.flushCount != 0 {
+		t.Fatalf("expected no body flush for non-streaming response, got %d", writer.flushCount)
+	}
+}
+
+func TestHTTP2StreamingResponsesFlush(t *testing.T) {
+	testCases := []struct {
+		name   string
+		header http.Header
+	}{
+		{
+			name: "sse",
+			header: http.Header{
+				"Content-Type":   []string{"text/event-stream"},
+				"Content-Length": []string{"1"},
+			},
+		},
+		{
+			name: "grpc",
+			header: http.Header{
+				"Content-Type":   []string{"application/grpc"},
+				"Content-Length": []string{"1"},
+			},
+		},
+		{
+			name: "ndjson",
+			header: http.Header{
+				"Content-Type":   []string{"application/x-ndjson"},
+				"Content-Length": []string{"1"},
+			},
+		},
+		{
+			name: "chunked",
+			header: http.Header{
+				"Content-Type":      []string{"application/json"},
+				"Content-Length":    []string{"-1"},
+				"Transfer-Encoding": []string{"chunked"},
+			},
+		},
+		{
+			name: "no-content-length",
+			header: http.Header{
+				"Content-Type": []string{"application/json"},
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			writer := &captureHTTP2Writer{}
+			flushState := &http2FlushState{}
+			respWriter := &http2ResponseWriter{
+				writer:     writer,
+				flusher:    writer,
+				flushState: flushState,
+			}
+
+			err := respWriter.WriteResponse(nil, encodeResponseHeaders(http.StatusOK, testCase.header))
+			if err != nil {
+				t.Fatal(err)
+			}
+			if writer.flushCount == 0 {
+				t.Fatal("expected header flush for streaming response")
+			}
+
+			stream := &http2DataStream{
+				writer:  writer,
+				flusher: writer,
+				state:   flushState,
+				logger:  log.NewNOPFactory().NewLogger("test"),
+			}
+			if _, err := stream.Write([]byte("chunk")); err != nil {
+				t.Fatal(err)
+			}
+			if writer.flushCount < 2 {
+				t.Fatalf("expected body flush for streaming response, got %d flushes", writer.flushCount)
+			}
+		})
+	}
+}
+
+func TestHTTP2DataStreamWriteRecoversPanic(t *testing.T) {
+	writer := &captureHTTP2Writer{panicWrite: true}
+	stream := &http2DataStream{
+		writer:  writer,
+		flusher: writer,
+		state:   &http2FlushState{shouldFlush: true},
+		logger:  log.NewNOPFactory().NewLogger("test"),
+	}
+
+	_, err := stream.Write([]byte("panic"))
+	if err != io.ErrClosedPipe {
+		t.Fatalf("expected io.ErrClosedPipe, got %v", err)
+	}
+}

protocol/cloudflare/connection_quic.go

@@ -60,6 +60,15 @@ type QUICConnection struct {
 	closeOnce         sync.Once
 }
 
+type quicStreamHandle interface {
+	io.Reader
+	io.Writer
+	io.Closer
+	CancelRead(code quic.StreamErrorCode)
+	CancelWrite(code quic.StreamErrorCode)
+	SetWriteDeadline(t time.Time) error
+}
+
 type quicConnection interface {
 	OpenStream() (*quic.Stream, error)
 	AcceptStream(ctx context.Context) (*quic.Stream, error)
@@ -80,7 +89,6 @@ func (c *closeableQUICConn) CloseWithError(code quic.ApplicationErrorCode, reaso
 	return err
 }
 
-
 // NewQUICConnection dials the edge and establishes a QUIC connection.
 func NewQUICConnection(
 	ctx context.Context,
@@ -240,13 +248,14 @@ func (q *QUICConnection) acceptStreams(ctx context.Context, handler StreamHandle
 	}
 }
 
-func (q *QUICConnection) handleStream(ctx context.Context, stream *quic.Stream, handler StreamHandler) {
+func (q *QUICConnection) handleStream(ctx context.Context, stream quicStreamHandle, handler StreamHandler) {
 	rwc := newStreamReadWriteCloser(stream)
 	defer rwc.Close()
 
 	streamType, err := ReadStreamSignature(rwc)
 	if err != nil {
 		q.logger.Debug("failed to read stream signature: ", err)
+		stream.CancelWrite(0)
 		return
 	}
 
@@ -255,6 +264,7 @@ func (q *QUICConnection) handleStream(ctx context.Context, stream *quic.Stream,
 		request, err := ReadConnectRequest(rwc)
 		if err != nil {
 			q.logger.Debug("failed to read connect request: ", err)
+			stream.CancelWrite(0)
 			return
 		}
 		handler.HandleDataStream(ctx, &nopCloserReadWriter{ReadWriteCloser: rwc}, request, q.connIndex)
@@ -365,11 +375,11 @@ type DatagramSender interface {
 // streamReadWriteCloser adapts a *quic.Stream to io.ReadWriteCloser
 // with mutex-protected writes and safe close semantics.
 type streamReadWriteCloser struct {
-	stream      *quic.Stream
+	stream      quicStreamHandle
 	writeAccess sync.Mutex
 }
 
-func newStreamReadWriteCloser(stream *quic.Stream) *streamReadWriteCloser {
+func newStreamReadWriteCloser(stream quicStreamHandle) *streamReadWriteCloser {
 	return &streamReadWriteCloser{stream: stream}
 }