Script: Protected _source inside update scripts (#88733)

CtxMap delegates all metadata keys to it's `Metadata` container and
all other keys to it's source map.  In most write contexts (update,
update by query, reindex), the source map should only contain one
key, `_source`, which has a `Map<String, Object>`.

This change adds validations to writes to the source map that rejects
insertion of invalid keys, illegal removal of the `_source` key, and
illegally overwriting the `_source` mapping with the wrong type.

docs/changelog/88733.yaml

@@ -0,0 +1,5 @@
+pr: 88733
+summary: "Script: Protected `_source` inside update scripts"
+area: Infra/Scripting
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/ingest/IngestCtxMap.java

@@ -56,6 +56,14 @@ class IngestCtxMap extends CtxMap {
         super(source, metadata);
     }
 
+    /**
+     * In ingest, all non-metadata keys are source keys, so the {@link #source} map is accessed directly from ctx.
+     */
+    @Override
+    protected boolean directSourceAccess() {
+        return true;
+    }
+
     /**
      * Fetch the timestamp from the ingestMetadata, if it exists
      * @return the timestamp for the document or null
@@ -72,15 +80,4 @@ class IngestCtxMap extends CtxMap {
         }
         return null;
     }
-
-    @Override
-    public Map<String, Object> getSource() {
-        return source;
-    }
-
-    @Override
-    protected Map<String, Object> wrapSource(Map<String, Object> source) {
-        // Not wrapped in Ingest
-        return source;
-    }
 }

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

@@ -8,14 +8,12 @@
 
 package org.elasticsearch.script;
 
-import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.common.util.set.Sets;
 
 import java.util.AbstractCollection;
 import java.util.AbstractMap;
 import java.util.AbstractSet;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.Map;
@@ -30,7 +28,7 @@ import java.util.stream.Collectors;
  */
 public class CtxMap extends AbstractMap<String, Object> {
     protected static final String SOURCE = "_source";
-    protected final Map<String, Object> source;
+    protected Map<String, Object> source;
     protected final Metadata metadata;
 
     /**
@@ -40,7 +38,7 @@ public class CtxMap extends AbstractMap<String, Object> {
      * @param metadata the metadata map
      */
     protected CtxMap(Map<String, Object> source, Metadata metadata) {
-        this.source = wrapSource(source != null ? source : new HashMap<>());
+        this.source = source;
         this.metadata = metadata;
         Set<String> badKeys = Sets.intersection(this.metadata.keySet(), this.source.keySet());
         if (badKeys.size() > 0) {
@@ -52,20 +50,21 @@ public class CtxMap extends AbstractMap<String, Object> {
         }
     }
 
-    protected Map<String, Object> wrapSource(Map<String, Object> source) {
-        Map<String, Object> wrapper = Maps.newHashMapWithExpectedSize(1);
-        wrapper.put(SOURCE, source);
-        return wrapper;
+    /**
+     * Does this access to the internal {@link #source} map occur directly via ctx? ie {@code ctx['myField']}.
+     * Or does it occur via the {@link #SOURCE} key? ie {@code ctx['_source']['myField']}.
+     *
+     * Defaults to indirect, {@code ctx['_source']}
+     */
+    protected boolean directSourceAccess() {
+        return false;
     }
 
     /**
      * get the source map, if externally modified then the guarantees of this class are not enforced
      */
-    @SuppressWarnings("unchecked")
-    public Map<String, Object> getSource() {
-        Object rawSource = source.get(SOURCE);
-        assert rawSource instanceof Map<?, ?> : " wrapped source of unexpected type";
-        return (Map<String, Object>) rawSource;
+    public final Map<String, Object> getSource() {
+        return source;
     }
 
     /**
@@ -81,7 +80,8 @@ public class CtxMap extends AbstractMap<String, Object> {
     @Override
     public Set<Map.Entry<String, Object>> entrySet() {
         // Make a copy of the Metadata.keySet() to avoid a ConcurrentModificationException when removing a value from the iterator
-        return new EntrySet(source.entrySet(), new HashSet<>(metadata.keySet()));
+        Set<Map.Entry<String, Object>> sourceEntries = directSourceAccess() ? source.entrySet() : Set.of(new IndirectSourceEntry());
+        return new EntrySet(sourceEntries, new HashSet<>(metadata.keySet()));
     }
 
     /**
@@ -93,7 +93,34 @@ public class CtxMap extends AbstractMap<String, Object> {
         if (metadata.isAvailable(key)) {
             return metadata.put(key, value);
         }
-        return source.put(key, value);
+        if (directSourceAccess()) {
+            return source.put(key, value);
+        } else if (SOURCE.equals(key)) {
+            return replaceSource(value);
+        }
+        throw new IllegalArgumentException("Cannot put key [" + key + "] with value [" + value + "] into ctx");
+    }
+
+    private Object replaceSource(Object value) {
+        if (value instanceof Map == false) {
+            throw new IllegalArgumentException(
+                "Expected ["
+                    + SOURCE
+                    + "] to be a Map, not ["
+                    + value
+                    + "]"
+                    + (value != null ? " with type [" + value.getClass().getName() + "]" : "")
+            );
+
+        }
+        var oldSource = source;
+        source = castSourceMap(value);
+        return oldSource;
+    }
+
+    @SuppressWarnings({ "unchecked", "raw" })
+    private static Map<String, Object> castSourceMap(Object value) {
+        return (Map<String, Object>) value;
     }
 
     /**
@@ -108,7 +135,11 @@ public class CtxMap extends AbstractMap<String, Object> {
                 return metadata.remove(str);
             }
         }
-        return source.remove(key);
+        if (directSourceAccess()) {
+            return source.remove(key);
+        } else {
+            throw new UnsupportedOperationException("Cannot remove key " + key + " from ctx");
+        }
     }
 
     /**
@@ -121,28 +152,33 @@ public class CtxMap extends AbstractMap<String, Object> {
         for (String key : metadata.keySet()) {
             metadata.remove(key);
         }
+        // TODO: this is just bogus, there isn't any case where metadata won't trip a failure above?
         source.clear();
     }
 
     @Override
     public int size() {
         // uses map directly to avoid creating an EntrySet via AbstractMaps implementation, which returns entrySet().size()
-        return source.size() + metadata.size();
+        final int sourceSize = directSourceAccess() ? source.size() : 1;
+        return sourceSize + metadata.size();
     }
 
     @Override
     public boolean containsValue(Object value) {
         // uses map directly to avoid AbstractMaps linear time implementation using entrySet()
-        return metadata.containsValue(value) || source.containsValue(value);
+        return metadata.containsValue(value) || (directSourceAccess() ? source.containsValue(value) : source.equals(value));
     }
 
     @Override
     public boolean containsKey(Object key) {
         // uses map directly to avoid AbstractMaps linear time implementation using entrySet()
         if (key instanceof String str) {
-            return metadata.containsKey(str) || source.containsKey(key);
+            if (metadata.isAvailable(str)) {
+                return metadata.containsKey(str);
+            }
+            return directSourceAccess() ? source.containsKey(key) : SOURCE.equals(key);
         }
-        return source.containsKey(key);
+        return false;
     }
 
     @Override
@@ -153,7 +189,7 @@ public class CtxMap extends AbstractMap<String, Object> {
                 return metadata.get(str);
             }
         }
-        return source.get(key);
+        return directSourceAccess() ? source.get(key) : (SOURCE.equals(key) ? source : null);
     }
 
     /**
@@ -239,13 +275,37 @@ public class CtxMap extends AbstractMap<String, Object> {
                 throw new IllegalStateException();
             }
             if (sourceCur) {
-                sourceIter.remove();
+                try {
+                    sourceIter.remove();
+                } catch (UnsupportedOperationException e) {
+                    // UnsupportedOperationException's message is "remove", rethrowing with more helpful message
+                    throw new UnsupportedOperationException("Cannot remove key [" + cur.getKey() + "] from ctx");
+                }
+
             } else {
                 metadata.remove(cur.getKey());
             }
         }
     }
 
+    private class IndirectSourceEntry implements Map.Entry<String, Object> {
+
+        @Override
+        public String getKey() {
+            return SOURCE;
+        }
+
+        @Override
+        public Object getValue() {
+            return source;
+        }
+
+        @Override
+        public Object setValue(Object value) {
+            return replaceSource(value);
+        }
+    }
+
     /**
      * Map.Entry that stores metadata key and calls into {@link #metadata} for {@link #setValue}
      */

server/src/test/java/org/elasticsearch/action/update/UpdateRequestTests.java

@@ -361,16 +361,6 @@ public class UpdateRequestTests extends ESTestCase {
             IndexRequest indexAction = (IndexRequest) action;
             assertEquals(nowInMillis, indexAction.sourceAsMap().get("update_timestamp"));
         }
-        {
-            UpdateRequest updateRequest = new UpdateRequest("test", "2").upsert(indexRequest)
-                .script(mockInlineScript("ctx._timestamp = ctx._now"))
-                .scriptedUpsert(true);
-            // We simulate that the document is not existing yet
-            GetResult getResult = new GetResult("test", "2", 0, 1, 0, true, new BytesArray("{}"), null, null);
-            UpdateHelper.Result result = updateHelper.prepare(new ShardId("test", "_na_", 0), updateRequest, getResult, () -> 42L);
-            Writeable action = result.action();
-            assertThat(action, instanceOf(IndexRequest.class));
-        }
     }
 
     public void testIndexTimeout() {

server/src/test/java/org/elasticsearch/ingest/IngestCtxMapTests.java

@@ -104,6 +104,29 @@ public class IngestCtxMapTests extends ESTestCase {
         map = new IngestCtxMap(source, new IngestDocMetadata(metadata, null));
     }
 
+    public void testRemoveSource() {
+        Map<String, Object> source = new HashMap<>();
+        source.put("abc", 123);
+        source.put("def", 456);
+        source.put("hij", 789);
+        map = new IngestCtxMap(source, new TestIngestCtxMetadata(new HashMap<>(), new HashMap<>()));
+
+        // Make sure there isn't a ConcurrentModificationException when removing a key from the iterator
+        String removedKey = null;
+        for (Iterator<Map.Entry<String, Object>> it = map.entrySet().iterator(); it.hasNext();) {
+            Map.Entry<String, Object> entry = it.next();
+            String key = entry.getKey();
+            if (removedKey == null && source.containsKey(key)) {
+                removedKey = key;
+                it.remove();
+            }
+        }
+
+        assertNotNull(removedKey);
+        assertFalse(source.containsKey(removedKey));
+        assertTrue(source.containsKey(removedKey.equals("abc") ? "def" : "abc"));
+    }
+
     public void testRemove() {
         String cannotRemove = "cannotRemove";
         String canRemove = "canRemove";

server/src/test/java/org/elasticsearch/script/CtxMapTests.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.script;
+
+import org.elasticsearch.test.ESTestCase;
+import org.junit.Before;
+
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+
+import static org.hamcrest.Matchers.containsString;
+
+public class CtxMapTests extends ESTestCase {
+    CtxMap map;
+
+    @Override
+    @Before
+    public void setUp() throws Exception {
+        super.setUp();
+        map = new CtxMap(new HashMap<>(), new Metadata(new HashMap<>(), new HashMap<>()));
+    }
+
+    public void testAddingJunkToCtx() {
+        IllegalArgumentException err = expectThrows(IllegalArgumentException.class, () -> map.put("junk", "stuff"));
+        assertEquals("Cannot put key [junk] with value [stuff] into ctx", err.getMessage());
+    }
+
+    public void testRemovingSource() {
+        UnsupportedOperationException err = expectThrows(UnsupportedOperationException.class, () -> map.remove("_source"));
+        assertEquals("Cannot remove key _source from ctx", err.getMessage());
+        err = expectThrows(UnsupportedOperationException.class, () -> iteratorAtSource().remove());
+        assertEquals("Cannot remove key [_source] from ctx", err.getMessage());
+    }
+
+    @SuppressWarnings("unchecked")
+    public void testReplacingSource() {
+        map.put("_source", Map.of("abc", 123));
+        assertEquals(123, ((Map<String, Object>) map.get("_source")).get("abc"));
+        sourceEntry().setValue(Map.of("def", 456));
+        assertEquals(456, ((Map<String, Object>) map.get("_source")).get("def"));
+    }
+
+    public void testInvalidReplacementOfSource() {
+        IllegalArgumentException err = expectThrows(
+            IllegalArgumentException.class,
+            () -> map.put("_source", List.of(1, 2, "buckle my shoe"))
+        );
+        assertThat(
+            err.getMessage(),
+            containsString("Expected [_source] to be a Map, not [[1, 2, buckle my shoe]] with type [java.util.ImmutableCollections$ListN]")
+        );
+        err = expectThrows(IllegalArgumentException.class, () -> sourceEntry().setValue(List.of(1, 2, "buckle my shoe")));
+        assertThat(
+            err.getMessage(),
+            containsString("Expected [_source] to be a Map, not [[1, 2, buckle my shoe]] with type [java.util.ImmutableCollections$ListN]")
+        );
+    }
+
+    protected Map.Entry<String, Object> sourceEntry() {
+        for (Map.Entry<String, Object> entry : map.entrySet()) {
+            if ("_source".equals(entry.getKey())) {
+                return entry;
+            }
+        }
+        fail("no _source");
+        return null;
+    }
+
+    protected Iterator<Map.Entry<String, Object>> iteratorAtSource() {
+        Iterator<Map.Entry<String, Object>> it = map.entrySet().iterator();
+        while (it.hasNext()) {
+            if ("_source".equals(it.next().getKey())) {
+                return it;
+            }
+        }
+        fail("no _source");
+        return null;
+    }
+}