Fix lost headers with chunked responses (#104845)

* Fix lost headers with chunked responses

In #99852 we introduced a layer of wrapping around the `RestResponse` to
capture the length of a chunked-encoded response, but this wrapping does
not preserve the headers of the original response. This commit fixes the
bug.

Backport of #104808 to 8.12

* Fix compile

Author: DaveCTurner

docs/changelog/104808.yaml

@@ -0,0 +1,5 @@
+pr: 104808
+summary: Fix lost headers with chunked responses
+area: Network
+type: bug
+issues: []

server/src/internalClusterTest/java/org/elasticsearch/rest/RestControllerIT.java

@@ -0,0 +1,94 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.rest;
+
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.internal.node.NodeClient;
+import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.node.DiscoveryNodes;
+import org.elasticsearch.common.collect.Iterators;
+import org.elasticsearch.common.settings.ClusterSettings;
+import org.elasticsearch.common.settings.IndexScopedSettings;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.settings.SettingsFilter;
+import org.elasticsearch.common.util.CollectionUtils;
+import org.elasticsearch.plugins.ActionPlugin;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.test.ESIntegTestCase;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.List;
+import java.util.function.Supplier;
+
+public class RestControllerIT extends ESIntegTestCase {
+    @Override
+    protected boolean addMockHttpTransport() {
+        return false; // enable HTTP
+    }
+
+    public void testHeadersEmittedWithChunkedResponses() throws IOException {
+        final var client = getRestClient();
+        final var response = client.performRequest(new Request("GET", ChunkedResponseWithHeadersPlugin.ROUTE));
+        assertEquals(200, response.getStatusLine().getStatusCode());
+        assertEquals(ChunkedResponseWithHeadersPlugin.HEADER_VALUE, response.getHeader(ChunkedResponseWithHeadersPlugin.HEADER_NAME));
+    }
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return CollectionUtils.appendToCopy(super.nodePlugins(), ChunkedResponseWithHeadersPlugin.class);
+    }
+
+    public static class ChunkedResponseWithHeadersPlugin extends Plugin implements ActionPlugin {
+
+        static final String ROUTE = "/_test/chunked_response_with_headers";
+        static final String HEADER_NAME = "test-header";
+        static final String HEADER_VALUE = "test-header-value";
+
+        @Override
+        public List<RestHandler> getRestHandlers(
+            Settings settings,
+            RestController restController,
+            ClusterSettings clusterSettings,
+            IndexScopedSettings indexScopedSettings,
+            SettingsFilter settingsFilter,
+            IndexNameExpressionResolver indexNameExpressionResolver,
+            Supplier<DiscoveryNodes> nodesInCluster
+        ) {
+            return List.of(new BaseRestHandler() {
+                @Override
+                public String getName() {
+                    return ChunkedResponseWithHeadersPlugin.class.getCanonicalName();
+                }
+
+                @Override
+                public List<Route> routes() {
+                    return List.of(new Route(RestRequest.Method.GET, ROUTE));
+                }
+
+                @Override
+                protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) {
+                    return channel -> {
+                        final var response = RestResponse.chunked(
+                            RestStatus.OK,
+                            ChunkedRestResponseBody.fromXContent(
+                                params -> Iterators.single((b, p) -> b.startObject().endObject()),
+                                request,
+                                channel,
+                                null
+                            )
+                        );
+                        response.addHeader(HEADER_NAME, HEADER_VALUE);
+                        channel.sendResponse(response);
+                    };
+                }
+            });
+        }
+    }
+}

server/src/main/java/org/elasticsearch/rest/RestController.java

@@ -791,10 +791,14 @@ public void sendResponse(RestResponse response) {
                 if (response.isChunked() == false) {
                     methodHandlers.addResponseStats(response.content().length());
                 } else {
-                    response = RestResponse.chunked(
-                        response.status(),
-                        new EncodedLengthTrackingChunkedRestResponseBody(response.chunkedContent(), methodHandlers)
-                    );
+                    final var wrapped = new EncodedLengthTrackingChunkedRestResponseBody(response.chunkedContent(), methodHandlers);
+                    final var headers = response.getHeaders();
+                    response = RestResponse.chunked(response.status(), wrapped);
+                    for (final var header : headers.entrySet()) {
+                        for (final var value : header.getValue()) {
+                            response.addHeader(header.getKey(), value);
+                        }
+                    }
                 }
                 delegate.sendResponse(response);
                 success = true;