Check for runtime field loops in queries

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/AbstractScriptMappedFieldType.java

@@ -8,17 +8,20 @@
 
 import org.apache.lucene.analysis.TokenStream;
 import org.apache.lucene.search.MultiTermQuery;
+import org.apache.lucene.search.MultiTermQuery.RewriteMethod;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.spans.SpanMultiTermQueryWrapper;
 import org.apache.lucene.search.spans.SpanQuery;
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.common.TriFunction;
 import org.elasticsearch.common.geo.ShapeRelation;
 import org.elasticsearch.common.time.DateMathParser;
 import org.elasticsearch.common.unit.Fuzziness;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.TextSearchInfo;
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.script.Script;
+import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
 import java.time.ZoneId;
@@ -31,12 +34,19 @@
 /**
  * Abstract base {@linkplain MappedFieldType} for scripted fields.
  */
-abstract class AbstractScriptMappedFieldType extends MappedFieldType {
+abstract class AbstractScriptMappedFieldType<LeafFactory> extends MappedFieldType {
     protected final Script script;
+    private final TriFunction<String, Map<String, Object>, SearchLookup, LeafFactory> factory;
 
-    AbstractScriptMappedFieldType(String name, Script script, Map<String, String> meta) {
+    AbstractScriptMappedFieldType(
+        String name,
+        Script script,
+        TriFunction<String, Map<String, Object>, SearchLookup, LeafFactory> factory,
+        Map<String, String> meta
+    ) {
         super(name, false, false, TextSearchInfo.SIMPLE_MATCH_ONLY, meta);
         this.script = script;
+        this.factory = factory;
     }
 
     protected abstract String runtimeType();
@@ -61,6 +71,26 @@ public final boolean isAggregatable() {
         return true;
     }
 
+    /**
+     * Create a script leaf factory.
+     */
+    protected final LeafFactory leafFactory(SearchLookup searchLookup) {
+        return factory.apply(name(), script.getParams(), searchLookup);
+    }
+
+    /**
+     * Create a script leaf factory for queries.
+     */
+    protected final LeafFactory leafFactory(QueryShardContext context) {
+        /*
+         * Forking here causes us to count this field in the field data loop
+         * detection code as though we were resolving field data for this field.
+         * We're not, but running the query is close enough.
+         */
+        return leafFactory(context.lookup().forkAndTrackFieldReferences(name()));
+    }
+
+    @Override
     public abstract Query termsQuery(List<?> values, QueryShardContext context);
 
     @Override
@@ -119,7 +149,7 @@ public Query regexpQuery(
         int syntaxFlags,
         int matchFlags,
         int maxDeterminizedStates,
-        MultiTermQuery.RewriteMethod method,
+        RewriteMethod method,
         QueryShardContext context
     ) {
         throw new IllegalArgumentException(unsupported("regexp", "keyword and text"));
@@ -149,7 +179,7 @@ public SpanQuery spanPrefixQuery(String value, SpanMultiTermQueryWrapper.SpanRew
     }
 
     private String unsupported(String query, String supported) {
-        String thisField = "[" + name() + "] which is of type [script] with runtime_type [" + runtimeType() + "]";
+        String thisField = "[" + name() + "] which is of type [runtime] with runtime_type [" + runtimeType() + "]";
         return "Can only use " + query + " queries on " + supported + " fields - not on " + thisField;
     }
 

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/RuntimeFieldMapper.java

@@ -51,7 +51,7 @@ public <FactoryType> FactoryType compile(Script script, ScriptContext<FactoryTyp
 
     protected RuntimeFieldMapper(
         String simpleName,
-        AbstractScriptMappedFieldType mappedFieldType,
+        AbstractScriptMappedFieldType<?> mappedFieldType,
         MultiFields multiFields,
         CopyTo copyTo,
         String runtimeType,
@@ -86,7 +86,7 @@ protected String contentType() {
 
     public static class Builder extends ParametrizedFieldMapper.Builder {
 
-        static final Map<String, BiFunction<Builder, BuilderContext, AbstractScriptMappedFieldType>> FIELD_TYPE_RESOLVER = Map.of(
+        static final Map<String, BiFunction<Builder, BuilderContext, AbstractScriptMappedFieldType<?>>> FIELD_TYPE_RESOLVER = Map.of(
             BooleanFieldMapper.CONTENT_TYPE,
             (builder, context) -> {
                 builder.formatAndLocaleNotSupported();
@@ -211,7 +211,7 @@ private static RuntimeFieldMapper toType(FieldMapper in) {
         private final Parameter<String> format = Parameter.stringParam(
             "format",
             true,
-            mapper -> ((AbstractScriptMappedFieldType) mapper.fieldType()).format(),
+            mapper -> ((AbstractScriptMappedFieldType<?>) mapper.fieldType()).format(),
             null
         ).setSerializer((b, n, v) -> {
             if (v != null && false == v.equals(DateFieldMapper.DEFAULT_DATE_TIME_FORMATTER.pattern())) {
@@ -223,7 +223,7 @@ private static RuntimeFieldMapper toType(FieldMapper in) {
             true,
             () -> null,
             (n, c, o) -> o == null ? null : LocaleUtils.parse(o.toString()),
-            mapper -> ((AbstractScriptMappedFieldType) mapper.fieldType()).formatLocale()
+            mapper -> ((AbstractScriptMappedFieldType<?>) mapper.fieldType()).formatLocale()
         ).setSerializer((b, n, v) -> {
             if (v != null && false == v.equals(Locale.ROOT)) {
                 b.field(n, v.toString());
@@ -244,7 +244,7 @@ protected List<Parameter<?>> getParameters() {
 
         @Override
         public RuntimeFieldMapper build(BuilderContext context) {
-            BiFunction<Builder, BuilderContext, AbstractScriptMappedFieldType> fieldTypeResolver = Builder.FIELD_TYPE_RESOLVER.get(
+            BiFunction<Builder, BuilderContext, AbstractScriptMappedFieldType<?>> fieldTypeResolver = Builder.FIELD_TYPE_RESOLVER.get(
                 runtimeType.getValue()
             );
             if (fieldTypeResolver == null) {

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/ScriptBooleanMappedFieldType.java

@@ -27,12 +27,9 @@
 import java.util.Map;
 import java.util.function.Supplier;
 
-public class ScriptBooleanMappedFieldType extends AbstractScriptMappedFieldType {
-    private final BooleanScriptFieldScript.Factory scriptFactory;
-
+public class ScriptBooleanMappedFieldType extends AbstractScriptMappedFieldType<BooleanScriptFieldScript.LeafFactory> {
     ScriptBooleanMappedFieldType(String name, Script script, BooleanScriptFieldScript.Factory scriptFactory, Map<String, String> meta) {
-        super(name, script, meta);
-        this.scriptFactory = scriptFactory;
+        super(name, script, scriptFactory::newFactory, meta);
     }
 
     @Override
@@ -71,14 +68,10 @@ public ScriptBooleanFieldData.Builder fielddataBuilder(String fullyQualifiedInde
         return new ScriptBooleanFieldData.Builder(name(), leafFactory(searchLookup.get()));
     }
 
-    private BooleanScriptFieldScript.LeafFactory leafFactory(SearchLookup searchLookup) {
-        return scriptFactory.newFactory(name(), script.getParams(), searchLookup);
-    }
-
     @Override
     public Query existsQuery(QueryShardContext context) {
         checkAllowExpensiveQueries(context);
-        return new BooleanScriptFieldExistsQuery(script, leafFactory(context.lookup()), name());
+        return new BooleanScriptFieldExistsQuery(script, leafFactory(context), name());
     }
 
     @Override
@@ -149,7 +142,7 @@ public Query rangeQuery(
     @Override
     public Query termQuery(Object value, QueryShardContext context) {
         checkAllowExpensiveQueries(context);
-        return new BooleanScriptFieldTermQuery(script, leafFactory(context.lookup()), name(), toBoolean(value));
+        return new BooleanScriptFieldTermQuery(script, leafFactory(context), name(), toBoolean(value));
     }
 
     @Override
@@ -176,11 +169,11 @@ private Query termsQuery(boolean trueAllowed, boolean falseAllowed, QueryShardCo
                 return existsQuery(context);
             }
             checkAllowExpensiveQueries(context);
-            return new BooleanScriptFieldTermQuery(script, leafFactory(context.lookup()), name(), true);
+            return new BooleanScriptFieldTermQuery(script, leafFactory(context), name(), true);
         }
         if (falseAllowed) {
             checkAllowExpensiveQueries(context);
-            return new BooleanScriptFieldTermQuery(script, leafFactory(context.lookup()), name(), false);
+            return new BooleanScriptFieldTermQuery(script, leafFactory(context), name(), false);
         }
         return new MatchNoDocsQuery("neither true nor false allowed");
     }

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/ScriptDateMappedFieldType.java

@@ -37,8 +37,7 @@
 import java.util.Map;
 import java.util.function.Supplier;
 
-public class ScriptDateMappedFieldType extends AbstractScriptMappedFieldType {
-    private final DateScriptFieldScript.Factory scriptFactory;
+public class ScriptDateMappedFieldType extends AbstractScriptMappedFieldType<DateScriptFieldScript.LeafFactory> {
     private final DateFormatter dateTimeFormatter;
 
     ScriptDateMappedFieldType(
@@ -48,8 +47,7 @@ public class ScriptDateMappedFieldType extends AbstractScriptMappedFieldType {
         DateFormatter dateTimeFormatter,
         Map<String, String> meta
     ) {
-        super(name, script, meta);
-        this.scriptFactory = scriptFactory;
+        super(name, script, (n, params, ctx) -> scriptFactory.newFactory(n, params, ctx, dateTimeFormatter), meta);
         this.dateTimeFormatter = dateTimeFormatter;
     }
 
@@ -84,10 +82,6 @@ public ScriptDateFieldData.Builder fielddataBuilder(String fullyQualifiedIndexNa
         return new ScriptDateFieldData.Builder(name(), leafFactory(lookup.get()));
     }
 
-    private DateScriptFieldScript.LeafFactory leafFactory(SearchLookup lookup) {
-        return scriptFactory.newFactory(name(), script.getParams(), lookup, dateTimeFormatter);
-    }
-
     @Override
     public Query distanceFeatureQuery(Object origin, String pivot, float boost, QueryShardContext context) {
         checkAllowExpensiveQueries(context);
@@ -103,7 +97,7 @@ public Query distanceFeatureQuery(Object origin, String pivot, float boost, Quer
             TimeValue pivotTime = TimeValue.parseTimeValue(pivot, "distance_feature.pivot");
             return new LongScriptFieldDistanceFeatureQuery(
                 script,
-                leafFactory(context.lookup())::newInstance,
+                leafFactory(context)::newInstance,
                 name(),
                 originLong,
                 pivotTime.getMillis(),
@@ -115,7 +109,7 @@ public Query distanceFeatureQuery(Object origin, String pivot, float boost, Quer
     @Override
     public Query existsQuery(QueryShardContext context) {
         checkAllowExpensiveQueries(context);
-        return new LongScriptFieldExistsQuery(script, leafFactory(context.lookup())::newInstance, name());
+        return new LongScriptFieldExistsQuery(script, leafFactory(context)::newInstance, name());
     }
 
     @Override
@@ -139,7 +133,7 @@ public Query rangeQuery(
             parser,
             context,
             DateFieldMapper.Resolution.MILLISECONDS,
-            (l, u) -> new LongScriptFieldRangeQuery(script, leafFactory(context.lookup())::newInstance, name(), l, u)
+            (l, u) -> new LongScriptFieldRangeQuery(script, leafFactory(context)::newInstance, name(), l, u)
         );
     }
 
@@ -155,7 +149,7 @@ public Query termQuery(Object value, QueryShardContext context) {
                 DateFieldMapper.Resolution.MILLISECONDS
             );
             checkAllowExpensiveQueries(context);
-            return new LongScriptFieldTermQuery(script, leafFactory(context.lookup())::newInstance, name(), l);
+            return new LongScriptFieldTermQuery(script, leafFactory(context)::newInstance, name(), l);
         });
     }
 
@@ -179,7 +173,7 @@ public Query termsQuery(List<?> values, QueryShardContext context) {
                 );
             }
             checkAllowExpensiveQueries(context);
-            return new LongScriptFieldTermsQuery(script, leafFactory(context.lookup())::newInstance, name(), terms);
+            return new LongScriptFieldTermsQuery(script, leafFactory(context)::newInstance, name(), terms);
         });
     }
 

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/ScriptDoubleMappedFieldType.java

@@ -29,12 +29,9 @@
 import java.util.Map;
 import java.util.function.Supplier;
 
-public class ScriptDoubleMappedFieldType extends AbstractScriptMappedFieldType {
-    private final DoubleScriptFieldScript.Factory scriptFactory;
-
+public class ScriptDoubleMappedFieldType extends AbstractScriptMappedFieldType<DoubleScriptFieldScript.LeafFactory> {
     ScriptDoubleMappedFieldType(String name, Script script, DoubleScriptFieldScript.Factory scriptFactory, Map<String, String> meta) {
-        super(name, script, meta);
-        this.scriptFactory = scriptFactory;
+        super(name, script, scriptFactory::newFactory, meta);
     }
 
     @Override
@@ -63,14 +60,10 @@ public ScriptDoubleFieldData.Builder fielddataBuilder(String fullyQualifiedIndex
         return new ScriptDoubleFieldData.Builder(name(), leafFactory(searchLookup.get()));
     }
 
-    private DoubleScriptFieldScript.LeafFactory leafFactory(SearchLookup searchLookup) {
-        return scriptFactory.newFactory(name(), script.getParams(), searchLookup);
-    }
-
     @Override
     public Query existsQuery(QueryShardContext context) {
         checkAllowExpensiveQueries(context);
-        return new DoubleScriptFieldExistsQuery(script, leafFactory(context.lookup()), name());
+        return new DoubleScriptFieldExistsQuery(script, leafFactory(context), name());
     }
 
     @Override
@@ -89,14 +82,14 @@ public Query rangeQuery(
             upperTerm,
             includeLower,
             includeUpper,
-            (l, u) -> new DoubleScriptFieldRangeQuery(script, leafFactory(context.lookup()), name(), l, u)
+            (l, u) -> new DoubleScriptFieldRangeQuery(script, leafFactory(context), name(), l, u)
         );
     }
 
     @Override
     public Query termQuery(Object value, QueryShardContext context) {
         checkAllowExpensiveQueries(context);
-        return new DoubleScriptFieldTermQuery(script, leafFactory(context.lookup()), name(), NumberType.objectToDouble(value));
+        return new DoubleScriptFieldTermQuery(script, leafFactory(context), name(), NumberType.objectToDouble(value));
     }
 
     @Override
@@ -109,6 +102,6 @@ public Query termsQuery(List<?> values, QueryShardContext context) {
             terms.add(Double.doubleToLongBits(NumberType.objectToDouble(value)));
         }
         checkAllowExpensiveQueries(context);
-        return new DoubleScriptFieldTermsQuery(script, leafFactory(context.lookup()), name(), terms);
+        return new DoubleScriptFieldTermsQuery(script, leafFactory(context), name(), terms);
     }
 }

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/ScriptIpMappedFieldType.java

@@ -37,15 +37,9 @@
 import java.util.Map;
 import java.util.function.Supplier;
 
-public final class ScriptIpMappedFieldType extends AbstractScriptMappedFieldType {
-
-    private final Script script;
-    private final IpScriptFieldScript.Factory scriptFactory;
-
+public final class ScriptIpMappedFieldType extends AbstractScriptMappedFieldType<IpScriptFieldScript.LeafFactory> {
     ScriptIpMappedFieldType(String name, Script script, IpScriptFieldScript.Factory scriptFactory, Map<String, String> meta) {
-        super(name, script, meta);
-        this.script = script;
-        this.scriptFactory = scriptFactory;
+        super(name, script, scriptFactory::newFactory, meta);
     }
 
     @Override
@@ -79,14 +73,10 @@ public ScriptIpFieldData.Builder fielddataBuilder(String fullyQualifiedIndexName
         return new ScriptIpFieldData.Builder(name(), leafFactory(searchLookup.get()));
     }
 
-    private IpScriptFieldScript.LeafFactory leafFactory(SearchLookup searchLookup) {
-        return scriptFactory.newFactory(name(), script.getParams(), searchLookup);
-    }
-
     @Override
     public Query existsQuery(QueryShardContext context) {
         checkAllowExpensiveQueries(context);
-        return new IpScriptFieldExistsQuery(script, leafFactory(context.lookup()), name());
+        return new IpScriptFieldExistsQuery(script, leafFactory(context), name());
     }
 
     @Override
@@ -107,7 +97,7 @@ public Query rangeQuery(
             includeUpper,
             (lower, upper) -> new IpScriptFieldRangeQuery(
                 script,
-                leafFactory(context.lookup()),
+                leafFactory(context),
                 name(),
                 new BytesRef(InetAddressPoint.encode(lower)),
                 new BytesRef(InetAddressPoint.encode(upper))
@@ -119,14 +109,18 @@ public Query rangeQuery(
     public Query termQuery(Object value, QueryShardContext context) {
         checkAllowExpensiveQueries(context);
         if (value instanceof InetAddress) {
-            return InetAddressPoint.newExactQuery(name(), (InetAddress) value);
+            return inetAddressQuery((InetAddress) value, context);
         }
         String term = BytesRefs.toString(value);
         if (term.contains("/")) {
             return cidrQuery(term, context);
         }
         InetAddress address = InetAddresses.forString(term);
-        return new IpScriptFieldTermQuery(script, leafFactory(context.lookup()), name(), new BytesRef(InetAddressPoint.encode(address)));
+        return inetAddressQuery(address, context);
+    }
+
+    private Query inetAddressQuery(InetAddress address, QueryShardContext context) {
+        return new IpScriptFieldTermQuery(script, leafFactory(context), name(), new BytesRef(InetAddressPoint.encode(address)));
     }
 
     @Override
@@ -149,7 +143,7 @@ public Query termsQuery(List<?> values, QueryShardContext context) {
             }
             cidrQueries.add(cidrQuery(term, context));
         }
-        Query termsQuery = new IpScriptFieldTermsQuery(script, leafFactory(context.lookup()), name(), terms);
+        Query termsQuery = new IpScriptFieldTermsQuery(script, leafFactory(context), name(), terms);
         if (cidrQueries == null) {
             return termsQuery;
         }
@@ -176,6 +170,6 @@ private Query cidrQuery(String term, QueryShardContext context) {
         // Force the terms into IPv6
         BytesRef lowerBytes = new BytesRef(InetAddressPoint.encode(InetAddressPoint.decode(lower)));
         BytesRef upperBytes = new BytesRef(InetAddressPoint.encode(InetAddressPoint.decode(upper)));
-        return new IpScriptFieldRangeQuery(script, leafFactory(context.lookup()), name(), lowerBytes, upperBytes);
+        return new IpScriptFieldRangeQuery(script, leafFactory(context), name(), lowerBytes, upperBytes);
     }
 }