Convert RestWrapper to an explicit Interceptor

Adds a new `RestInterceptor` interface and converts
`RestServerActionPlugin.getRestHandlerInterceptor` to return this new
type instead of a wrapping function.

This has the following benefits:
- Less object creation, there is 1 instance of the interceptor class
  (see `SecurityRestFilter`) rather than an instance per handler
- More control over the sequence of steps in processing a request.
  The explicit interceptor separates it from the deprecation handler
  or any validation that might be needed, and the controller can be
  intentional about the order in which these operations are applied.

qa/custom-rest-controller/src/javaRestTest/java/co/elastic/elasticsearch/test/CustomRestPlugin.java

@@ -8,6 +8,7 @@
 
 package co.elastic.elasticsearch.test;
 
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
@@ -18,12 +19,11 @@
 import org.elasticsearch.rest.RestChannel;
 import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestHandler;
+import org.elasticsearch.rest.RestInterceptor;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.telemetry.tracing.Tracer;
 import org.elasticsearch.usage.UsageService;
 
-import java.util.function.UnaryOperator;
-
 public class CustomRestPlugin extends Plugin implements RestServerActionPlugin {
 
     private static final Logger logger = LogManager.getLogger(CustomRestPlugin.class);
@@ -35,34 +35,33 @@ private static void echoHeader(String name, RestRequest request, ThreadContext t
         }
     }
 
-    public static class CustomInterceptor implements RestHandler {
+    public static class CustomInterceptor implements RestInterceptor {
 
         private final ThreadContext threadContext;
-        private final RestHandler delegate;
 
-        public CustomInterceptor(ThreadContext threadContext, RestHandler delegate) {
+        public CustomInterceptor(ThreadContext threadContext) {
             this.threadContext = threadContext;
-            this.delegate = delegate;
         }
 
         @Override
-        public void handleRequest(RestRequest request, RestChannel channel, NodeClient client) throws Exception {
+        public void intercept(RestRequest request, RestChannel channel, RestHandler targetHandler, ActionListener<Boolean> listener)
+            throws Exception {
             logger.info("intercept request {} {}", request.method(), request.uri());
             echoHeader("x-test-interceptor", request, threadContext);
-            delegate.handleRequest(request, channel, client);
+            listener.onResponse(Boolean.TRUE);
         }
 
     }
 
     public static class CustomController extends RestController {
         public CustomController(
-            UnaryOperator<RestHandler> handlerWrapper,
+            RestInterceptor interceptor,
             NodeClient client,
             CircuitBreakerService circuitBreakerService,
             UsageService usageService,
             Tracer tracer
         ) {
-            super(handlerWrapper, client, circuitBreakerService, usageService, tracer);
+            super(interceptor, client, circuitBreakerService, usageService, tracer);
         }
 
         @Override
@@ -74,19 +73,19 @@ public void dispatchRequest(RestRequest request, RestChannel channel, ThreadCont
     }
 
     @Override
-    public UnaryOperator<RestHandler> getRestHandlerInterceptor(ThreadContext threadContext) {
-        return handler -> new CustomInterceptor(threadContext, handler);
+    public RestInterceptor getRestHandlerInterceptor(ThreadContext threadContext) {
+        return new CustomInterceptor(threadContext);
     }
 
     @Override
     public RestController getRestController(
-        UnaryOperator<RestHandler> handlerWrapper,
+        RestInterceptor interceptor,
         NodeClient client,
         CircuitBreakerService circuitBreakerService,
         UsageService usageService,
         Tracer tracer
     ) {
-        return new CustomController(handlerWrapper, client, circuitBreakerService, usageService, tracer);
+        return new CustomController(interceptor, client, circuitBreakerService, usageService, tracer);
     }
 
 }

server/src/main/java/org/elasticsearch/action/ActionModule.java

@@ -264,6 +264,7 @@
 import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestHandler;
 import org.elasticsearch.rest.RestHeaderDefinition;
+import org.elasticsearch.rest.RestInterceptor;
 import org.elasticsearch.rest.RestUtils;
 import org.elasticsearch.rest.action.RestFieldCapabilitiesAction;
 import org.elasticsearch.rest.action.admin.cluster.RestAddVotingConfigExclusionAction;
@@ -425,7 +426,6 @@
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
-import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -501,7 +501,7 @@ public ActionModule(
                 new RestHeaderDefinition(Task.X_ELASTIC_PRODUCT_ORIGIN_HTTP_HEADER, false)
             )
         ).collect(Collectors.toSet());
-        UnaryOperator<RestHandler> restInterceptor = getRestServerComponent(
+        final RestInterceptor restInterceptor = getRestServerComponent(
             "REST interceptor",
             actionPlugins,
             restPlugin -> restPlugin.getRestHandlerInterceptor(threadPool.getThreadContext())

server/src/main/java/org/elasticsearch/plugins/interceptor/RestServerActionPlugin.java

@@ -14,7 +14,7 @@
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
 import org.elasticsearch.plugins.ActionPlugin;
 import org.elasticsearch.rest.RestController;
-import org.elasticsearch.rest.RestHandler;
+import org.elasticsearch.rest.RestInterceptor;
 import org.elasticsearch.telemetry.tracing.Tracer;
 import org.elasticsearch.usage.UsageService;
 
@@ -46,15 +46,15 @@ public interface RestServerActionPlugin extends ActionPlugin {
      *
      * Note: Only one installed plugin may implement a rest interceptor.
      */
-    UnaryOperator<RestHandler> getRestHandlerInterceptor(ThreadContext threadContext);
+    RestInterceptor getRestHandlerInterceptor(ThreadContext threadContext);
 
     /**
      * Returns a replacement {@link RestController} to be used in the server.
      * Note: Only one installed plugin may override the rest controller.
      */
     @Nullable
     default RestController getRestController(
-        @Nullable UnaryOperator<RestHandler> handlerWrapper,
+        @Nullable RestInterceptor interceptor,
         NodeClient client,
         CircuitBreakerService circuitBreakerService,
         UsageService usageService,

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

@@ -13,6 +13,8 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.ElasticsearchStatusException;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.breaker.CircuitBreaker;
@@ -57,7 +59,6 @@
 import java.util.TreeMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
-import java.util.function.UnaryOperator;
 
 import static org.elasticsearch.indices.SystemIndices.EXTERNAL_SYSTEM_INDEX_ACCESS_CONTROL_HEADER_KEY;
 import static org.elasticsearch.indices.SystemIndices.SYSTEM_INDEX_ACCESS_CONTROL_HEADER_KEY;
@@ -95,7 +96,7 @@ public class RestController implements HttpServerTransport.Dispatcher {
 
     private final PathTrie<MethodHandlers> handlers = new PathTrie<>(RestUtils.REST_DECODER);
 
-    private final UnaryOperator<RestHandler> handlerWrapper;
+    private final RestInterceptor interceptor;
 
     private final NodeClient client;
 
@@ -107,18 +108,18 @@ public class RestController implements HttpServerTransport.Dispatcher {
     private final ServerlessApiProtections apiProtections;
 
     public RestController(
-        UnaryOperator<RestHandler> handlerWrapper,
+        RestInterceptor restInterceptor,
         NodeClient client,
         CircuitBreakerService circuitBreakerService,
         UsageService usageService,
         Tracer tracer
     ) {
         this.usageService = usageService;
         this.tracer = tracer;
-        if (handlerWrapper == null) {
-            handlerWrapper = h -> h; // passthrough if no wrapper set
+        if (restInterceptor == null) {
+            restInterceptor = (request, channel, targetHandler, listener) -> listener.onResponse(Boolean.TRUE);
         }
-        this.handlerWrapper = handlerWrapper;
+        this.interceptor = restInterceptor;
         this.client = client;
         this.circuitBreakerService = circuitBreakerService;
         registerHandlerNoWrap(RestRequest.Method.GET, "/favicon.ico", RestApiVersion.current(), new RestFavIconHandler());
@@ -264,7 +265,7 @@ protected void registerHandler(RestRequest.Method method, String path, RestApiVe
         if (handler instanceof BaseRestHandler) {
             usageService.addRestHandler((BaseRestHandler) handler);
         }
-        registerHandlerNoWrap(method, path, version, handlerWrapper.apply(handler));
+        registerHandlerNoWrap(method, path, version, handler);
     }
 
     private void registerHandlerNoWrap(RestRequest.Method method, String path, RestApiVersion version, RestHandler handler) {
@@ -325,7 +326,7 @@ public void dispatchRequest(RestRequest request, RestChannel channel, ThreadCont
             tryAllHandlers(request, channel, threadContext);
         } catch (Exception e) {
             try {
-                channel.sendResponse(new RestResponse(channel, e));
+                sendFailure(channel, e);
             } catch (Exception inner) {
                 inner.addSuppressed(e);
                 logger.error(() -> "failed to send failure response for uri [" + request.uri() + "]", inner);
@@ -348,7 +349,7 @@ public void dispatchBadRequest(final RestChannel channel, final ThreadContext th
             // unless it's a http headers validation error, we consider any exceptions encountered so far during request processing
             // to be a problem of invalid/malformed request (hence the RestStatus#BAD_REQEST (400) HTTP response code)
             if (e instanceof HttpHeadersValidationException) {
-                channel.sendResponse(new RestResponse(channel, (Exception) e.getCause()));
+                sendFailure(channel, (Exception) e.getCause());
             } else {
                 channel.sendResponse(new RestResponse(channel, BAD_REQUEST, e));
             }
@@ -438,12 +439,44 @@ private void dispatchRequest(
             } else {
                 threadContext.putHeader(SYSTEM_INDEX_ACCESS_CONTROL_HEADER_KEY, Boolean.TRUE.toString());
             }
-            handler.handleRequest(request, responseChannel, client);
+            final var finalChannel = responseChannel;
+            this.interceptor.intercept(request, responseChannel, handler.getConcreteRestHandler(), new ActionListener<>() {
+                @Override
+                public void onResponse(Boolean processRequest) {
+                    if (processRequest) {
+                        try {
+                            validateRequest(request, handler, client);
+                            handler.handleRequest(request, finalChannel, client);
+                        } catch (Exception e) {
+                            onFailure(e);
+                        }
+                    }
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    try {
+                        sendFailure(finalChannel, e);
+                    } catch (IOException ex) {
+                        logger.info("Failed to send error [{}] to HTTP client", ex.toString());
+                    }
+                }
+            });
         } catch (Exception e) {
-            responseChannel.sendResponse(new RestResponse(responseChannel, e));
+            sendFailure(responseChannel, e);
         }
     }
 
+    /**
+     * Validates that the request should be allowed. Throws an exception if the request should be rejected.
+     */
+    @SuppressWarnings("unused")
+    protected void validateRequest(RestRequest request, RestHandler handler, NodeClient client) throws ElasticsearchStatusException {}
+
+    private static void sendFailure(RestChannel responseChannel, Exception e) throws IOException {
+        responseChannel.sendResponse(new RestResponse(responseChannel, e));
+    }
+
     /**
      * in order to prevent CSRF we have to reject all media types that are from a browser safelist
      * see https://fetch.spec.whatwg.org/#cors-safelisted-request-header
@@ -691,7 +724,7 @@ public static void handleBadRequest(String uri, RestRequest.Method method, RestC
     public static void handleServerlessRequestToProtectedResource(String uri, RestRequest.Method method, RestChannel channel)
         throws IOException {
         String msg = "uri [" + uri + "] with method [" + method + "] exists but is not available when running in serverless mode";
-        channel.sendResponse(new RestResponse(channel, new ApiNotAvailableException(msg)));
+        sendFailure(channel, new ApiNotAvailableException(msg));
     }
 
     /**

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

@@ -0,0 +1,27 @@
+/*
+ * 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.action.ActionListener;
+
+/**
+ * Wraps the execution of a {@link RestHandler}
+ */
+@FunctionalInterface
+public interface RestInterceptor {
+
+    /**
+     * @param listener The interceptor responds with {@code True} if the handler should be called,
+     *                 or {@code False} if the request has been entirely handled by the interceptor.
+     *                 In the case of {@link ActionListener#onFailure(Exception)}, the target handler
+     *                 will not be called, the request will be treated as unhandled, and the regular
+     *                 rest exception handling will be performed
+     */
+    void intercept(RestRequest request, RestChannel channel, RestHandler targetHandler, ActionListener<Boolean> listener) throws Exception;
+}

server/src/test/java/org/elasticsearch/action/ActionModuleTests.java

@@ -30,6 +30,7 @@
 import org.elasticsearch.rest.RestChannel;
 import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestHandler;
+import org.elasticsearch.rest.RestInterceptor;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.action.admin.cluster.RestNodesInfoAction;
 import org.elasticsearch.tasks.Task;
@@ -45,7 +46,6 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.function.Supplier;
-import java.util.function.UnaryOperator;
 
 import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
@@ -362,24 +362,24 @@ class SecPlugin implements ActionPlugin, RestServerActionPlugin {
         }
 
         @Override
-        public UnaryOperator<RestHandler> getRestHandlerInterceptor(ThreadContext threadContext) {
+        public RestInterceptor getRestHandlerInterceptor(ThreadContext threadContext) {
             if (installInterceptor) {
-                return UnaryOperator.identity();
+                return (request, channel, targetHandler, listener) -> listener.onResponse(true);
             } else {
                 return null;
             }
         }
 
         @Override
         public RestController getRestController(
-            UnaryOperator<RestHandler> handlerWrapper,
+            RestInterceptor interceptor,
             NodeClient client,
             CircuitBreakerService circuitBreakerService,
             UsageService usageService,
             Tracer tracer
         ) {
             if (installController) {
-                return new RestController(handlerWrapper, client, circuitBreakerService, usageService, tracer);
+                return new RestController(interceptor, client, circuitBreakerService, usageService, tracer);
             } else {
                 return null;
             }

server/src/test/java/org/elasticsearch/rest/RestControllerTests.java

@@ -286,22 +286,25 @@ public void testRegisterSecondMethodWithDifferentNamedWildcard() {
         assertThat(exception.getMessage(), equalTo("Trying to use conflicting wildcard names for same path: wildcard1 and wildcard2"));
     }
 
-    public void testRestHandlerWrapper() throws Exception {
+    public void testRestInterceptor() throws Exception {
         AtomicBoolean handlerCalled = new AtomicBoolean(false);
         AtomicBoolean wrapperCalled = new AtomicBoolean(false);
+        final boolean callHandler = randomBoolean();
         final RestHandler handler = (RestRequest request, RestChannel channel, NodeClient client) -> handlerCalled.set(true);
         final HttpServerTransport httpServerTransport = new TestHttpServerTransport();
-        final RestController restController = new RestController(h -> {
-            assertSame(handler, h);
-            return (RestRequest request, RestChannel channel, NodeClient client) -> wrapperCalled.set(true);
-        }, client, circuitBreakerService, usageService, tracer);
+        final RestInterceptor interceptor = (request, channel, targetHandler, listener) -> {
+            assertSame(handler, targetHandler);
+            wrapperCalled.set(true);
+            listener.onResponse(callHandler);
+        };
+        final RestController restController = new RestController(interceptor, client, circuitBreakerService, usageService, tracer);
         restController.registerHandler(new Route(GET, "/wrapped"), handler);
         RestRequest request = testRestRequest("/wrapped", "{}", XContentType.JSON);
         AssertingChannel channel = new AssertingChannel(request, true, RestStatus.BAD_REQUEST);
         restController.dispatchRequest(request, channel, client.threadPool().getThreadContext());
         httpServerTransport.start();
-        assertTrue(wrapperCalled.get());
-        assertFalse(handlerCalled.get());
+        assertThat(wrapperCalled.get(), is(true));
+        assertThat(handlerCalled.get(), is(callHandler));
     }
 
     public void testDispatchRequestAddsAndFreesBytesOnSuccess() {

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/LocalStateCompositeXPackPlugin.java

@@ -83,6 +83,7 @@
 import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestHandler;
 import org.elasticsearch.rest.RestHeaderDefinition;
+import org.elasticsearch.rest.RestInterceptor;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.search.internal.ShardSearchRequest;
 import org.elasticsearch.snapshots.Snapshot;
@@ -381,10 +382,9 @@ public List<BootstrapCheck> getBootstrapChecks() {
     }
 
     @Override
-    public UnaryOperator<RestHandler> getRestHandlerInterceptor(ThreadContext threadContext) {
-
+    public RestInterceptor getRestHandlerInterceptor(ThreadContext threadContext) {
         // There can be only one.
-        List<UnaryOperator<RestHandler>> items = filterPlugins(ActionPlugin.class).stream()
+        List<RestInterceptor> items = filterPlugins(ActionPlugin.class).stream()
             .filter(RestServerActionPlugin.class::isInstance)
             .map(RestServerActionPlugin.class::cast)
             .map(p -> p.getRestHandlerInterceptor(threadContext))