Make RestController pluggable (#98187)

This commit changes the ActionModules to allow  the RestController to be
provided by an internal plugin.

It renames  `RestInterceptorActionPlugin` to `RestServerActionPlugin`
and adds a new `getRestController` method to it.

There may be multiple RestServerActionPlugins installed on a node, but
only 1 may provide a Rest Wrapper (getRestHandlerInterceptor) and only 1
may provide a RestController (getRestController).

docs/changelog/98187.yaml

@@ -0,0 +1,5 @@
+pr: 98187
+summary: Make `RestController` pluggable
+area: Infra/REST API
+type: enhancement
+issues: []

qa/custom-rest-controller/build.gradle

@@ -0,0 +1,10 @@
+/*
+ * 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.
+ */
+
+apply plugin: 'elasticsearch.internal-java-rest-test'
+

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

@@ -0,0 +1,92 @@
+/*
+ * 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 co.elastic.elasticsearch.test;
+
+import org.elasticsearch.client.internal.node.NodeClient;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
+import org.elasticsearch.indices.breaker.CircuitBreakerService;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.plugins.interceptor.RestServerActionPlugin;
+import org.elasticsearch.rest.RestChannel;
+import org.elasticsearch.rest.RestController;
+import org.elasticsearch.rest.RestHandler;
+import org.elasticsearch.rest.RestRequest;
+import org.elasticsearch.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);
+
+    private static void echoHeader(String name, RestRequest request, ThreadContext threadContext) {
+        var value = request.header(name);
+        if (value != null) {
+            threadContext.addResponseHeader(name, value);
+        }
+    }
+
+    public static class CustomInterceptor implements RestHandler {
+
+        private final ThreadContext threadContext;
+        private final RestHandler delegate;
+
+        public CustomInterceptor(ThreadContext threadContext, RestHandler delegate) {
+            this.threadContext = threadContext;
+            this.delegate = delegate;
+        }
+
+        @Override
+        public void handleRequest(RestRequest request, RestChannel channel, NodeClient client) throws Exception {
+            logger.info("intercept request {} {}", request.method(), request.uri());
+            echoHeader("x-test-interceptor", request, threadContext);
+            delegate.handleRequest(request, channel, client);
+        }
+
+    }
+
+    public static class CustomController extends RestController {
+        public CustomController(
+            UnaryOperator<RestHandler> handlerWrapper,
+            NodeClient client,
+            CircuitBreakerService circuitBreakerService,
+            UsageService usageService,
+            Tracer tracer
+        ) {
+            super(handlerWrapper, client, circuitBreakerService, usageService, tracer);
+        }
+
+        @Override
+        public void dispatchRequest(RestRequest request, RestChannel channel, ThreadContext threadContext) {
+            logger.info("dispatch request {} {}", request.method(), request.uri());
+            echoHeader("x-test-controller", request, threadContext);
+            super.dispatchRequest(request, channel, threadContext);
+        }
+    }
+
+    @Override
+    public UnaryOperator<RestHandler> getRestHandlerInterceptor(ThreadContext threadContext) {
+        return handler -> new CustomInterceptor(threadContext, handler);
+    }
+
+    @Override
+    public RestController getRestController(
+        UnaryOperator<RestHandler> handlerWrapper,
+        NodeClient client,
+        CircuitBreakerService circuitBreakerService,
+        UsageService usageService,
+        Tracer tracer
+    ) {
+        return new CustomController(handlerWrapper, client, circuitBreakerService, usageService, tracer);
+    }
+
+}

qa/custom-rest-controller/src/javaRestTest/java/org/elasticsearch/plugins/interceptor/CustomRestPluginIT.java

@@ -0,0 +1,64 @@
+/*
+ * 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.plugins.interceptor;
+
+import co.elastic.elasticsearch.test.CustomRestPlugin;
+
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.RequestOptions;
+import org.elasticsearch.client.Response;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.test.ESIntegTestCase;
+
+import java.io.IOException;
+import java.util.Collection;
+
+import static org.elasticsearch.common.util.CollectionUtils.appendToCopyNoNullElements;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.notNullValue;
+
+public class CustomRestPluginIT extends ESIntegTestCase {
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return appendToCopyNoNullElements(super.nodePlugins(), CustomRestPlugin.class);
+    }
+
+    @Override
+    protected boolean addMockHttpTransport() {
+        return false; // enable http
+    }
+
+    public void testInterceptor() throws Exception {
+        var headerValue = randomAlphaOfLengthBetween(4, 12);
+        assertThat(doRequest("x-test-interceptor", headerValue), equalTo(headerValue));
+        assertThat(doRequest("x-test-interceptor", null), equalTo(null));
+    }
+
+    public void testController() throws Exception {
+        var headerValue = randomAlphaOfLengthBetween(4, 12);
+        assertThat(doRequest("x-test-controller", headerValue), equalTo(headerValue));
+        assertThat(doRequest("x-test-controller", null), equalTo(null));
+    }
+
+    private String doRequest(String headerName, String headerValue) throws IOException {
+        assertThat(headerName, notNullValue());
+
+        var client = getRestClient();
+        RequestOptions.Builder options = RequestOptions.DEFAULT.toBuilder();
+        if (headerValue != null) {
+            options.addHeader(headerName, headerValue);
+        }
+        var request = new Request("GET", "/_nodes/_local/plugins");
+        request.setOptions(options);
+
+        final Response response = client.performRequest(request);
+        return response.getHeader(headerName);
+    }
+}

server/src/main/java/module-info.java

@@ -295,7 +295,7 @@ module org.elasticsearch.server {
     exports org.elasticsearch.persistent;
     exports org.elasticsearch.persistent.decider;
     exports org.elasticsearch.plugins;
-    exports org.elasticsearch.plugins.interceptor to org.elasticsearch.security;
+    exports org.elasticsearch.plugins.interceptor to org.elasticsearch.security, org.elasticsearch.serverless.rest;
     exports org.elasticsearch.plugins.spi;
     exports org.elasticsearch.repositories;
     exports org.elasticsearch.repositories.blobstore;

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

@@ -305,7 +305,7 @@ import org.elasticsearch.persistent.StartPersistentTaskAction;
 import org.elasticsearch.persistent.UpdatePersistentTaskStatusAction;
 import org.elasticsearch.plugins.ActionPlugin;
 import org.elasticsearch.plugins.ActionPlugin.ActionHandler;
-import org.elasticsearch.plugins.interceptor.RestInterceptorActionPlugin;
+import org.elasticsearch.plugins.interceptor.RestServerActionPlugin;
 import org.elasticsearch.reservedstate.ReservedClusterStateHandler;
 import org.elasticsearch.reservedstate.service.ReservedClusterStateService;
 import org.elasticsearch.rest.RestController;
@@ -469,6 +469,7 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
@@ -545,37 +546,68 @@ public class ActionModule extends AbstractModule {
                 new RestHeaderDefinition(Task.X_ELASTIC_PRODUCT_ORIGIN_HTTP_HEADER, false)
             )
         ).collect(Collectors.toSet());
-        UnaryOperator<RestHandler> restInterceptor = null;
+        UnaryOperator<RestHandler> restInterceptor = getRestServerComponent(
+            "REST interceptor",
+            actionPlugins,
+            restPlugin -> restPlugin.getRestHandlerInterceptor(threadPool.getThreadContext())
+        );
+        mappingRequestValidators = new RequestValidators<>(
+            actionPlugins.stream().flatMap(p -> p.mappingRequestValidators().stream()).toList()
+        );
+        indicesAliasesRequestRequestValidators = new RequestValidators<>(
+            actionPlugins.stream().flatMap(p -> p.indicesAliasesRequestValidators().stream()).toList()
+        );
+        headersToCopy = headers;
+
+        var customController = getRestServerComponent(
+            "REST controller",
+            actionPlugins,
+            restPlugin -> restPlugin.getRestController(restInterceptor, nodeClient, circuitBreakerService, usageService, tracer)
+        );
+        if (customController != null) {
+            restController = customController;
+        } else {
+            restController = new RestController(restInterceptor, nodeClient, circuitBreakerService, usageService, tracer);
+        }
+        reservedClusterStateService = new ReservedClusterStateService(clusterService, reservedStateHandlers);
+    }
+
+    private static <T> T getRestServerComponent(
+        String type,
+        List<ActionPlugin> actionPlugins,
+        Function<RestServerActionPlugin, T> function
+    ) {
+        T result = null;
         for (ActionPlugin plugin : actionPlugins) {
-            if (plugin instanceof RestInterceptorActionPlugin riplugin) {
-                UnaryOperator<RestHandler> newRestInterceptor = riplugin.getRestHandlerInterceptor(threadPool.getThreadContext());
-                if (newRestInterceptor != null) {
-                    logger.debug("Using REST interceptor from plugin " + plugin.getClass().getName());
-                    if (plugin.getClass().getCanonicalName() == null
-                        || plugin.getClass().getCanonicalName().startsWith("org.elasticsearch.xpack") == false) {
+            if (plugin instanceof RestServerActionPlugin restPlugin) {
+                var newInstance = function.apply(restPlugin);
+                if (newInstance != null) {
+                    logger.debug("Using custom {} from plugin {}", type, plugin.getClass().getName());
+                    if (isInternalPlugin(plugin) == false) {
                         throw new IllegalArgumentException(
                             "The "
                                 + plugin.getClass().getName()
-                                + " plugin tried to install a custom REST "
-                                + "interceptor. This functionality is not available anymore."
+                                + " plugin tried to install a custom "
+                                + type
+                                + ". This functionality is not available to external plugins."
                         );
                     }
-                    if (restInterceptor != null) {
-                        throw new IllegalArgumentException("Cannot have more than one plugin implementing a REST interceptor");
+                    if (result != null) {
+                        throw new IllegalArgumentException("Cannot have more than one plugin implementing a " + type);
                     }
-                    restInterceptor = newRestInterceptor;
+                    result = newInstance;
                 }
             }
         }
-        mappingRequestValidators = new RequestValidators<>(
-            actionPlugins.stream().flatMap(p -> p.mappingRequestValidators().stream()).toList()
-        );
-        indicesAliasesRequestRequestValidators = new RequestValidators<>(
-            actionPlugins.stream().flatMap(p -> p.indicesAliasesRequestValidators().stream()).toList()
-        );
-        headersToCopy = headers;
-        restController = new RestController(restInterceptor, nodeClient, circuitBreakerService, usageService, tracer);
-        reservedClusterStateService = new ReservedClusterStateService(clusterService, reservedStateHandlers);
+        return result;
+    }
+
+    private static boolean isInternalPlugin(ActionPlugin plugin) {
+        final String canonicalName = plugin.getClass().getCanonicalName();
+        if (canonicalName == null) {
+            return false;
+        }
+        return canonicalName.startsWith("org.elasticsearch.xpack.") || canonicalName.startsWith("co.elastic.elasticsearch.");
     }
 
     /**

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

@@ -8,16 +8,22 @@
 
 package org.elasticsearch.plugins.interceptor;
 
+import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.indices.breaker.CircuitBreakerService;
 import org.elasticsearch.plugins.ActionPlugin;
+import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestHandler;
+import org.elasticsearch.tracing.Tracer;
+import org.elasticsearch.usage.UsageService;
 
 import java.util.function.UnaryOperator;
 
 /**
  * An action plugin that intercepts incoming the REST requests.
  */
-public interface RestInterceptorActionPlugin extends ActionPlugin {
+public interface RestServerActionPlugin extends ActionPlugin {
 
     /**
      * Returns a function used to intercept each rest request before handling the request.
@@ -41,4 +47,19 @@ public interface RestInterceptorActionPlugin extends ActionPlugin {
      * Note: Only one installed plugin may implement a rest interceptor.
      */
     UnaryOperator<RestHandler> 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,
+        NodeClient client,
+        CircuitBreakerService circuitBreakerService,
+        UsageService usageService,
+        Tracer tracer
+    ) {
+        return null;
+    }
 }

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

@@ -23,9 +23,10 @@ import org.elasticsearch.common.settings.SettingsFilter;
 import org.elasticsearch.common.settings.SettingsModule;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.indices.TestIndexNameExpressionResolver;
+import org.elasticsearch.indices.breaker.CircuitBreakerService;
 import org.elasticsearch.plugins.ActionPlugin;
 import org.elasticsearch.plugins.ActionPlugin.ActionHandler;
-import org.elasticsearch.plugins.interceptor.RestInterceptorActionPlugin;
+import org.elasticsearch.plugins.interceptor.RestServerActionPlugin;
 import org.elasticsearch.rest.RestChannel;
 import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestHandler;
@@ -36,6 +37,7 @@ import org.elasticsearch.tasks.TaskManager;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.tracing.Tracer;
 import org.elasticsearch.usage.UsageService;
 import org.hamcrest.Matchers;
 
@@ -259,7 +261,7 @@ public class ActionModuleTests extends ESTestCase {
 
         SettingsModule settingsModule = new SettingsModule(Settings.EMPTY);
         ThreadPool threadPool = new TestThreadPool(getTestName());
-        ActionPlugin secPlugin = new SecPlugin();
+        ActionPlugin secPlugin = new SecPlugin(true, false);
         try {
             UsageService usageService = new UsageService();
 
@@ -286,7 +288,45 @@ public class ActionModuleTests extends ESTestCase {
                 e.getMessage(),
                 Matchers.equalTo(
                     "The org.elasticsearch.action.ActionModuleTests$SecPlugin plugin tried to "
-                        + "install a custom REST interceptor. This functionality is not available anymore."
+                        + "install a custom REST interceptor. This functionality is not available to external plugins."
+                )
+            );
+        } finally {
+            threadPool.shutdown();
+        }
+    }
+
+    public void test3rdPartyRestControllerIsNotInstalled() {
+        SettingsModule settingsModule = new SettingsModule(Settings.EMPTY);
+        ThreadPool threadPool = new TestThreadPool(getTestName());
+        ActionPlugin secPlugin = new SecPlugin(false, true);
+        try {
+            UsageService usageService = new UsageService();
+
+            Exception e = expectThrows(
+                IllegalArgumentException.class,
+                () -> new ActionModule(
+                    settingsModule.getSettings(),
+                    TestIndexNameExpressionResolver.newInstance(threadPool.getThreadContext()),
+                    settingsModule.getIndexScopedSettings(),
+                    settingsModule.getClusterSettings(),
+                    settingsModule.getSettingsFilter(),
+                    threadPool,
+                    Arrays.asList(secPlugin),
+                    null,
+                    null,
+                    usageService,
+                    null,
+                    null,
+                    mock(ClusterService.class),
+                    List.of()
+                )
+            );
+            assertThat(
+                e.getMessage(),
+                Matchers.equalTo(
+                    "The org.elasticsearch.action.ActionModuleTests$SecPlugin plugin tried to install a custom REST controller."
+                        + " This functionality is not available to external plugins."
                 )
             );
         } finally {
@@ -304,10 +344,37 @@ public class ActionModuleTests extends ESTestCase {
         public void handleRequest(RestRequest request, RestChannel channel, NodeClient client) throws Exception {}
     }
 
-    class SecPlugin implements ActionPlugin, RestInterceptorActionPlugin {
+    class SecPlugin implements ActionPlugin, RestServerActionPlugin {
+        private final boolean installInterceptor;
+        private final boolean installController;
+
+        SecPlugin(boolean installInterceptor, boolean installController) {
+            this.installInterceptor = installInterceptor;
+            this.installController = installController;
+        }
+
         @Override
         public UnaryOperator<RestHandler> getRestHandlerInterceptor(ThreadContext threadContext) {
-            return UnaryOperator.identity();
+            if (installInterceptor) {
+                return UnaryOperator.identity();
+            } else {
+                return null;
+            }
         }
-    };
+
+        @Override
+        public RestController getRestController(
+            UnaryOperator<RestHandler> handlerWrapper,
+            NodeClient client,
+            CircuitBreakerService circuitBreakerService,
+            UsageService usageService,
+            Tracer tracer
+        ) {
+            if (installController) {
+                return new RestController(handlerWrapper, client, circuitBreakerService, usageService, tracer);
+            } else {
+                return null;
+            }
+        }
+    }
 }

server/src/test/java/org/elasticsearch/module/BasicServerModuleTests.java

@@ -38,9 +38,17 @@ public class BasicServerModuleTests extends ESTestCase {
     public void testQualifiedExports() {
         var md = getServerDescriptor();
 
-        // The package containing the RestInterceptor type, org.elasticsearch.plugins.interceptor,
-        // should only be exported to security.
-        assertThat(md.exports(), hasItem(exportsOf("org.elasticsearch.plugins.interceptor", Set.of("org.elasticsearch.security"))));
+        // The package containing the RestServerActionPlugin (RestInterceptor) type, org.elasticsearch.plugins.interceptor,
+        // should only be exported to security or serverless (rest controller)
+        assertThat(
+            md.exports(),
+            hasItem(
+                exportsOf(
+                    "org.elasticsearch.plugins.interceptor",
+                    Set.of("org.elasticsearch.security", "org.elasticsearch.serverless.rest")
+                )
+            )
+        );
 
         // additional qualified export constraint go here
     }

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

@@ -80,7 +80,7 @@ import org.elasticsearch.plugins.ScriptPlugin;
 import org.elasticsearch.plugins.SearchPlugin;
 import org.elasticsearch.plugins.ShutdownAwarePlugin;
 import org.elasticsearch.plugins.SystemIndexPlugin;
-import org.elasticsearch.plugins.interceptor.RestInterceptorActionPlugin;
+import org.elasticsearch.plugins.interceptor.RestServerActionPlugin;
 import org.elasticsearch.repositories.RepositoriesService;
 import org.elasticsearch.repositories.Repository;
 import org.elasticsearch.rest.RestController;
@@ -138,7 +138,7 @@ public class LocalStateCompositeXPackPlugin extends XPackPlugin
         SystemIndexPlugin,
         SearchPlugin,
         ShutdownAwarePlugin,
-        RestInterceptorActionPlugin {
+        RestServerActionPlugin {
 
     private XPackLicenseState licenseState;
     private SSLService sslService;
@@ -441,8 +441,8 @@ public class LocalStateCompositeXPackPlugin extends XPackPlugin
 
         // There can be only one.
         List<UnaryOperator<RestHandler>> items = filterPlugins(ActionPlugin.class).stream()
-            .filter(RestInterceptorActionPlugin.class::isInstance)
-            .map(RestInterceptorActionPlugin.class::cast)
+            .filter(RestServerActionPlugin.class::isInstance)
+            .map(RestServerActionPlugin.class::cast)
             .map(p -> p.getRestHandlerInterceptor(threadContext))
             .filter(Objects::nonNull)
             .toList();

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/Security.java

@@ -79,7 +79,7 @@ import org.elasticsearch.plugins.NetworkPlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.plugins.SearchPlugin;
 import org.elasticsearch.plugins.SystemIndexPlugin;
-import org.elasticsearch.plugins.interceptor.RestInterceptorActionPlugin;
+import org.elasticsearch.plugins.interceptor.RestServerActionPlugin;
 import org.elasticsearch.repositories.RepositoriesService;
 import org.elasticsearch.reservedstate.ReservedClusterStateHandler;
 import org.elasticsearch.rest.RestController;
@@ -419,7 +419,7 @@ public class Security extends Plugin
         MapperPlugin,
         ExtensiblePlugin,
         SearchPlugin,
-        RestInterceptorActionPlugin {
+        RestServerActionPlugin {
 
     public static final String SECURITY_CRYPTO_THREAD_POOL_NAME = XPackField.SECURITY + "-crypto";
 

x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/SecurityTests.java

@@ -756,7 +756,7 @@ public class SecurityTests extends ESTestCase {
                     "Security rest interceptor",
                     ActionModule.class.getName(),
                     Level.DEBUG,
-                    "Using REST interceptor from plugin org.elasticsearch.xpack.security.Security"
+                    "Using custom REST interceptor from plugin org.elasticsearch.xpack.security.Security"
                 )
             );