Skip applying field masking if search request is size 0 and only contains cardinality or count aggregations

.github/workflows/ci.yml

@@ -152,7 +152,7 @@ jobs:
       with:
         cache-disabled: true
         arguments: |
-            integrationTest -Dbuild.snapshot=false --tests org.opensearch.security.ResourceFocusedTests
+            integrationTest -Dbuild.snapshot=false --tests org.opensearch.security.stress.ResourceFocusedTests
 
   backward-compatibility-build:
     runs-on: ubuntu-latest

src/integrationTest/java/org/opensearch/security/FlsAndFieldMaskingTests.java

@@ -53,7 +53,9 @@
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.plugin.mapper.MapperSizePlugin;
 import org.opensearch.search.aggregations.Aggregation;
+import org.opensearch.search.aggregations.AggregationBuilders;
 import org.opensearch.search.aggregations.metrics.ParsedAvg;
+import org.opensearch.search.aggregations.metrics.ParsedCardinality;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.search.sort.SortOrder;
 import org.opensearch.test.framework.TestSecurityConfig;
@@ -65,6 +67,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.allOf;
 import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.everyItem;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.hasItem;
@@ -119,6 +122,7 @@ public class FlsAndFieldMaskingTests {
     static final String FIRST_INDEX_ID_SONG_2 = "INDEX_1_S2";
     static final String FIRST_INDEX_ID_SONG_3 = "INDEX_1_S3";
     static final String FIRST_INDEX_ID_SONG_4 = "INDEX_1_S4";
+    static final String FIRST_INDEX_ID_SONG_5 = "INDEX_1_S5";
     static final String SECOND_INDEX_ID_SONG_1 = "INDEX_2_S1";
     static final String SECOND_INDEX_ID_SONG_2 = "INDEX_2_S2";
     static final String SECOND_INDEX_ID_SONG_3 = "INDEX_2_S3";
@@ -174,6 +178,19 @@ public class FlsAndFieldMaskingTests {
                 .on(SECOND_INDEX_NAME)
         );
 
+    /**
+     * User who is allowed to see all fields on indices {@link #FIRST_INDEX_NAME} except artist
+     * <ul>
+     *     <li>values of the artist fields should be masked on index {@link #FIRST_INDEX_NAME}</li>
+     * </ul>
+     */
+    static final TestSecurityConfig.User MASKED_ARTIST_READER = new TestSecurityConfig.User("masked_title_lyrics_reader").roles(
+        new TestSecurityConfig.Role("masked_title_lyrics_reader").clusterPermissions("cluster_composite_ops_ro")
+            .indexPermissions("read")
+            .maskedFields(FIELD_ARTIST)
+            .on(FIRST_INDEX_NAME)
+    );
+
     /**
     * Function that converts field value to value masked with {@link #MASK_VALUE}
     */
@@ -304,6 +321,7 @@ public class FlsAndFieldMaskingTests {
             ALL_INDICES_STRING_ARTIST_READER,
             ALL_INDICES_STARS_LESS_THAN_ZERO_READER,
             TWINS_FIRST_ARTIST_READER,
+            MASKED_ARTIST_READER,
             USER_ONLY_FIELD_TITLE_FLS,
             USER_NO_FIELD_TITLE_FLS,
             USER_ONLY_FIELD_TITLE_MASKED,
@@ -346,6 +364,7 @@ public class FlsAndFieldMaskingTests {
             put(FIRST_INDEX_ID_SONG_2, SONGS[1]);
             put(FIRST_INDEX_ID_SONG_3, SONGS[2]);
             put(FIRST_INDEX_ID_SONG_4, SONGS[3]);
+            put(FIRST_INDEX_ID_SONG_5, SONGS[4]);
         }
     };
 
@@ -918,6 +937,22 @@ public void getFieldCapabilities() throws IOException {
         }
     }
 
+    @Test
+    public void getCardinalityAggregationOnMaskedField() throws IOException {
+        // FIELD MASKING
+        try (RestHighLevelClient restHighLevelClient = cluster.getRestHighLevelClient(MASKED_ARTIST_READER)) {
+            SearchRequest searchRequest = new SearchRequest(FIRST_INDEX_NAME);
+            SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+            sourceBuilder.aggregation(AggregationBuilders.cardinality("unique_artists_agg").field("artist.keyword"));
+            sourceBuilder.size(0);
+            searchRequest.source(sourceBuilder);
+            SearchResponse response = restHighLevelClient.search(searchRequest, DEFAULT);
+            ParsedCardinality parsedCardinality = response.getAggregations().get("unique_artists_agg");
+            long cardinality = parsedCardinality.getValue();
+            assertThat(cardinality, equalTo(4L));
+        }
+    }
+
     @Test
     public void testGetDocumentWithNoTitleFieldOrOnlyTitleFieldFLSRestrictions() throws IOException, Exception {
         GetRequest getRequest = new GetRequest(FIRST_INDEX_NAME, FIRST_INDEX_ID_SONG_1);
@@ -1106,9 +1141,9 @@ private void assertSearchForFLSRestrictions(
         SearchResponse searchResponse = restHighLevelClient.search(searchRequest, DEFAULT);
 
         assertThat(searchResponse, isSuccessfulSearchResponse());
-        assertThat(searchResponse, numberOfTotalHitsIsEqualTo(4));
+        assertThat(searchResponse, numberOfTotalHitsIsEqualTo(FIRST_INDEX_SONGS_BY_ID.keySet().size()));
 
-        IntStream.range(0, 4).forEach(hitIndex -> {
+        IntStream.range(0, FIRST_INDEX_SONGS_BY_ID.keySet().size()).forEach(hitIndex -> {
             assertThat(
                 searchResponse,
                 shouldShowFieldTitle
@@ -1244,8 +1279,8 @@ private void assertProperSearchResponseForTitleFieldMaskingRestriction(
         SearchResponse searchResponse = restHighLevelClient.search(searchRequest, DEFAULT);
 
         assertThat(searchResponse, isSuccessfulSearchResponse());
-        assertThat(searchResponse, numberOfTotalHitsIsEqualTo(4));
-        IntStream.range(0, 4).forEach(hitIndex -> {
+        assertThat(searchResponse, numberOfTotalHitsIsEqualTo(FIRST_INDEX_SONGS_BY_ID.keySet().size()));
+        IntStream.range(0, FIRST_INDEX_SONGS_BY_ID.keySet().size()).forEach(hitIndex -> {
             assertThat(
                 searchResponse,
                 searchHitContainsFieldWithValue(hitIndex, FIELD_TITLE, VALUE_TO_MASKED_VALUE.apply(SONGS[hitIndex].getTitle()))
@@ -1350,10 +1385,10 @@ private void assertProperSearchResponseForOnlyAndNoTitleFLSRestrictions(
         SearchResponse searchResponse = restHighLevelClient.search(searchRequest, DEFAULT);
 
         assertThat(searchResponse, isSuccessfulSearchResponse());
-        assertThat(searchResponse, numberOfTotalHitsIsEqualTo(4));
+        assertThat(searchResponse, numberOfTotalHitsIsEqualTo(FIRST_INDEX_SONGS_BY_ID.keySet().size()));
 
         // since the roles are overlapping, the role with less permissions is the only one that is used- which is no title
-        IntStream.range(0, 4).forEach(hitIndex -> {
+        IntStream.range(0, FIRST_INDEX_SONGS_BY_ID.keySet().size()).forEach(hitIndex -> {
             assertThat(searchResponse, searchHitDoesNotContainField(hitIndex, FIELD_TITLE));
             assertThat(searchResponse, searchHitDoesContainField(hitIndex, FIELD_ARTIST));
             assertThat(searchResponse, searchHitDoesContainField(hitIndex, FIELD_LYRICS));
@@ -1484,8 +1519,8 @@ private void assertProperSearchResponseForTitleFieldMaskingAndOnlyTitleFLSRestri
         SearchResponse searchResponse = restHighLevelClient.search(searchRequest, DEFAULT);
 
         assertThat(searchResponse, isSuccessfulSearchResponse());
-        assertThat(searchResponse, numberOfTotalHitsIsEqualTo(4));
-        IntStream.range(0, 4).forEach(hitIndex -> {
+        assertThat(searchResponse, numberOfTotalHitsIsEqualTo(FIRST_INDEX_SONGS_BY_ID.keySet().size()));
+        IntStream.range(0, FIRST_INDEX_SONGS_BY_ID.keySet().size()).forEach(hitIndex -> {
             assertThat(
                 searchResponse,
                 searchHitContainsFieldWithValue(hitIndex, FIELD_TITLE, VALUE_TO_MASKED_VALUE.apply(SONGS[hitIndex].getTitle()))
@@ -1594,8 +1629,8 @@ private void assertProperSearchResponseForTitleFieldMaskingAndNoTitleFLSRestrict
         SearchResponse searchResponse = restHighLevelClient.search(searchRequest, DEFAULT);
 
         assertThat(searchResponse, isSuccessfulSearchResponse());
-        assertThat(searchResponse, numberOfTotalHitsIsEqualTo(4));
-        IntStream.range(0, 4).forEach(hitIndex -> {
+        assertThat(searchResponse, numberOfTotalHitsIsEqualTo(FIRST_INDEX_SONGS_BY_ID.keySet().size()));
+        IntStream.range(0, FIRST_INDEX_SONGS_BY_ID.keySet().size()).forEach(hitIndex -> {
             assertThat(searchResponse, searchHitDoesNotContainField(hitIndex, FIELD_TITLE));
             assertThat(searchResponse, searchHitContainsFieldWithValue(hitIndex, FIELD_ARTIST, SONGS[hitIndex].getArtist()));
             assertThat(searchResponse, searchHitContainsFieldWithValue(hitIndex, FIELD_LYRICS, SONGS[hitIndex].getLyrics()));
@@ -1719,11 +1754,11 @@ private void assertProperSearchResponseForTitleFieldMaskingAndNoTitleFieldAndOnl
         SearchResponse searchResponse = restHighLevelClient.search(searchRequest, DEFAULT);
 
         assertThat(searchResponse, isSuccessfulSearchResponse());
-        assertThat(searchResponse, numberOfTotalHitsIsEqualTo(4));
+        assertThat(searchResponse, numberOfTotalHitsIsEqualTo(FIRST_INDEX_SONGS_BY_ID.keySet().size()));
 
         // since the roles are overlapping, the role with less permissions is the only one that is used- which is no title, and since there
         // is no title the masking role has no effect
-        IntStream.range(0, 4).forEach(hitIndex -> {
+        IntStream.range(0, FIRST_INDEX_SONGS_BY_ID.keySet().size()).forEach(hitIndex -> {
             assertThat(searchResponse, searchHitDoesNotContainField(hitIndex, FIELD_TITLE));
             assertThat(searchResponse, searchHitDoesContainField(hitIndex, FIELD_ARTIST));
             assertThat(searchResponse, searchHitDoesContainField(hitIndex, FIELD_LYRICS));

src/integrationTest/java/org/opensearch/security/Song.java

@@ -23,6 +23,7 @@ public class Song {
     public static final String ARTIST_STRING = "String";
     public static final String ARTIST_TWINS = "Twins";
     public static final String TITLE_MAGNUM_OPUS = "Magnum Opus";
+    public static final String TITLE_ABC = "ABC";
     public static final String TITLE_SONG_1_PLUS_1 = "Song 1+1";
     public static final String TITLE_NEXT_SONG = "Next song";
     public static final String ARTIST_NO = "No!";
@@ -41,6 +42,7 @@ public class Song {
     public static final String LYRICS_4 = "Much too much";
     public static final String LYRICS_5 = "Little to little";
     public static final String LYRICS_6 = "confidential secret classified";
+    public static final String LYRICS_7 = "abcdefghijklmnopqrstuvwxyz";
 
     public static final String GENRE_ROCK = "rock";
     public static final String GENRE_JAZZ = "jazz";
@@ -55,6 +57,7 @@ public class Song {
         new Song(ARTIST_STRING, TITLE_SONG_1_PLUS_1, LYRICS_2, 2, GENRE_BLUES),
         new Song(ARTIST_TWINS, TITLE_NEXT_SONG, LYRICS_3, 3, GENRE_JAZZ),
         new Song(ARTIST_NO, TITLE_POISON, LYRICS_4, 4, GENRE_ROCK),
+        new Song(ARTIST_FIRST, TITLE_ABC, LYRICS_7, 7, GENRE_ROCK),
         new Song(ARTIST_YES, TITLE_AFFIRMATIVE, LYRICS_5, 5, GENRE_BLUES),
         new Song(ARTIST_UNKNOWN, TITLE_CONFIDENTIAL, LYRICS_6, 6, GENRE_JAZZ) };
 

src/integrationTest/java/org/opensearch/security/stress/FieldMaskingStressTests.java

@@ -0,0 +1,110 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+package org.opensearch.security.stress;
+
+import java.io.IOException;
+import java.util.concurrent.TimeUnit;
+
+import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope;
+import org.apache.http.HttpStatus;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import org.opensearch.action.search.SearchRequest;
+import org.opensearch.action.search.SearchResponse;
+import org.opensearch.client.RestHighLevelClient;
+import org.opensearch.search.aggregations.AggregationBuilders;
+import org.opensearch.search.aggregations.metrics.ParsedCardinality;
+import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.test.framework.AsyncActions;
+import org.opensearch.test.framework.TestSecurityConfig;
+import org.opensearch.test.framework.cluster.ClusterManager;
+import org.opensearch.test.framework.cluster.LocalCluster;
+import org.opensearch.test.framework.cluster.TestRestClient;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+import static org.opensearch.client.RequestOptions.DEFAULT;
+import static org.opensearch.security.Song.FIELD_TITLE;
+import static org.opensearch.test.framework.TestSecurityConfig.AuthcDomain.AUTHC_HTTPBASIC_INTERNAL;
+import static org.opensearch.test.framework.TestSecurityConfig.Role.ALL_ACCESS;
+
+@RunWith(com.carrotsearch.randomizedtesting.RandomizedRunner.class)
+@ThreadLeakScope(ThreadLeakScope.Scope.NONE)
+public class FieldMaskingStressTests {
+    private final static Logger LOG = LogManager.getLogger(FieldMaskingStressTests.class);
+    private static final long BYTE_TO_MB_CONVERSION_VALUE = 1024 * 1024;
+    private static final TestSecurityConfig.User ADMIN_USER = new TestSecurityConfig.User("admin").roles(ALL_ACCESS);
+
+    /**
+     * User who is allowed to see all fields on indices songs except artist
+     * <ul>
+     *     <li>values of the artist fields should be masked on index songs</li>
+     * </ul>
+     */
+    static final TestSecurityConfig.User MASKED_TITLE_READER = new TestSecurityConfig.User("masked_title_reader").roles(
+        new TestSecurityConfig.Role("masked_title_reader").clusterPermissions("cluster_composite_ops_ro")
+            .indexPermissions("read")
+            .maskedFields(FIELD_TITLE)
+            .on("songs")
+    );
+
+    private String bar = "bar".repeat(100);
+    private String TEST_DOC = "{\"artist\": \"foo\", \"title\": \"" + bar + "\"}";
+
+    @ClassRule
+    public static LocalCluster cluster = new LocalCluster.Builder().clusterManager(ClusterManager.THREE_CLUSTER_MANAGERS)
+        .authc(AUTHC_HTTPBASIC_INTERNAL)
+        .users(ADMIN_USER, MASKED_TITLE_READER)
+        .anonymousAuth(false)
+        .doNotFailOnForbidden(true)
+        .build();
+
+    @Test
+    public void testCardinalityAggregationWithFieldMaskingOnLargeIndex() throws IOException {
+        final int totalNumberOfDocs = 10000;
+
+        try (TestRestClient client = cluster.getRestClient(ADMIN_USER)) {
+            final var requests = AsyncActions.generate(() -> {
+                TestRestClient.HttpResponse indexDocResponse = client.postJson("songs/_doc", TEST_DOC);
+                return indexDocResponse.getStatusCode();
+            }, 10, totalNumberOfDocs);
+
+            AsyncActions.getAll(requests, 2, TimeUnit.MINUTES)
+                .forEach((responseCode) -> { assertThat(responseCode, equalTo(HttpStatus.SC_CREATED)); });
+        }
+
+        System.gc();
+        long memoryUsageBeforeLoadingData = getCurrentlyUsedMemory();
+        System.out.println("Used memory before loading some data: " + memoryUsageBeforeLoadingData + " MB");
+        try (RestHighLevelClient restHighLevelClient = cluster.getRestHighLevelClient(MASKED_TITLE_READER)) {
+            SearchRequest searchRequest = new SearchRequest("songs");
+            SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+            sourceBuilder.aggregation(AggregationBuilders.cardinality("unique_titles_agg").field("title.keyword").precisionThreshold(100));
+            sourceBuilder.size(1);
+            searchRequest.source(sourceBuilder);
+            SearchResponse response = restHighLevelClient.search(searchRequest, DEFAULT);
+            System.out.println("SearchResponse: " + response);
+            ParsedCardinality parsedCardinality = response.getAggregations().get("unique_titles_agg");
+            long cardinality = parsedCardinality.getValue();
+            // assertThat(cardinality, equalTo(1L));
+        }
+        long memoryUsageAfterLoadingData = getCurrentlyUsedMemory();
+        System.out.println("Used memory after loading some data: " + memoryUsageAfterLoadingData + " MB");
+        System.out.println("Difference: " + (memoryUsageAfterLoadingData - memoryUsageBeforeLoadingData) + " MB");
+    }
+
+    private long getCurrentlyUsedMemory() {
+        return (Runtime.getRuntime().totalMemory() - Runtime.getRuntime().freeMemory()) / BYTE_TO_MB_CONVERSION_VALUE;
+    }
+}

src/integrationTest/java/org/opensearch/security/ResourceFocusedTests.java → src/integrationTest/java/org/opensearch/security/stress/ResourceFocusedTests.java

@@ -8,7 +8,7 @@
  *
  */
 
-package org.opensearch.security;
+package org.opensearch.security.stress;
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;

src/main/java/org/opensearch/security/configuration/DlsFlsValveImpl.java

@@ -534,12 +534,30 @@ private void setFlsHeaders(EvaluatedDlsFlsConfig dlsFls, ActionRequest request)
                         }
                     }
                 } else {
-                    threadContext.putHeader(
-                        ConfigConstants.OPENDISTRO_SECURITY_MASKED_FIELD_HEADER,
-                        Base64Helper.serializeObject((Serializable) maskedFieldsMap)
-                    );
-                    if (log.isDebugEnabled()) {
-                        log.debug("attach masked fields info: {}", maskedFieldsMap);
+                    boolean shouldIncludeFlsHeaders = true;
+                    if (request instanceof SearchRequest) {
+                        SearchRequest searchRequest = (SearchRequest) request;
+                        boolean hasOnlyCountOrCardinality = true;
+                        if (searchRequest.source().aggregations() != null) {
+                            for (AggregationBuilder af : searchRequest.source().aggregations().getAggregatorFactories()) {
+                                if (!(af.getType().equals("cardinality") || af.getType().equals("count"))) {
+                                    hasOnlyCountOrCardinality = false;
+                                }
+                            }
+                        }
+                        if (hasOnlyCountOrCardinality && searchRequest.source().size() == 0) {
+                            shouldIncludeFlsHeaders = false;
+                        }
+                    }
+
+                    if (shouldIncludeFlsHeaders) {
+                        threadContext.putHeader(
+                            ConfigConstants.OPENDISTRO_SECURITY_MASKED_FIELD_HEADER,
+                            Base64Helper.serializeObject((Serializable) maskedFieldsMap)
+                        );
+                        if (log.isDebugEnabled()) {
+                            log.debug("attach masked fields info: {}", maskedFieldsMap);
+                        }
                     }
                 }
             }