Improve serialization speeds

bwc-test/src/test/java/org/opensearch/security/bwc/SecurityBackwardsCompatibilityIT.java

@@ -31,14 +31,18 @@
 import org.apache.hc.core5.http.nio.ssl.TlsStrategy;
 import org.apache.hc.core5.reactor.ssl.TlsDetails;
 import org.apache.hc.core5.ssl.SSLContextBuilder;
+import org.junit.AfterClass;
+import org.junit.Assert;
 import org.junit.Assume;
 import org.junit.Before;
 import org.opensearch.client.Response;
 import org.opensearch.client.ResponseException;
 import org.opensearch.client.RestClient;
 import org.opensearch.client.RestClientBuilder;
+import org.opensearch.common.Randomness;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.common.util.io.IOUtils;
 import org.opensearch.security.bwc.helper.RestHelper;
 import org.opensearch.test.rest.OpenSearchRestTestCase;
 import org.opensearch.Version;
@@ -56,13 +60,22 @@ public class SecurityBackwardsCompatibilityIT extends OpenSearchRestTestCase {
     private final String TEST_USER = "user";
     private final String TEST_PASSWORD = "290735c0-355d-4aaf-9b42-1aaa1f2a3cee";
     private final String TEST_ROLE = "test-dls-fls-role";
+    private static RestClient testUserRestClient = null;
 
     @Before
     public void testSetup() {
         final String bwcsuiteString = System.getProperty("tests.rest.bwcsuite");
         Assume.assumeTrue("Test cannot be run outside the BWC gradle task 'bwcTestSuite' or its dependent tasks", bwcsuiteString != null);
         CLUSTER_TYPE = ClusterType.parse(bwcsuiteString);
         CLUSTER_NAME = System.getProperty("tests.clustername");
+        if (testUserRestClient == null) {
+            testUserRestClient = buildClient(
+                super.restClientSettings(),
+                super.getClusterHosts().toArray(new HttpHost[0]),
+                TEST_USER,
+                TEST_PASSWORD
+            );
+        }
     }
 
     @Override
@@ -101,18 +114,21 @@ protected final Settings restClientSettings() {
             .build();
     }
 
+    protected RestClient buildClient(Settings settings, HttpHost[] hosts, String username, String password) {
+        RestClientBuilder builder = RestClient.builder(hosts);
+        configureHttpsClient(builder, settings, username, password);
+        boolean strictDeprecationMode = settings.getAsBoolean("strictDeprecationMode", true);
+        builder.setStrictDeprecationMode(strictDeprecationMode);
+        return builder.build();
+    }
+
     @Override
     protected RestClient buildClient(Settings settings, HttpHost[] hosts) {
         String username = Optional.ofNullable(System.getProperty("tests.opensearch.username"))
             .orElseThrow(() -> new RuntimeException("user name is missing"));
         String password = Optional.ofNullable(System.getProperty("tests.opensearch.password"))
             .orElseThrow(() -> new RuntimeException("password is missing"));
-
-        RestClientBuilder builder = RestClient.builder(hosts);
-        configureHttpsClient(builder, settings, username, password);
-        boolean strictDeprecationMode = settings.getAsBoolean("strictDeprecationMode", true);
-        builder.setStrictDeprecationMode(strictDeprecationMode);
-        return builder.build();
+        return buildClient(super.restClientSettings(), super.getClusterHosts().toArray(new HttpHost[0]), username, password);
     }
 
     private static void configureHttpsClient(RestClientBuilder builder, Settings settings, String userName, String password) {
@@ -180,6 +196,11 @@ public void testDataIngestionAndSearchBackwardsCompatibility() throws Exception
         searchMatchAll(index);
     }
 
+    public void testNodeStats() throws IOException {
+        List<Response> responses = RestHelper.requestAgainstAllNodes(client(), "GET", "_nodes/stats", null);
+        responses.forEach(r -> Assert.assertEquals(200, r.getStatusLine().getStatusCode()));
+    }
+
     @SuppressWarnings("unchecked")
     private void assertPluginUpgrade(String uri) throws Exception {
         Map<String, Map<String, Object>> responseMap = (Map<String, Map<String, Object>>) getAsMap(uri).get("nodes");
@@ -205,19 +226,26 @@ private void assertPluginUpgrade(String uri) throws Exception {
     private void ingestData(String index) throws IOException {
         StringBuilder bulkRequestBody = new StringBuilder();
         ObjectMapper objectMapper = new ObjectMapper();
-        for (Song song : Song.SONGS) {
-            Map<String, Map<String, String>> indexRequest = new HashMap<>();
-            indexRequest.put("index", new HashMap<>() {
-                {
-                    put("_index", index);
-                }
-            });
-            bulkRequestBody.append(objectMapper.writeValueAsString(indexRequest) + "\n");
-            bulkRequestBody.append(objectMapper.writeValueAsString(song.asJson()) + "\n");
+        int numberOfRequests = Randomness.get().nextInt(10);
+        while (numberOfRequests-- > 0) {
+            for (int i = 0; i < Randomness.get().nextInt(100); i++) {
+                Map<String, Map<String, String>> indexRequest = new HashMap<>();
+                indexRequest.put("index", new HashMap<>() {
+                    {
+                        put("_index", index);
+                    }
+                });
+                bulkRequestBody.append(objectMapper.writeValueAsString(indexRequest) + "\n");
+                bulkRequestBody.append(objectMapper.writeValueAsString(Song.randomSong().asJson()) + "\n");
+            }
+            List<Response> responses = RestHelper.requestAgainstAllNodes(
+                testUserRestClient,
+                "POST",
+                "_bulk?refresh=wait_for",
+                RestHelper.toHttpEntity(bulkRequestBody.toString())
+            );
+            responses.forEach(r -> assertEquals(200, r.getStatusLine().getStatusCode()));
         }
-
-        Response response = RestHelper.makeRequest(client(), "POST", "_bulk", RestHelper.toHttpEntity(bulkRequestBody.toString()));
-        assertThat(response.getStatusLine().getStatusCode(), equalTo(200));
     }
 
     /**
@@ -226,10 +254,16 @@ private void ingestData(String index) throws IOException {
      */
     private void searchMatchAll(String index) throws IOException {
         String matchAllQuery = "{\n" + "    \"query\": {\n" + "        \"match_all\": {}\n" + "    }\n" + "}";
-
-        Response response = RestHelper.makeRequest(client(), "POST", index + "/_search", RestHelper.toHttpEntity(matchAllQuery));
-
-        assertThat(response.getStatusLine().getStatusCode(), equalTo(200));
+        int numberOfRequests = Randomness.get().nextInt(10);
+        while (numberOfRequests-- > 0) {
+            List<Response> responses = RestHelper.requestAgainstAllNodes(
+                testUserRestClient,
+                "POST",
+                index + "/_search",
+                RestHelper.toHttpEntity(matchAllQuery)
+            );
+            responses.forEach(r -> assertEquals(200, r.getStatusLine().getStatusCode()));
+        }
     }
 
     /**
@@ -324,4 +358,10 @@ private void createUserIfNotExists(String user, String password, String role) th
             assertThat(response.getStatusLine().getStatusCode(), equalTo(201));
         }
     }
+
+    @AfterClass
+    public static void cleanUp() throws IOException {
+        OpenSearchRestTestCase.closeClients();
+        IOUtils.close(testUserRestClient);
+    }
 }

bwc-test/src/test/java/org/opensearch/security/bwc/Song.java

@@ -11,9 +11,11 @@
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import org.opensearch.common.Randomness;
 
 import java.util.Map;
 import java.util.Objects;
+import java.util.UUID;
 
 public class Song {
 
@@ -102,4 +104,14 @@ public Map<String, Object> asMap() {
     public String asJson() throws JsonProcessingException {
         return new ObjectMapper().writeValueAsString(this.asMap());
     }
+
+    public static Song randomSong() {
+        return new Song(
+            UUID.randomUUID().toString(),
+            UUID.randomUUID().toString(),
+            UUID.randomUUID().toString(),
+            Randomness.get().nextInt(5),
+            UUID.randomUUID().toString()
+        );
+    }
 }

bwc-test/src/test/java/org/opensearch/security/bwc/helper/RestHelper.java

@@ -9,6 +9,7 @@
 package org.opensearch.security.bwc.helper;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 
 import org.apache.hc.core5.http.Header;
@@ -63,6 +64,26 @@ public static Response makeRequest(RestClient client, String method, String endp
         return response;
     }
 
+    public static List<Response> requestAgainstAllNodes(RestClient client, String method, String endpoint, HttpEntity entity)
+        throws IOException {
+        return requestAgainstAllNodes(client, method, endpoint, entity, null);
+    }
+
+    public static List<Response> requestAgainstAllNodes(
+        RestClient client,
+        String method,
+        String endpoint,
+        HttpEntity entity,
+        List<Header> headers
+    ) throws IOException {
+        int nodeCount = client.getNodes().size();
+        List<Response> responses = new ArrayList<>();
+        while (nodeCount-- > 0) {
+            responses.add(makeRequest(client, method, endpoint, entity, headers));
+        }
+        return responses;
+    }
+
     public static Header getAuthorizationHeader(String username, String password) {
         return new BasicHeader("Authorization", "Basic " + username + ":" + password);
     }

src/main/java/org/opensearch/security/OpenSearchSecurityPlugin.java

@@ -779,7 +779,6 @@ public <T extends TransportRequest> TransportRequestHandler<T> interceptHandler(
 
                         @Override
                         public void messageReceived(T request, TransportChannel channel, Task task) throws Exception {
-                            threadContext.putTransient(ConfigConstants.USE_JDK_SERIALIZATION, channel.getVersion().before(Version.V_3_0_0));
                             si.getHandler(action, actualHandler).messageReceived(request, channel, task);
                         }
                     };

src/main/java/org/opensearch/security/ssl/transport/SecuritySSLRequestHandler.java

@@ -83,8 +83,14 @@ protected ThreadContext getThreadContext() {
 
     @Override
     public final void messageReceived(T request, TransportChannel channel, Task task) throws Exception {
+
         ThreadContext threadContext = getThreadContext();
 
+        threadContext.putTransient(
+            ConfigConstants.USE_JDK_SERIALIZATION,
+            channel.getVersion().before(ConfigConstants.FIRST_CUSTOM_SERIALIZATION_SUPPORTED_OS_VERSION)
+        );
+
         if (SSLRequestHelper.containsBadHeader(threadContext, "_opendistro_security_ssl_")) {
             final Exception exception = ExceptionUtils.createBadHeaderException();
             channel.sendResponse(exception);

src/main/java/org/opensearch/security/support/Base64CustomHelper.java

@@ -35,7 +35,7 @@
  * Provides support for Serialization/Deserialization of objects of supported classes into/from Base64 encoded stream
  * using the OpenSearch's custom serialization protocol implemented by the StreamInput/StreamOutput classes.
  */
 public class Base64CustomHelper {
 
     private enum CustomSerializationFormat {
 
@@ -58,7 +58,7 @@
                 case 3:
                     return GENERIC;
                 default:
                     throw new IllegalArgumentException(String.format("%d is not a valid id", id));
             }
         }
 
@@ -94,12 +94,12 @@
                     streamOutput.writeGenericValue(object);
                     break;
                 default:
                     throw new IllegalArgumentException(
                         String.format("Could not determine custom serialization mode for class %s", clazz.getName())
                     );
             }
         } catch (final Exception e) {
             throw new OpenSearchException("Instance {} of class {} is not serializable", e, object, object.getClass());
         }
         final byte[] bytes = streamOutput.bytes().toBytesRef().bytes;
         streamOutput.close();
@@ -108,7 +108,7 @@
 
     protected static Serializable deserializeObject(final String string) {
 
-        Preconditions.checkArgument(!Strings.isNullOrEmpty(string), "string must not be null or empty");
+        Preconditions.checkArgument(!Strings.isNullOrEmpty(string), "object must not be null or empty");
         final byte[] bytes = BaseEncoding.base64().decode(string);
         Serializable obj = null;
         try (final BytesStreamInput streamInput = new SafeBytesStreamInput(bytes)) {
@@ -126,12 +126,12 @@
                     obj = (Serializable) streamInput.readGenericValue();
                     break;
                 default:
                     throw new IllegalArgumentException("Could not determine custom deserialization mode");
             }
             prohibitUnsafeClasses(obj.getClass());
             return obj;
         } catch (final Exception e) {
             throw new OpenSearchException(e);
         }
     }
 
@@ -146,10 +146,10 @@
      */
     protected static Integer getWriteableClassID(Class<?> clazz) {
         if (!isWriteable(clazz)) {
             throw new OpenSearchException("clazz should implement Writeable ", clazz);
         }
         if (!writeableClassToIdMap.containsKey(clazz)) {
             throw new OpenSearchException("Writeable clazz not registered ", clazz);
         }
         return writeableClassToIdMap.get(clazz);
     }
@@ -165,7 +165,7 @@
      */
     private static void registerWriteable(Class<? extends Writeable> clazz) {
         if (writeableClassToIdMap.containsKey(clazz)) {
             throw new OpenSearchException("writeable clazz is already registered ", clazz.getName());
         }
         int id = writeableClassToIdMap.size() + 1;
         writeableClassToIdMap.put(clazz, id);

src/main/java/org/opensearch/security/support/Base64Helper.java

@@ -39,7 +39,7 @@
     }
 
     public static Serializable deserializeObject(final String string) {
         return deserializeObject(string, false);
     }
 
     public static Serializable deserializeObject(final String string, final boolean useJDKDeserialization) {
@@ -58,15 +58,15 @@
     public static String ensureJDKSerialized(final String string) {
         Serializable serializable;
         try {
             serializable = Base64Helper.deserializeObject(string, false);
         } catch (Exception e) {
             // We received an exception when de-serializing the given string. It is probably JDK serialized.
             // Try to deserialize using JDK
             Base64Helper.deserializeObject(string, true);
             // Since we could deserialize the object using JDK, the string is already JDK serialized, return as is
             return string;
         }
         // If we see an exception now, we want the caller to see it -
-        return Base64Helper.serializeObject(serializable, false);
+        return Base64Helper.serializeObject(serializable, true);
     }
 }

src/main/java/org/opensearch/security/support/Base64JDKHelper.java

@@ -52,7 +52,7 @@
  * Provides support for Serialization/Deserialization of objects of supported classes into/from Base64 encoded stream
  * using JDK's in-built serialization protocol implemented by the ObjectOutputStream and ObjectInputStream classes.
  */
 public class Base64JDKHelper {
 
     private final static class SafeObjectOutputStream extends ObjectOutputStream {
 
@@ -63,15 +63,15 @@
             SecurityManager sm = System.getSecurityManager();
             if (sm != null) {
                 try {
                     sm.checkPermission(new SpecialPermission());
 
                     AccessController.doPrivileged((PrivilegedAction<Void>) () -> {
                         AccessController.checkPermission(SUBSTITUTION_PERMISSION);
                         return null;
                     });
                 } catch (SecurityException e) {
                     return false;
                 }
             }
             return true;
         }
@@ -79,11 +79,11 @@
         static ObjectOutputStream create(ByteArrayOutputStream out) throws IOException {
             try {
                 return useSafeObjectOutputStream ? new SafeObjectOutputStream(out) : new ObjectOutputStream(out);
             } catch (SecurityException e) {
                 // As we try to create SafeObjectOutputStream only when necessary permissions are granted, we should
                 // not reach here, but if we do, we can still return ObjectOutputStream after resetting ByteArrayOutputStream
                 out.reset();
                 return new ObjectOutputStream(out);
             }
         }
 
@@ -93,7 +93,7 @@
 
             SecurityManager sm = System.getSecurityManager();
             if (sm != null) {
                 sm.checkPermission(new SpecialPermission());
             }
 
             AccessController.doPrivileged((PrivilegedAction<Boolean>) () -> enableReplaceObject(true));
@@ -105,7 +105,7 @@
             if (isSafeClass(clazz)) {
                 return obj;
             }
             throw new IOException("Unauthorized serialization attempt " + clazz.getName());
         }
     }
 
@@ -116,8 +116,8 @@
         final ByteArrayOutputStream bos = new ByteArrayOutputStream();
         try (final ObjectOutputStream out = SafeObjectOutputStream.create(bos)) {
             out.writeObject(object);
         } catch (final Exception e) {
             throw new OpenSearchException("Instance {} of class {} is not serializable", e, object, object.getClass());
         }
         final byte[] bytes = bos.toByteArray();
         return BaseEncoding.base64().encode(bytes);
@@ -125,14 +125,14 @@
 
     public static Serializable deserializeObject(final String string) {
 
-        Preconditions.checkArgument(!Strings.isNullOrEmpty(string), "string must not be null or empty");
+        Preconditions.checkArgument(!Strings.isNullOrEmpty(string), "object must not be null or empty");
 
         final byte[] bytes = BaseEncoding.base64().decode(string);
         final ByteArrayInputStream bis = new ByteArrayInputStream(bytes);
         try (SafeObjectInputStream in = new SafeObjectInputStream(bis)) {
             return (Serializable) in.readObject();
         } catch (final Exception e) {
             throw new OpenSearchException(e);
         }
     }
 
@@ -150,7 +150,7 @@
                 return clazz;
             }
 
             throw new InvalidClassException("Unauthorized deserialization attempt ", clazz.getName());
         }
     }
 }

src/main/java/org/opensearch/security/support/ConfigConstants.java

@@ -35,6 +35,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 
+import org.opensearch.Version;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.security.auditlog.impl.AuditCategory;
 
@@ -325,6 +326,7 @@ public enum RolesMappingResolution {
     public static final String TENANCY_GLOBAL_TENANT_DEFAULT_NAME = "";
 
     public static final String USE_JDK_SERIALIZATION = "plugins.security.use_jdk_serialization";
+    public static final Version FIRST_CUSTOM_SERIALIZATION_SUPPORTED_OS_VERSION = Version.V_3_0_0;
 
     // On-behalf-of endpoints settings
     // CS-SUPPRESS-SINGLE: RegexpSingleline get Extensions Settings

src/main/java/org/opensearch/security/transport/SecurityInterceptor.java

@@ -38,7 +38,6 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
-import org.opensearch.Version;
 import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsAction;
 import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsResponse;
 import org.opensearch.action.get.GetRequest;
@@ -149,7 +148,7 @@
         final String origCCSTransientMf = getThreadContext().getTransient(ConfigConstants.OPENDISTRO_SECURITY_MASKED_FIELD_CCS);
 
         final boolean isDebugEnabled = log.isDebugEnabled();
-        final boolean useJDKSerialization = connection.getVersion().before(Version.V_3_0_0);
+        final boolean useJDKSerialization = connection.getVersion().before(ConfigConstants.FIRST_CUSTOM_SERIALIZATION_SUPPORTED_OS_VERSION);
         final boolean isSameNodeRequest = localNode != null && localNode.equals(connection.getNode());
 
         try (ThreadContext.StoredContext stashedContext = getThreadContext().stashContext()) {
@@ -228,12 +227,12 @@
             }
 
             if (useJDKSerialization) {
                 Map<String, String> jdkSerializedHeaders = new HashMap<>();
                 HeaderHelper.getAllSerializedHeaderNames()
                     .stream()
                     .filter(k -> headerMap.get(k) != null)
                     .forEach(k -> jdkSerializedHeaders.put(k, Base64Helper.ensureJDKSerialized(headerMap.get(k))));
                 headerMap.putAll(jdkSerializedHeaders);
             }
 
             getThreadContext().putHeader(headerMap);

src/test/java/org/opensearch/security/support/Base64HelperTest.java

@@ -40,4 +40,14 @@ public void testSerde() {
         Assert.assertEquals(test, dsJDK(test));
     }
 
+    @Test
+    public void testEnsureJDKSerialized() {
+        String test = "string";
+        String jdkSerialized = Base64Helper.serializeObject(test, true);
+        String customSerialized = Base64Helper.serializeObject(test, false);
+        Assert.assertEquals(jdkSerialized, Base64Helper.ensureJDKSerialized(jdkSerialized));
+        Assert.assertEquals(jdkSerialized, Base64Helper.ensureJDKSerialized(customSerialized));
+
+    }
+
 }