Fix DLS over Runtime Fields (#112260) (#112318)

There is a DLS query referencing a runtime field loaded from _source, when we create the collector manager we retrieve the numDocs which triggers going through all segments and executing the script for each document. StoredFieldSourceProvider relies on leaf ordinals to build an array, but those ordinals are not populated when computing the numDocs via BaseCompositeReader, as that goes through the subreaders contexts, and not the context leaves (there is a subtle difference that bites us there).

Fixes #111637

Author: original-brownbear

docs/changelog/112260.yaml

@@ -0,0 +1,6 @@
+pr: 112260
+summary: Fix DLS over Runtime Fields
+area: "Authorization"
+type: bug
+issues:
+ - 111637

server/src/main/java/org/elasticsearch/search/lookup/StoredFieldSourceProvider.java

@@ -8,53 +8,36 @@
 
 package org.elasticsearch.search.lookup;
 
-import org.apache.lucene.index.IndexReaderContext;
 import org.apache.lucene.index.LeafReaderContext;
+import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.index.fieldvisitor.LeafStoredFieldLoader;
 import org.elasticsearch.index.fieldvisitor.StoredFieldLoader;
 
 import java.io.IOException;
+import java.util.Map;
 
 // NB This is written under the assumption that individual segments are accessed by a single
 // thread, even if separate segments may be searched concurrently.  If we ever implement
 // within-segment concurrency this will have to work entirely differently.
 class StoredFieldSourceProvider implements SourceProvider {
 
     private final StoredFieldLoader storedFieldLoader;
-    private volatile LeafStoredFieldSourceProvider[] leaves;
+    private final Map<Object, LeafStoredFieldSourceProvider> leaves = ConcurrentCollections.newConcurrentMap();
 
     StoredFieldSourceProvider(StoredFieldLoader storedFieldLoader) {
         this.storedFieldLoader = storedFieldLoader;
     }
 
     @Override
     public Source getSource(LeafReaderContext ctx, int doc) throws IOException {
-        LeafStoredFieldSourceProvider[] leaves = getLeavesUnderLock(findParentContext(ctx));
-        if (leaves[ctx.ord] == null) {
-            // individual segments are currently only accessed on one thread so there's no need
-            // for locking here.
-            leaves[ctx.ord] = new LeafStoredFieldSourceProvider(storedFieldLoader.getLoader(ctx, null));
+        final Object id = ctx.id();
+        var provider = leaves.get(id);
+        if (provider == null) {
+            provider = new LeafStoredFieldSourceProvider(storedFieldLoader.getLoader(ctx, null));
+            var existing = leaves.put(id, provider);
+            assert existing == null : "unexpected source provider [" + existing + "]";
         }
-        return leaves[ctx.ord].getSource(doc);
-    }
-
-    private static IndexReaderContext findParentContext(LeafReaderContext ctx) {
-        if (ctx.parent != null) {
-            return ctx.parent;
-        }
-        assert ctx.isTopLevel;
-        return ctx;
-    }
-
-    private LeafStoredFieldSourceProvider[] getLeavesUnderLock(IndexReaderContext parentCtx) {
-        if (leaves == null) {
-            synchronized (this) {
-                if (leaves == null) {
-                    leaves = new LeafStoredFieldSourceProvider[parentCtx.leaves().size()];
-                }
-            }
-        }
-        return leaves;
+        return provider.getSource(doc);
     }
 
     private static class LeafStoredFieldSourceProvider {

x-pack/plugin/security/src/internalClusterTest/java/org/elasticsearch/integration/DocumentLevelSecurityRandomTests.java

@@ -13,13 +13,16 @@
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.test.SecurityIntegTestCase;
 import org.elasticsearch.test.SecuritySettingsSourceField;
+import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xpack.core.XPackSettings;
+import org.junit.BeforeClass;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
 import static org.elasticsearch.xpack.core.security.authc.support.UsernamePasswordToken.BASIC_AUTH_HEADER;
 import static org.elasticsearch.xpack.core.security.authc.support.UsernamePasswordToken.basicAuthHeaderValue;
@@ -29,9 +32,12 @@ public class DocumentLevelSecurityRandomTests extends SecurityIntegTestCase {
 
     protected static final SecureString USERS_PASSWD = SecuritySettingsSourceField.TEST_PASSWORD_SECURE_STRING;
 
-    // can't add a second test method, because each test run creates a new instance of this class and that will will result
-    // in a new random value:
-    private final int numberOfRoles = scaledRandomIntBetween(3, 99);
+    private static volatile int numberOfRoles;
+
+    @BeforeClass
+    public static void setupRoleCount() throws Exception {
+        numberOfRoles = scaledRandomIntBetween(3, 99);
+    }
 
     @Override
     protected String configUsers() {
@@ -119,4 +125,38 @@ public void testDuelWithAliasFilters() throws Exception {
         }
     }
 
+    public void testWithRuntimeFields() throws Exception {
+        assertAcked(
+            indicesAdmin().prepareCreate("test")
+                .setMapping(
+                    XContentFactory.jsonBuilder()
+                        .startObject()
+                        .startObject("runtime")
+                        .startObject("field1")
+                        .field("type", "keyword")
+                        .endObject()
+                        .endObject()
+                        .startObject("properties")
+                        .startObject("field2")
+                        .field("type", "keyword")
+                        .endObject()
+                        .endObject()
+                        .endObject()
+                )
+        );
+        List<IndexRequestBuilder> requests = new ArrayList<>(47);
+        for (int i = 1; i <= 42; i++) {
+            requests.add(prepareIndex("test").setSource("field1", "value1", "field2", "foo" + i));
+        }
+        for (int i = 42; i <= 57; i++) {
+            requests.add(prepareIndex("test").setSource("field1", "value2", "field2", "foo" + i));
+        }
+        indexRandom(true, requests);
+        assertHitCount(
+            client().filterWithHeader(Collections.singletonMap(BASIC_AUTH_HEADER, basicAuthHeaderValue("user1", USERS_PASSWD)))
+                .prepareSearch("test"),
+            42L
+        );
+    }
+
 }