Fix: avoid further dispatch when ThreadContext population fails (#121665)

Author: ldematte

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

@@ -10,6 +10,7 @@
 package org.elasticsearch.rest;
 
 import org.elasticsearch.client.Request;
+import org.elasticsearch.client.RequestOptions;
 import org.elasticsearch.client.ResponseException;
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
@@ -20,6 +21,7 @@
 import org.elasticsearch.common.settings.IndexScopedSettings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.settings.SettingsFilter;
+import org.elasticsearch.common.xcontent.ChunkedToXContentHelper;
 import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
@@ -30,15 +32,21 @@
 import org.elasticsearch.telemetry.Measurement;
 import org.elasticsearch.telemetry.TestTelemetryPlugin;
 import org.elasticsearch.test.ESIntegTestCase;
+import org.elasticsearch.xcontent.XContentParser;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 
+import static org.elasticsearch.test.rest.ESRestTestCase.responseAsParser;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasEntry;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.instanceOf;
@@ -58,6 +66,49 @@ public void testHeadersEmittedWithChunkedResponses() throws IOException {
  assertEquals(ChunkedResponseWithHeadersPlugin.HEADER_VALUE, response.getHeader(ChunkedResponseWithHeadersPlugin.HEADER_NAME));
  }
 
+ public void testHeadersAreCollapsed() throws IOException {
+ final var client = getRestClient();
+ final var request = new Request("GET", TestEchoHeadersPlugin.ROUTE);
+ request.setOptions(RequestOptions.DEFAULT.toBuilder().addHeader("X-Foo", "1").addHeader("X-Foo", "2").build());
+ final var response = client.performRequest(request);
+ var responseMap = responseAsParser(response).map(HashMap::new, XContentParser::list);
+ assertThat(responseMap, hasEntry(equalTo("X-Foo"), containsInAnyOrder("1", "2")));
+ }
+
+ public void testHeadersTreatedCaseInsensitive() throws IOException {
+ final var client = getRestClient();
+ final var request = new Request("GET", TestEchoHeadersPlugin.ROUTE);
+ request.setOptions(RequestOptions.DEFAULT.toBuilder().addHeader("X-Foo", "1").addHeader("x-foo", "2").build());
+ final var response = client.performRequest(request);
+ var responseMap = responseAsParser(response).map(HashMap::new, XContentParser::list);
+ assertThat(responseMap, hasEntry(equalTo("x-foo"), containsInAnyOrder("1", "2")));
+ assertThat(responseMap, hasEntry(equalTo("X-Foo"), containsInAnyOrder("1", "2")));
+ }
+
+ public void testThreadContextPopulationFromMultipleHeadersFailsWithCorrectError() {
+ final var client = getRestClient();
+ final var sameCaseRequest = new Request("GET", TestEchoHeadersPlugin.ROUTE);
+ sameCaseRequest.setOptions(
+ RequestOptions.DEFAULT.toBuilder()
+ .addHeader("x-elastic-product-origin", "elastic")
+ .addHeader("x-elastic-product-origin", "other")
+ );
+ var exception1 = expectThrows(ResponseException.class, () -> client.performRequest(sameCaseRequest));
+ assertThat(exception1.getMessage(), containsString("multiple values for single-valued header [X-elastic-product-origin]"));
+ }
+
+ public void testMultipleProductOriginHeadersWithDifferentCaseFailsWithCorrectError() {
+ final var client = getRestClient();
+ final var differentCaseRequest = new Request("GET", TestEchoHeadersPlugin.ROUTE);
+ differentCaseRequest.setOptions(
+ RequestOptions.DEFAULT.toBuilder()
+ .addHeader("X-elastic-product-origin", "elastic")
+ .addHeader("x-elastic-product-origin", "other")
+ );
+ var exception2 = expectThrows(ResponseException.class, () -> client.performRequest(differentCaseRequest));
+ assertThat(exception2.getMessage(), containsString("multiple values for single-valued header [X-elastic-product-origin]"));
+ }
+
  public void testMetricsEmittedOnSuccess() throws Exception {
  final var client = getRestClient();
  final var request = new Request("GET", TestEchoStatusCodePlugin.ROUTE);
@@ -125,7 +176,12 @@ private void assertMeasurement(Consumer<Measurement> measurementConsumer) throws
 
  @Override
  protected Collection<Class<? extends Plugin>> nodePlugins() {
- return List.of(ChunkedResponseWithHeadersPlugin.class, TestEchoStatusCodePlugin.class, TestTelemetryPlugin.class);
+ return List.of(
+ ChunkedResponseWithHeadersPlugin.class,
+ TestEchoStatusCodePlugin.class,
+ TestEchoHeadersPlugin.class,
+ TestTelemetryPlugin.class
+ );
  }
 
  public static class TestEchoStatusCodePlugin extends Plugin implements ActionPlugin {
@@ -181,6 +237,62 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
  }
  }
 
+ public static class TestEchoHeadersPlugin extends Plugin implements ActionPlugin {
+ static final String ROUTE = "/_test/echo_headers";
+ static final String NAME = "test_echo_headers";
+
+ private static final Logger logger = LogManager.getLogger(TestEchoStatusCodePlugin.class);
+
+ @Override
+ public Collection<RestHandler> getRestHandlers(
+ Settings settings,
+ NamedWriteableRegistry namedWriteableRegistry,
+ RestController restController,
+ ClusterSettings clusterSettings,
+ IndexScopedSettings indexScopedSettings,
+ SettingsFilter settingsFilter,
+ IndexNameExpressionResolver indexNameExpressionResolver,
+ Supplier<DiscoveryNodes> nodesInCluster,
+ Predicate<NodeFeature> clusterSupportsFeature
+ ) {
+ return List.of(new BaseRestHandler() {
+ @Override
+ public String getName() {
+ return NAME;
+ }
+
+ @Override
+ public List<Route> routes() {
+ return List.of(new Route(RestRequest.Method.GET, ROUTE), new Route(RestRequest.Method.POST, ROUTE));
+ }
+
+ @Override
+ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) {
+ var headers = request.getHeaders();
+ logger.info("received header echo request for [{}]", String.join(",", headers.keySet()));
+
+ return channel -> {
+ final var response = RestResponse.chunked(
+ RestStatus.OK,
+ ChunkedRestResponseBodyPart.fromXContent(
+ params -> Iterators.concat(
+ ChunkedToXContentHelper.startObject(),
+ Iterators.map(headers.entrySet().iterator(), e -> (b, p) -> b.field(e.getKey(), e.getValue())),
+ ChunkedToXContentHelper.endObject()
+ ),
+ request,
+ channel
+ ),
+ null
+ );
+ channel.sendResponse(response);
+ logger.info("sent response");
+ };
+ }
+ });
+ }
+ }
+
  public static class ChunkedResponseWithHeadersPlugin extends Plugin implements ActionPlugin {
 
  static final String ROUTE = "/_test/chunked_response_with_headers";

server/src/main/java/org/elasticsearch/http/AbstractHttpServerTransport.java

@@ -38,7 +38,6 @@
 import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.rest.RestChannel;
 import org.elasticsearch.rest.RestRequest;
-import org.elasticsearch.rest.RestResponse;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.telemetry.tracing.Tracer;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -484,25 +483,22 @@ void dispatchRequest(final RestRequest restRequest, final RestChannel channel, f
  if (badRequestCause != null) {
  dispatcher.dispatchBadRequest(channel, threadContext, badRequestCause);
  } else {
- populatePerRequestThreadContext0(restRequest, channel, threadContext);
+ try {
+ populatePerRequestThreadContext(restRequest, threadContext);
+ } catch (Exception e) {
+ try {
+ dispatcher.dispatchBadRequest(channel, threadContext, e);
+ } catch (Exception inner) {
+ inner.addSuppressed(e);
+ logger.error(() -> "failed to send failure response for uri [" + restRequest.uri() + "]", inner);
+ }
+ return;
+ }
  dispatcher.dispatchRequest(restRequest, channel, threadContext);
  }
  }
  }
 
- private void populatePerRequestThreadContext0(RestRequest restRequest, RestChannel channel, ThreadContext threadContext) {
- try {
- populatePerRequestThreadContext(restRequest, threadContext);
- } catch (Exception e) {
- try {
- channel.sendResponse(new RestResponse(channel, e));
- } catch (Exception inner) {
- inner.addSuppressed(e);
- logger.error(() -> "failed to send failure response for uri [" + restRequest.uri() + "]", inner);
- }
- }
- }
-
  protected void populatePerRequestThreadContext(RestRequest restRequest, ThreadContext threadContext) {}
 
  private void handleIncomingRequest(final HttpRequest httpRequest, final HttpChannel httpChannel, final Exception exception) {

server/src/test/java/org/elasticsearch/http/AbstractHttpServerTransportTests.java

@@ -355,8 +355,6 @@ public void dispatchBadRequest(final RestChannel channel, final ThreadContext th
  }
 
  };
- // the set of headers to copy
- Set<RestHeaderDefinition> headers = Set.of(new RestHeaderDefinition(Task.TRACE_PARENT_HTTP_HEADER, false));
  // sample request headers to test with
  Map<String, List<String>> restHeaders = new HashMap<>();
  restHeaders.put(Task.TRACE_PARENT_HTTP_HEADER, Collections.singletonList(traceParentValue));
@@ -397,7 +395,7 @@ public HttpStats stats() {
 
  @Override
  protected void populatePerRequestThreadContext(RestRequest restRequest, ThreadContext threadContext) {
- getFakeActionModule(headers).copyRequestHeadersToThreadContext(restRequest.getHttpRequest(), threadContext);
+ getFakeActionModule(Set.of()).copyRequestHeadersToThreadContext(restRequest.getHttpRequest(), threadContext);
  }
  }
  ) {

x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/transport/netty4/SecurityNetty4HttpServerTransportTests.java

@@ -370,7 +370,21 @@ public void testHttpHeaderAuthnBypassHeaderValidator() throws Exception {
  new NetworkService(List.of()),
  testThreadPool,
  xContentRegistry(),
- new NullDispatcher(),
+ new HttpServerTransport.Dispatcher() {
+ @Override
+ public void dispatchRequest(RestRequest request, RestChannel channel, ThreadContext threadContext) {
+ fail("Request should not be dispatched");
+ }
+
+ @Override
+ public void dispatchBadRequest(RestChannel channel, ThreadContext threadContext, Throwable cause) {
+ try {
+ channel.sendResponse(new RestResponse(channel, (Exception) cause));
+ } catch (IOException e) {
+ fail(e, "Unexpected exception dispatching bad request");
+ }
+ }
+ },
  randomClusterSettings(),
  new SharedGroupFactory(settings),
  Tracer.NOOP,