Restrict handler metadata headers

connect_ext_test.go

@@ -2631,6 +2631,66 @@ func TestBlankImportCodeGeneration(t *testing.T) {
 	assert.NotNil(t, desc)
 }
 
+// TestSetProtocolHeaders tests that headers required by the protocols are set
+// overriding user provided headers.
+func TestSetProtocolHeaders(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name              string
+		clientOpts        []connect.ClientOption
+		proxyClientOpts   []connect.ClientOption
+		expectContentType string
+	}{{
+		name:              "connect",
+		clientOpts:        []connect.ClientOption{connect.WithGRPC()},
+		proxyClientOpts:   nil,
+		expectContentType: "application/proto",
+	}, {
+		name:              "grpc",
+		clientOpts:        nil,
+		proxyClientOpts:   []connect.ClientOption{connect.WithGRPC()},
+		expectContentType: "application/grpc",
+	}, {
+		name:              "grpcweb",
+		clientOpts:        nil,
+		proxyClientOpts:   []connect.ClientOption{connect.WithGRPCWeb()},
+		expectContentType: "application/grpc-web+proto",
+	}}
+	for _, tt := range tests {
+		testcase := tt
+		t.Run(testcase.name, func(t *testing.T) {
+			t.Parallel()
+			pingServer := &pingServer{}
+			mux := http.NewServeMux()
+			mux.Handle(pingv1connect.NewPingServiceHandler(pingServer))
+			server := memhttptest.NewServer(t, mux)
+
+			client := pingv1connect.NewPingServiceClient(server.Client(), server.URL(), testcase.clientOpts...)
+
+			pingProxyServer := &pluggablePingServer{
+				ping: func(ctx context.Context, request *connect.Request[pingv1.PingRequest]) (*connect.Response[pingv1.PingResponse], error) {
+					return client.Ping(ctx, request)
+				},
+			}
+			proxyMux := http.NewServeMux()
+			proxyMux.Handle(pingv1connect.NewPingServiceHandler(pingProxyServer))
+			proxyServer := memhttptest.NewServer(t, proxyMux)
+
+			proxyClient := pingv1connect.NewPingServiceClient(proxyServer.Client(), proxyServer.URL(), testcase.proxyClientOpts...)
+
+			request := connect.NewRequest(&pingv1.PingRequest{Number: 42})
+			request.Header().Set("X-Test", t.Name())
+			response, err := proxyClient.Ping(context.Background(), request)
+			if !assert.Nil(t, err) {
+				return
+			}
+			// Assert the Content-Type is set for the proxy clients protocol and not the client's.
+			assert.Equal(t, response.Header().Get("Content-Type"), testcase.expectContentType)
+			assert.Equal(t, len(response.Header().Values("Content-Type")), 1)
+		})
+	}
+}
+
 type unflushableWriter struct {
 	w http.ResponseWriter
 }

handler.go

@@ -71,8 +71,8 @@ func NewUnaryHandler[Req, Res any](
 		if err != nil {
 			return err
 		}
-		mergeHeaders(conn.ResponseHeader(), response.Header())
-		mergeHeaders(conn.ResponseTrailer(), response.Trailer())
+		mergeMetadataHeaders(conn.ResponseHeader(), response.Header())
+		mergeMetadataHeaders(conn.ResponseTrailer(), response.Trailer())
 		return conn.Send(response.Any())
 	}
 

header.go

@@ -45,6 +45,7 @@ func DecodeBinaryHeader(data string) ([]byte, error) {
 	return base64.StdEncoding.DecodeString(data)
 }
 
+// mergeHeaders merges the headers from the "from" header into the "into" header.
 func mergeHeaders(into, from http.Header) {
 	for key, vals := range from {
 		if len(vals) == 0 {
@@ -68,30 +69,7 @@ func mergeMetadataHeaders(into, from http.Header) {
 			// are no actual values for those keys, we skip them.
 			continue
 		}
-		switch http.CanonicalHeaderKey(key) {
-		case headerContentType,
-			headerContentLength,
-			headerContentEncoding,
-			headerHost,
-			headerUserAgent,
-			headerTrailer,
-			headerDate:
-			// HTTP headers.
-		case connectUnaryHeaderAcceptCompression,
-			connectUnaryTrailerPrefix,
-			connectStreamingHeaderCompression,
-			connectStreamingHeaderAcceptCompression,
-			connectHeaderTimeout,
-			connectHeaderProtocolVersion:
-			// Connect headers.
-		case grpcHeaderCompression,
-			grpcHeaderAcceptCompression,
-			grpcHeaderTimeout,
-			grpcHeaderStatus,
-			grpcHeaderMessage,
-			grpcHeaderDetails:
-			// gRPC headers.
-		default:
+		if !isProtocolHeader(key) {
 			into[key] = append(into[key], vals...)
 		}
 	}
@@ -124,3 +102,32 @@ func setHeaderCanonical(h http.Header, key, value string) {
 func delHeaderCanonical(h http.Header, key string) {
 	delete(h, key)
 }
+
+func isProtocolHeader(key string) bool {
+	switch http.CanonicalHeaderKey(key) {
+	case headerContentType,
+		headerContentLength,
+		headerContentEncoding,
+		headerHost,
+		headerUserAgent,
+		headerTrailer,
+		headerDate:
+		return true // HTTP headers.
+	case connectUnaryHeaderAcceptCompression,
+		connectUnaryTrailerPrefix,
+		connectStreamingHeaderCompression,
+		connectStreamingHeaderAcceptCompression,
+		connectHeaderTimeout,
+		connectHeaderProtocolVersion:
+		return true // Connect headers.
+	case grpcHeaderCompression,
+		grpcHeaderAcceptCompression,
+		grpcHeaderTimeout,
+		grpcHeaderStatus,
+		grpcHeaderMessage,
+		grpcHeaderDetails:
+		return true // gRPC headers.
+	default:
+		return false
+	}
+}