From bc9e32b8c307e280549380493cb9fd4504e606e0 Mon Sep 17 00:00:00 2001 From: David Turner Date: Mon, 29 Jan 2024 09:24:58 +0000 Subject: [PATCH] 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 --- docs/changelog/104808.yaml | 5 + .../elasticsearch/rest/RestControllerIT.java | 94 +++++++++++++++++++ .../elasticsearch/rest/RestController.java | 12 ++- 3 files changed, 107 insertions(+), 4 deletions(-) create mode 100644 docs/changelog/104808.yaml create mode 100644 server/src/internalClusterTest/java/org/elasticsearch/rest/RestControllerIT.java diff --git a/docs/changelog/104808.yaml b/docs/changelog/104808.yaml new file mode 100644 index 0000000000000..7682db085c7a9 --- /dev/null +++ b/docs/changelog/104808.yaml @@ -0,0 +1,5 @@ +pr: 104808 +summary: Fix lost headers with chunked responses +area: Network +type: bug +issues: [] diff --git a/server/src/internalClusterTest/java/org/elasticsearch/rest/RestControllerIT.java b/server/src/internalClusterTest/java/org/elasticsearch/rest/RestControllerIT.java new file mode 100644 index 0000000000000..81f8b962a9727 --- /dev/null +++ b/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> 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 getRestHandlers( + Settings settings, + RestController restController, + ClusterSettings clusterSettings, + IndexScopedSettings indexScopedSettings, + SettingsFilter settingsFilter, + IndexNameExpressionResolver indexNameExpressionResolver, + Supplier nodesInCluster + ) { + return List.of(new BaseRestHandler() { + @Override + public String getName() { + return ChunkedResponseWithHeadersPlugin.class.getCanonicalName(); + } + + @Override + public List 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); + }; + } + }); + } + } +} diff --git a/server/src/main/java/org/elasticsearch/rest/RestController.java b/server/src/main/java/org/elasticsearch/rest/RestController.java index 6a5d6f99df64b..d7563f3eb8ebd 100644 --- a/server/src/main/java/org/elasticsearch/rest/RestController.java +++ b/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;