TUN-6676: Add suport for trailers in http2 connections

connection/connection.go

@@ -24,9 +24,16 @@ const (
 	LogFieldConnIndex      = "connIndex"
 	MaxGracePeriod         = time.Minute * 3
 	MaxConcurrentStreams   = math.MaxUint32
+
+	contentTypeHeader = "content-type"
+	sseContentType    = "text/event-stream"
+	grpcContentType   = "application/grpc"
 )
 
-var switchingProtocolText = fmt.Sprintf("%d %s", http.StatusSwitchingProtocols, http.StatusText(http.StatusSwitchingProtocols))
+var (
+	switchingProtocolText = fmt.Sprintf("%d %s", http.StatusSwitchingProtocols, http.StatusText(http.StatusSwitchingProtocols))
+	flushableContentTypes = []string{sseContentType, grpcContentType}
+)
 
 type Orchestrator interface {
 	UpdateConfig(version int32, config []byte) *pogs.UpdateConfigurationResponse
@@ -190,6 +197,7 @@ func (h *HTTPResponseReadWriteAcker) AckConnection(tracePropagation string) erro
 
 type ResponseWriter interface {
 	WriteRespHeaders(status int, header http.Header) error
+	AddTrailer(trailerName, trailerValue string)
 	io.Writer
 }
 
@@ -198,10 +206,18 @@ type ConnectedFuse interface {
 	IsConnected() bool
 }
 
-func IsServerSentEvent(headers http.Header) bool {
-	if contentType := headers.Get("content-type"); contentType != "" {
-		return strings.HasPrefix(strings.ToLower(contentType), "text/event-stream")
+// Helper method to let the caller know what content-types should require a flush on every
+// write to a ResponseWriter.
+func shouldFlush(headers http.Header) bool {
+	if contentType := headers.Get(contentTypeHeader); contentType != "" {
+		contentType = strings.ToLower(contentType)
+		for _, c := range flushableContentTypes {
+			if strings.HasPrefix(contentType, c) {
+				return true
+			}
+		}
 	}
+
 	return false
 }
 

connection/connection_test.go

@@ -6,11 +6,9 @@ import (
 	"io"
 	"math/rand"
 	"net/http"
-	"testing"
 	"time"
 
 	"github.com/rs/zerolog"
-	"github.com/stretchr/testify/assert"
 
 	"github.com/cloudflare/cloudflared/tracing"
 	tunnelpogs "github.com/cloudflare/cloudflared/tunnelrpc/pogs"
@@ -197,40 +195,3 @@ func (mcf mockConnectedFuse) Connected() {}
 func (mcf mockConnectedFuse) IsConnected() bool {
 	return true
 }
-
-func TestIsEventStream(t *testing.T) {
-	tests := []struct {
-		headers       http.Header
-		isEventStream bool
-	}{
-		{
-			headers:       newHeader("Content-Type", "text/event-stream"),
-			isEventStream: true,
-		},
-		{
-			headers:       newHeader("content-type", "text/event-stream"),
-			isEventStream: true,
-		},
-		{
-			headers:       newHeader("Content-Type", "text/event-stream; charset=utf-8"),
-			isEventStream: true,
-		},
-		{
-			headers:       newHeader("Content-Type", "application/json"),
-			isEventStream: false,
-		},
-		{
-			headers:       http.Header{},
-			isEventStream: false,
-		},
-	}
-	for _, test := range tests {
-		assert.Equal(t, test.isEventStream, IsServerSentEvent(test.headers))
-	}
-}
-
-func newHeader(key, value string) http.Header {
-	header := http.Header{}
-	header.Add(key, value)
-	return header
-}

connection/h2mux.go

@@ -259,6 +259,10 @@ type h2muxRespWriter struct {
 	*h2mux.MuxedStream
 }
 
+func (rp *h2muxRespWriter) AddTrailer(trailerName, trailerValue string) {
+	// do nothing. we don't support trailers over h2mux
+}
+
 func (rp *h2muxRespWriter) WriteRespHeaders(status int, header http.Header) error {
 	headers := H1ResponseToH2ResponseHeaders(status, header)
 	headers = append(headers, h2mux.Header{Name: ResponseMetaHeader, Value: responseMetaHeaderOrigin})

connection/http2.go

@@ -191,11 +191,12 @@ func (c *HTTP2Connection) close() {
 }
 
 type http2RespWriter struct {
-	r           io.Reader
-	w           http.ResponseWriter
-	flusher     http.Flusher
-	shouldFlush bool
-	log         *zerolog.Logger
+	r             io.Reader
+	w             http.ResponseWriter
+	flusher       http.Flusher
+	shouldFlush   bool
+	statusWritten bool
+	log           *zerolog.Logger
 }
 
 func NewHTTP2RespWriter(r *http.Request, w http.ResponseWriter, connType Type, log *zerolog.Logger) (*http2RespWriter, error) {
@@ -219,11 +220,20 @@ func NewHTTP2RespWriter(r *http.Request, w http.ResponseWriter, connType Type, l
 	}, nil
 }
 
+func (rp *http2RespWriter) AddTrailer(trailerName, trailerValue string) {
+	if !rp.statusWritten {
+		rp.log.Warn().Msg("Tried to add Trailer to response before status written. Ignoring...")
+		return
+	}
+
+	rp.w.Header().Add(http2.TrailerPrefix+trailerName, trailerValue)
+}
+
 func (rp *http2RespWriter) WriteRespHeaders(status int, header http.Header) error {
 	dest := rp.w.Header()
 	userHeaders := make(http.Header, len(header))
 	for name, values := range header {
-		// Since these are http2 headers, they're required to be lowercase
+		// lowercase headers for simplicity check
 		h2name := strings.ToLower(name)
 
 		if h2name == "content-length" {
@@ -234,7 +244,7 @@ func (rp *http2RespWriter) WriteRespHeaders(status int, header http.Header) erro
 
 		if h2name == tracing.IntCloudflaredTracingHeader {
 			// Add cf-int-cloudflared-tracing header outside of serialized userHeaders
-			rp.w.Header()[tracing.CanonicalCloudflaredTracingHeader] = values
+			dest[tracing.CanonicalCloudflaredTracingHeader] = values
 			continue
 		}
 
@@ -247,18 +257,21 @@ func (rp *http2RespWriter) WriteRespHeaders(status int, header http.Header) erro
 
 	// Perform user header serialization and set them in the single header
 	dest.Set(CanonicalResponseUserHeaders, SerializeHeaders(userHeaders))
+
 	rp.setResponseMetaHeader(responseMetaHeaderOrigin)
 	// HTTP2 removes support for 101 Switching Protocols https://tools.ietf.org/html/rfc7540#section-8.1.1
 	if status == http.StatusSwitchingProtocols {
 		status = http.StatusOK
 	}
 	rp.w.WriteHeader(status)
-	if IsServerSentEvent(header) {
+	if shouldFlush(header) {
 		rp.shouldFlush = true
 	}
 	if rp.shouldFlush {
 		rp.flusher.Flush()
 	}
+
+	rp.statusWritten = true
 	return nil
 }
 

connection/quic.go

@@ -329,6 +329,10 @@ func newHTTPResponseAdapter(s *quicpogs.RequestServerStream) httpResponseAdapter
 	return httpResponseAdapter{s}
 }
 
+func (hrw httpResponseAdapter) AddTrailer(trailerName, trailerValue string) {
+	// we do not support trailers over QUIC
+}
+
 func (hrw httpResponseAdapter) WriteRespHeaders(status int, header http.Header) error {
 	metadata := make([]quicpogs.Metadata, 0)
 	metadata = append(metadata, quicpogs.Metadata{Key: "HttpStatus", Val: strconv.Itoa(status)})