Fix Source Return Bug in Scripting (#56831)

This change ensures that when a user returns _source directly no matter where 
accessed within scripting, the value is a Map of the converted source as 
opposed to a SourceLookup.

Fixes #52103

modules/lang-painless/src/test/java/org/elasticsearch/painless/ScriptedMetricAggContextsTests.java

@@ -91,6 +91,33 @@ public class ScriptedMetricAggContextsTests extends ScriptTestCase {
         assertEquals(1.0, state.get("testField"));
     }
 
+    public void testReturnSource() throws IOException {
+        ScriptedMetricAggContexts.MapScript.Factory factory = scriptEngine.compile("test",
+                "state._source = params._source", ScriptedMetricAggContexts.MapScript.CONTEXT, Collections.emptyMap());
+
+        Map<String, Object> params = new HashMap<>();
+        Map<String, Object> state = new HashMap<>();
+
+        MemoryIndex index = new MemoryIndex();
+        // we don't need a real index, just need to construct a LeafReaderContext which cannot be mocked
+        LeafReaderContext leafReaderContext = index.createSearcher().getIndexReader().leaves().get(0);
+
+        SearchLookup lookup = mock(SearchLookup.class);
+        LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
+        when(lookup.getLeafSearchLookup(leafReaderContext)).thenReturn(leafLookup);
+        SourceLookup sourceLookup = mock(SourceLookup.class);
+        when(leafLookup.asMap()).thenReturn(Collections.singletonMap("_source", sourceLookup));
+        when(sourceLookup.loadSourceIfNeeded()).thenReturn(Collections.singletonMap("test", 1));
+        ScriptedMetricAggContexts.MapScript.LeafFactory leafFactory = factory.newFactory(params, state, lookup);
+        ScriptedMetricAggContexts.MapScript script = leafFactory.newInstance(leafReaderContext);
+
+        script.execute();
+
+        assertTrue(state.containsKey("_source"));
+        assertTrue(state.get("_source") instanceof Map && ((Map)state.get("_source")).containsKey("test"));
+        assertEquals(1, ((Map)state.get("_source")).get("test"));
+    }
+
     public void testMapSourceAccess() throws IOException {
         ScriptedMetricAggContexts.MapScript.Factory factory = scriptEngine.compile("test",
             "state.testField = params._source.three", ScriptedMetricAggContexts.MapScript.CONTEXT, Collections.emptyMap());
@@ -107,13 +134,13 @@ public class ScriptedMetricAggContextsTests extends ScriptTestCase {
         when(lookup.getLeafSearchLookup(leafReaderContext)).thenReturn(leafLookup);
         SourceLookup sourceLookup = mock(SourceLookup.class);
         when(leafLookup.asMap()).thenReturn(Collections.singletonMap("_source", sourceLookup));
-        when(sourceLookup.get("three")).thenReturn(3);
+        when(sourceLookup.loadSourceIfNeeded()).thenReturn(Collections.singletonMap("three", 3));
         ScriptedMetricAggContexts.MapScript.LeafFactory leafFactory = factory.newFactory(params, state, lookup);
         ScriptedMetricAggContexts.MapScript script = leafFactory.newInstance(leafReaderContext);
 
         script.execute();
 
-        assert(state.containsKey("testField"));
+        assertTrue(state.containsKey("testField"));
         assertEquals(3, state.get("testField"));
     }
 

server/src/main/java/org/elasticsearch/script/AbstractSortScript.java

@@ -27,6 +27,7 @@ import org.elasticsearch.common.lucene.ScorerAware;
 import org.elasticsearch.index.fielddata.ScriptDocValues;
 import org.elasticsearch.search.lookup.LeafSearchLookup;
 import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.search.lookup.SourceLookup;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -49,7 +50,9 @@ abstract class AbstractSortScript implements ScorerAware {
                         "Accessing variable [doc] via [params._doc] from within an sort-script "
                                 + "is deprecated in favor of directly accessing [doc].");
                 return value;
-            });
+            },
+            "_source", value -> ((SourceLookup)value).loadSourceIfNeeded()
+    );
 
     /**
      * The generic runtime parameters for the script.

server/src/main/java/org/elasticsearch/script/AggregationScript.java

@@ -27,6 +27,7 @@ import org.elasticsearch.common.lucene.ScorerAware;
 import org.elasticsearch.index.fielddata.ScriptDocValues;
 import org.elasticsearch.search.lookup.LeafSearchLookup;
 import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.search.lookup.SourceLookup;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -53,7 +54,9 @@ public abstract class AggregationScript implements ScorerAware {
                         "Accessing variable [doc] via [params._doc] from within an aggregation-script "
                                 + "is deprecated in favor of directly accessing [doc].");
                 return value;
-            });
+            },
+            "_source", value -> ((SourceLookup)value).loadSourceIfNeeded()
+    );
 
     /**
      * The generic runtime parameters for the script.

server/src/main/java/org/elasticsearch/script/FieldScript.java

@@ -25,6 +25,7 @@ import org.elasticsearch.common.logging.DeprecationLogger;
 import org.elasticsearch.index.fielddata.ScriptDocValues;
 import org.elasticsearch.search.lookup.LeafSearchLookup;
 import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.search.lookup.SourceLookup;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -52,7 +53,9 @@ public abstract class FieldScript {
                         "Accessing variable [doc] via [params._doc] from within an field-script "
                                 + "is deprecated in favor of directly accessing [doc].");
                 return value;
-            });
+            },
+            "_source", value -> ((SourceLookup)value).loadSourceIfNeeded()
+    );
 
     /** The generic runtime parameters for the script. */
     private final Map<String, Object> params;

server/src/main/java/org/elasticsearch/script/ScoreScript.java

@@ -27,6 +27,7 @@ import org.elasticsearch.common.logging.DeprecationLogger;
 import org.elasticsearch.index.fielddata.ScriptDocValues;
 import org.elasticsearch.search.lookup.LeafSearchLookup;
 import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.search.lookup.SourceLookup;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
@@ -78,7 +79,9 @@ public abstract class ScoreScript {
                         "Accessing variable [doc] via [params._doc] from within an score-script "
                                 + "is deprecated in favor of directly accessing [doc].");
                 return value;
-            });
+            },
+            "_source", value -> ((SourceLookup)value).loadSourceIfNeeded()
+    );
 
     public static final String[] PARAMETERS = new String[]{ "explanation" };
 

server/src/main/java/org/elasticsearch/script/ScriptedMetricAggContexts.java

@@ -27,6 +27,7 @@ import org.elasticsearch.common.logging.DeprecationLogger;
 import org.elasticsearch.index.fielddata.ScriptDocValues;
 import org.elasticsearch.search.lookup.LeafSearchLookup;
 import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.search.lookup.SourceLookup;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -84,7 +85,9 @@ public class ScriptedMetricAggContexts {
                             "Accessing variable [_agg] via [params._agg] from within a scripted metric agg map script "
                                     + "is deprecated in favor of using [state].");
                     return value;
-                });
+                },
+                "_source", value -> ((SourceLookup)value).loadSourceIfNeeded()
+        );
 
         private final Map<String, Object> params;
         private final Map<String, Object> state;

server/src/main/java/org/elasticsearch/script/TermsSetQueryScript.java

@@ -24,6 +24,7 @@ import org.elasticsearch.common.logging.DeprecationLogger;
 import org.elasticsearch.index.fielddata.ScriptDocValues;
 import org.elasticsearch.search.lookup.LeafSearchLookup;
 import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.search.lookup.SourceLookup;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -50,7 +51,9 @@ public abstract class TermsSetQueryScript {
                         "Accessing variable [doc] via [params._doc] from within an terms-set-query-script "
                                 + "is deprecated in favor of directly accessing [doc].");
                 return value;
-            });
+            },
+            "_source", value -> ((SourceLookup)value).loadSourceIfNeeded()
+    );
 
     /**
      * The generic runtime parameters for the script.

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

@@ -54,7 +54,11 @@ public class SourceLookup implements Map<String, Object> {
         return sourceContentType;
     }
 
-    private Map<String, Object> loadSourceIfNeeded() {
+    // Scripting requires this method to be public. Using source()
+    // is not possible because certain checks use source == null as
+    // as a determination if source is enabled/disabled, but it should
+    // never be a null Map for scripting even when disabled.
+    public Map<String, Object> loadSourceIfNeeded() {
         if (source != null) {
             return source;
         }