Use auth tokens passed from core and introduce extension and user REST clients

src/main/java/org/opensearch/sdk/ExtensionSettings.java

@@ -21,6 +21,7 @@
 
 import org.yaml.snakeyaml.Yaml;
 
+import static org.opensearch.sdk.ssl.SSLConfigConstants.SSL_HTTP_ENABLED;
 import static org.opensearch.sdk.ssl.SSLConfigConstants.SSL_TRANSPORT_CLIENT_PEMCERT_FILEPATH;
 import static org.opensearch.sdk.ssl.SSLConfigConstants.SSL_TRANSPORT_CLIENT_PEMKEY_FILEPATH;
 import static org.opensearch.sdk.ssl.SSLConfigConstants.SSL_TRANSPORT_CLIENT_PEMTRUSTEDCAS_FILEPATH;
@@ -62,6 +63,7 @@ public class ExtensionSettings {
      */
     public static final Set<String> SECURITY_SETTINGS_KEYS = Set.of(
         "path.home", // TODO Find the right place to put this setting
+        SSL_HTTP_ENABLED,
         SSL_TRANSPORT_CLIENT_PEMCERT_FILEPATH,
         SSL_TRANSPORT_CLIENT_PEMKEY_FILEPATH,
         SSL_TRANSPORT_CLIENT_PEMTRUSTEDCAS_FILEPATH,

src/main/java/org/opensearch/sdk/ExtensionsRunner.java

@@ -17,8 +17,10 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.action.ActionType;
 import org.opensearch.action.support.TransportAction;
+import org.opensearch.client.opensearch.OpenSearchClient;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
 import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
+import org.opensearch.discovery.InitializeExtensionSecurityRequest;
 import org.opensearch.extensions.rest.ExtensionRestRequest;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Settings;
@@ -125,6 +127,7 @@ public class ExtensionsRunner {
     private final SDKNamedXContentRegistry sdkNamedXContentRegistry;
     private final SDKNamedWriteableRegistry sdkNamedWriteableRegistry;
     private final SDKClient sdkClient;
+    private OpenSearchClient extensionRestClient;
     private final SDKClusterService sdkClusterService;
     private final SDKTransportService sdkTransportService;
     private final SDKActionModule sdkActionModule;
@@ -344,6 +347,19 @@ public void setExtensionNode(DiscoveryExtensionNode extensionNode) {
         this.extensionNode = extensionNode;
     }
 
+    /**
+     * Initializes a REST Client for this extension to interact with an OpenSearch cluster on its own behalf
+     *
+     * @param serviceAccountToken Access token that permits an extension to make requests on its own behalf.
+     *                            Common examples of usages of service account tokens include interacting with
+     *                            an extension's reserved indices.
+     */
+    public void initializeExtensionRestClient(String serviceAccountToken) {
+        OpenSearchClient restClient = getSdkClient()
+                .initializeJavaClientWithHeaders(Map.of("Authorization", "Bearer " + serviceAccountToken));
+        this.extensionRestClient = restClient;
+    }
+
     /**
      * Returns the discovery extension node set during extension initialization
      *
@@ -403,6 +419,15 @@ public void startTransportService(TransportService transportService) {
             (request, channel, task) -> channel.sendResponse(extensionsInitRequestHandler.handleExtensionInitRequest(request))
         );
 
+        transportService.registerRequestHandler(
+            ExtensionsManager.REQUEST_EXTENSION_REGISTER_SECURITY_SETTINGS,
+            ThreadPool.Names.GENERIC,
+            false,
+            false,
+            InitializeExtensionSecurityRequest::new,
+            (request, channel, task) -> channel.sendResponse(extensionsInitRequestHandler.handleExtensionSecurityInitRequest(request))
+        );
+
         transportService.registerRequestHandler(
             ExtensionsManager.REQUEST_REST_EXECUTE_ON_EXTENSION_ACTION,
             ThreadPool.Names.GENERIC,
@@ -528,6 +553,15 @@ public SDKClient getSdkClient() {
         return sdkClient;
     }
 
+    /**
+     * Returns the Extension rest client instance used by this extension.
+     *
+     * @return The Extension rest client instance.
+     */
+    public OpenSearchClient getExtensionRestClient() {
+        return extensionRestClient;
+    }
+
     /**
      * @return The SDKClusterService instance associated with this object.
      */

src/main/java/org/opensearch/sdk/SDKClient.java

@@ -23,11 +23,13 @@
 import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
 import com.fasterxml.jackson.datatype.guava.GuavaModule;
 import org.apache.hc.core5.function.Factory;
+import org.apache.hc.core5.http.Header;
 import org.apache.hc.core5.http.HttpHost;
 import org.apache.hc.client5.http.ssl.NoopHostnameVerifier;
 import org.apache.hc.client5.http.ssl.ClientTlsStrategyBuilder;
 import org.apache.hc.client5.http.impl.nio.PoolingAsyncClientConnectionManager;
 import org.apache.hc.client5.http.impl.nio.PoolingAsyncClientConnectionManagerBuilder;
+import org.apache.hc.core5.http.message.BasicHeader;
 import org.apache.hc.core5.http.nio.ssl.TlsStrategy;
 import org.apache.hc.core5.reactor.ssl.TlsDetails;
 import org.apache.hc.core5.ssl.SSLContextBuilder;
@@ -95,6 +97,9 @@
 
 import javax.net.ssl.SSLEngine;
 
+import static org.opensearch.sdk.ssl.SSLConfigConstants.SSL_HTTP_ENABLED;
+import static org.opensearch.sdk.ssl.SSLConfigConstants.SSL_TRANSPORT_ENABLED;
+
 /**
  * This class creates SDKClient for an extension to make requests to OpenSearch
  */
@@ -164,8 +169,11 @@ public void updateOpenSearchNodeSettings(String address, String httpPort) {
      * @param port The port the client should connect to
      * @return An instance of the builder
      */
-    private static RestClientBuilder builder(String hostAddress, int port) {
-        RestClientBuilder builder = RestClient.builder(new HttpHost(hostAddress, port));
+    private static RestClientBuilder builder(String hostAddress, int port, ExtensionSettings extensionSettings) {
+        boolean httpsEnabled = extensionSettings.getSecuritySettings().containsKey(SSL_HTTP_ENABLED)
+                && "true".equals(extensionSettings.getSecuritySettings().get(SSL_HTTP_ENABLED));
+        String scheme = httpsEnabled ? "https" : "http";
+        RestClientBuilder builder = RestClient.builder(new HttpHost(scheme, hostAddress, port));
         builder.setStrictDeprecationMode(true);
         builder.setHttpClientConfigCallback(httpClientBuilder -> {
             try {
@@ -201,8 +209,9 @@ public TlsDetails create(final SSLEngine sslEngine) {
      * @param port The port of OpenSearch cluster
      * @return The OpenSearchTransport implementation of RestClientTransport.
      */
-    private OpenSearchTransport initializeTransport(String hostAddress, int port) {
-        RestClientBuilder builder = builder(hostAddress, port);
+    private OpenSearchTransport initializeTransport(String hostAddress, int port, Map<String, String> headers) {
+        RestClientBuilder builder = builder(hostAddress, port, extensionSettings);
+        builder.setDefaultHeaders(headers.keySet().stream().map(k -> new BasicHeader(k, headers.get(k))).toArray(Header[]::new));
 
         restClient = builder.build();
         ObjectMapper mapper = new ObjectMapper();
@@ -227,6 +236,20 @@ public OpenSearchClient initializeJavaClient() {
         return initializeJavaClient(extensionSettings.getOpensearchAddress(), Integer.parseInt(extensionSettings.getOpensearchPort()));
     }
 
+    /**
+     * Initializes an OpenSearchClient using OpenSearch JavaClient
+     *
+     * @return The SDKClient implementation of OpenSearchClient. The user is responsible for calling
+     *         {@link #doCloseJavaClients()} when finished with the client
+     */
+    public OpenSearchClient initializeJavaClientWithHeaders(Map<String, String> headers) {
+        return initializeJavaClientWithHeaders(
+            extensionSettings.getOpensearchAddress(),
+            Integer.parseInt(extensionSettings.getOpensearchPort()),
+            headers
+        );
+    }
+
     /**
      * Initializes an OpenSearchClient using OpenSearch JavaClient
      *
@@ -236,7 +259,21 @@ public OpenSearchClient initializeJavaClient() {
      *         {@link #doCloseJavaClients()} when finished with the client
      */
     public OpenSearchClient initializeJavaClient(String hostAddress, int port) {
-        OpenSearchTransport transport = initializeTransport(hostAddress, port);
+        OpenSearchTransport transport = initializeTransport(hostAddress, port, Map.of());
+        javaClient = new OpenSearchClient(transport);
+        return javaClient;
+    }
+
+    /**
+     * Initializes an OpenSearchClient using OpenSearch JavaClient
+     *
+     * @param hostAddress The address of OpenSearch cluster, client can connect to
+     * @param port The port of OpenSearch cluster
+     * @return The SDKClient implementation of OpenSearchClient. The user is responsible for calling
+     *         {@link #doCloseJavaClients()} when finished with the client
+     */
+    public OpenSearchClient initializeJavaClientWithHeaders(String hostAddress, int port, Map<String, String> headers) {
+        OpenSearchTransport transport = initializeTransport(hostAddress, port, headers);
         javaClient = new OpenSearchClient(transport);
         return javaClient;
     }
@@ -260,7 +297,7 @@ public OpenSearchAsyncClient initializeJavaAsyncClient() {
      *         {@link #doCloseJavaClients()} when finished with the client
      */
     public OpenSearchAsyncClient initalizeJavaAsyncClient(String hostAddress, int port) {
-        OpenSearchTransport transport = initializeTransport(hostAddress, port);
+        OpenSearchTransport transport = initializeTransport(hostAddress, port, Map.of());
         javaAsyncClient = new OpenSearchAsyncClient(transport);
         return javaAsyncClient;
     }
@@ -300,7 +337,7 @@ public SDKRestClient initializeRestClient() {
      */
     @Deprecated
     public SDKRestClient initializeRestClient(String hostAddress, int port) {
-        this.sdkRestClient = new SDKRestClient(this, new RestHighLevelClient(builder(hostAddress, port)));
+        this.sdkRestClient = new SDKRestClient(this, new RestHighLevelClient(builder(hostAddress, port, extensionSettings)));
         return this.sdkRestClient;
     }
 
@@ -556,7 +593,7 @@ public void bulk(BulkRequest request, ActionListener<BulkResponse> listener) {
          * @return the response returned by OpenSearch
          * @throws IOException in case of a problem or the connection was aborted
          */
-        public Response performRequest(Request request) throws IOException {
+        public Response uest(Request request) throws IOException {
             return restHighLevelClient.getLowLevelClient().performRequest(request);
         }
 

src/main/java/org/opensearch/sdk/handlers/ExtensionsInitRequestHandler.java

@@ -9,16 +9,32 @@
 
 package org.opensearch.sdk.handlers;
 
+import com.fasterxml.jackson.databind.JsonNode;
 import org.apache.logging.log4j.LogManager;
 
 import org.apache.logging.log4j.Logger;
+import org.opensearch.client.RequestOptions;
+import org.opensearch.client.WarningFailureException;
+import org.opensearch.client.opensearch.OpenSearchClient;
+import org.opensearch.client.opensearch.core.IndexRequest;
+import org.opensearch.client.opensearch.core.SearchRequest;
+import org.opensearch.client.opensearch.core.SearchResponse;
+import org.opensearch.client.opensearch.indices.CreateIndexRequest;
+import org.opensearch.client.opensearch.indices.DeleteIndexRequest;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.discovery.InitializeExtensionRequest;
 import org.opensearch.discovery.InitializeExtensionResponse;
+import org.opensearch.discovery.InitializeExtensionSecurityRequest;
+import org.opensearch.discovery.InitializeExtensionSecurityResponse;
 import org.opensearch.sdk.ExtensionsRunner;
 import org.opensearch.sdk.SDKTransportService;
 import org.opensearch.transport.TransportService;
 
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+import java.util.Map;
+
 import static org.opensearch.sdk.ExtensionsRunner.NODE_NAME_SETTING;
 
 /**
@@ -94,4 +110,20 @@ public InitializeExtensionResponse handleExtensionInitRequest(InitializeExtensio
             extensionsRunner.getSdkClusterService().getClusterSettings().sendPendingSettingsUpdateConsumers();
         }
     }
+
+    /**
+     * Handles a extension request from OpenSearch. This is the first request for the transport communication and will initialize the extension and will be a part of OpenSearch bootstrap.
+     *
+     * @param extensionInitSecurityRequest  The request to handle.
+     * @return A response to OpenSearch validating that this is an extension.
+     */
+    public InitializeExtensionSecurityResponse handleExtensionSecurityInitRequest(
+        InitializeExtensionSecurityRequest extensionInitSecurityRequest
+    ) {
+        logger.info("Registering Extension Request received from OpenSearch");
+
+        extensionsRunner.initializeExtensionRestClient(extensionInitSecurityRequest.getServiceAccountToken());
+
+        return new InitializeExtensionSecurityResponse(extensionsRunner.getExtensionNode().getId());
+    }
 }

src/main/java/org/opensearch/sdk/handlers/ExtensionsRestRequestHandler.java

@@ -9,6 +9,7 @@
 
 package org.opensearch.sdk.handlers;
 
+import joptsimple.internal.Strings;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.core.common.bytes.BytesReference;
@@ -22,6 +23,10 @@
 import org.opensearch.sdk.rest.SDKHttpRequest;
 import org.opensearch.sdk.rest.SDKRestRequest;
 
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.emptyList;
@@ -68,11 +73,19 @@ public RestExecuteOnExtensionResponse handleRestExecuteOnExtensionRequest(Extens
             );
         }
 
+        String oboToken = request.getRequestIssuerIdentity();
+        Map<String, List<String>> headers = new HashMap<>();
+        headers.putAll(request.headers());
+        System.out.println("oboToken: " + oboToken);
+        if (!Strings.isNullOrEmpty(oboToken)) {
+            headers.put("Authorization", List.of("Bearer " + oboToken));
+        }
+
         SDKRestRequest sdkRestRequest = new SDKRestRequest(
             sdkNamedXContentRegistry.getRegistry(),
             request.params(),
             request.path(),
-            request.headers(),
+            headers,
             new SDKHttpRequest(request),
             null
         );

src/main/java/org/opensearch/sdk/rest/BaseExtensionRestHandler.java

@@ -11,7 +11,9 @@
 
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 
@@ -25,6 +27,7 @@
 import static org.apache.hc.core5.http.ContentType.APPLICATION_JSON;
 
 import org.opensearch.OpenSearchException;
+import org.opensearch.client.opensearch.OpenSearchClient;
 import org.opensearch.common.logging.DeprecationLogger;
 import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.core.common.Strings;
@@ -38,6 +41,7 @@
 import org.opensearch.rest.RestRequest.Method;
 import org.opensearch.rest.RestResponse;
 import org.opensearch.core.rest.RestStatus;
+import org.opensearch.sdk.SDKClient;
 
 /**
  * Provides convenience methods to reduce boilerplate code in an {@link ExtensionRestHandler} implementation.
@@ -48,6 +52,14 @@ public abstract class BaseExtensionRestHandler implements ExtensionRestHandler {
 
     private String routeNamePrefix;
 
+    private SDKClient sdkClient;
+
+    protected OpenSearchClient userRestClient;
+
+    public BaseExtensionRestHandler(SDKClient sdkClient) {
+        this.sdkClient = sdkClient;
+    }
+
     /**
      * Constant for JSON content type
      */
@@ -114,6 +126,15 @@ public List<ReplacedRoute> replacedRoutes() {
 
     @Override
     public ExtensionRestResponse handleRequest(RestRequest request) {
+        if (request instanceof SDKRestRequest) {
+            SDKRestRequest sdkRestRequest = (SDKRestRequest) request;
+            List<String> authorizationHeaders = sdkRestRequest.getHeaders().get("Authorization");
+            Map<String, String> headers = new HashMap<>();
+            if (!authorizationHeaders.isEmpty()) {
+                headers.put("Authorization", authorizationHeaders.get(0));
+            }
+            this.userRestClient = sdkClient.initializeJavaClientWithHeaders(headers);
+        }
         Optional<NamedRoute> route = routes().stream()
             .filter(rh -> rh.getMethod().equals(request.method()))
             .filter(rh -> restPathMatches(request.path(), rh.getPath()))

src/main/java/org/opensearch/sdk/sample/helloworld/HelloWorldExtension.java

@@ -56,7 +56,7 @@ public HelloWorldExtension() {
 
     @Override
     public List<ExtensionRestHandler> getExtensionRestHandlers() {
-        return List.of(new RestHelloAction(), new RestRemoteHelloAction(extensionsRunner()));
+        return List.of(new RestHelloAction(extensionsRunner().getSdkClient()), new RestRemoteHelloAction(extensionsRunner()));
     }
 
     @Override

src/main/java/org/opensearch/sdk/sample/helloworld/rest/RestHelloAction.java

@@ -18,6 +18,7 @@
 import org.opensearch.rest.NamedRoute;
 import org.opensearch.rest.RestRequest;
 import org.opensearch.rest.RestResponse;
+import org.opensearch.sdk.SDKClient;
 import org.opensearch.sdk.rest.BaseExtensionRestHandler;
 import org.opensearch.sdk.rest.ExtensionRestHandler;
 
@@ -48,7 +49,7 @@
 public class RestHelloAction extends BaseExtensionRestHandler {
 
     private static final String TEXT_CONTENT_TYPE = "text/plain; charset=UTF-8";
-    private static final String GREETING = "Hello, %s!";
+    public static final String GREETING = "Hello, %s!";
     private static final String DEFAULT_NAME = "World";
 
     private String worldName = DEFAULT_NAME;
@@ -58,7 +59,9 @@ public class RestHelloAction extends BaseExtensionRestHandler {
     /**
      * Instantiate this action
      */
-    public RestHelloAction() {}
+    public RestHelloAction(SDKClient sdkClient) {
+        super(sdkClient);
+    }
 
     @Override
     public List<NamedRoute> routes() {