From 8e1a6155e24a7763f01477ba834cac660062b520 Mon Sep 17 00:00:00 2001
From: Stuart Tettemer <stu@elastic.co>
Date: Thu, 4 Mar 2021 13:09:08 -0600
Subject: [PATCH] Script: no compile rate limit for ingest templates (#69841)
 (#69987)

Remove the compilation rate limit for ingest templates.

Creates a new context, `ingest_template`, with an unlimited
compilation rate limit.

The `template` context is used in many places so it cannot be
exempted from rate limits.

Fixes: #64595
Backport: 9370b1c
---
 .../script/mustache/MustacheScriptEngine.java             | 2 +-
 .../main/java/org/elasticsearch/ingest/ValueSource.java   | 2 +-
 .../main/java/org/elasticsearch/script/ScriptModule.java  | 1 +
 .../java/org/elasticsearch/script/TemplateScript.java     | 8 ++++++++
 4 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/MustacheScriptEngine.java b/modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/MustacheScriptEngine.java
index 400f398f2196f..b2f683f9b1981 100644
--- a/modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/MustacheScriptEngine.java
+++ b/modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/MustacheScriptEngine.java
@@ -75,7 +75,7 @@ public <T> T compile(
 
     @Override
     public Set<ScriptContext<?>> getSupportedContexts() {
-        return Collections.singleton(TemplateScript.CONTEXT);
+        return org.elasticsearch.common.collect.Set.of(TemplateScript.CONTEXT, TemplateScript.INGEST_CONTEXT);
     }
 
     private CustomMustacheFactory createMustacheFactory(Map<String, String> options) {
diff --git a/server/src/main/java/org/elasticsearch/ingest/ValueSource.java b/server/src/main/java/org/elasticsearch/ingest/ValueSource.java
index 825bfc59a4e78..11f2b50db4918 100644
--- a/server/src/main/java/org/elasticsearch/ingest/ValueSource.java
+++ b/server/src/main/java/org/elasticsearch/ingest/ValueSource.java
@@ -71,7 +71,7 @@ static ValueSource wrap(Object value, ScriptService scriptService, Map<String, S
                 Script script = new Script(
                     ScriptType.INLINE, DEFAULT_TEMPLATE_LANG, (String) value, scriptOptions, org.elasticsearch.common.collect.Map.of()
                 );
-                return new TemplatedValue(scriptService.compile(script, TemplateScript.CONTEXT));
+                return new TemplatedValue(scriptService.compile(script, TemplateScript.INGEST_CONTEXT));
             } else {
                 return new ObjectValue(value);
             }
diff --git a/server/src/main/java/org/elasticsearch/script/ScriptModule.java b/server/src/main/java/org/elasticsearch/script/ScriptModule.java
index 8485ecd722e30..1c1094123470f 100644
--- a/server/src/main/java/org/elasticsearch/script/ScriptModule.java
+++ b/server/src/main/java/org/elasticsearch/script/ScriptModule.java
@@ -53,6 +53,7 @@ public class ScriptModule {
             SimilarityScript.CONTEXT,
             SimilarityWeightScript.CONTEXT,
             TemplateScript.CONTEXT,
+            TemplateScript.INGEST_CONTEXT,
             MovingFunctionScript.CONTEXT,
             ScriptedMetricAggContexts.InitScript.CONTEXT,
             ScriptedMetricAggContexts.MapScript.CONTEXT,
diff --git a/server/src/main/java/org/elasticsearch/script/TemplateScript.java b/server/src/main/java/org/elasticsearch/script/TemplateScript.java
index a22ca0abc58bc..6738df2536689 100644
--- a/server/src/main/java/org/elasticsearch/script/TemplateScript.java
+++ b/server/src/main/java/org/elasticsearch/script/TemplateScript.java
@@ -8,6 +8,8 @@
 
 package org.elasticsearch.script;
 
+import org.elasticsearch.common.unit.TimeValue;
+
 import java.util.Map;
 
 /**
@@ -35,4 +37,10 @@ public interface Factory {
     }
 
     public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("template", Factory.class);
+
+    // Remove compilation rate limit for ingest.  Ingest pipelines may use many mustache templates, triggering compilation
+    // rate limiting.  MustacheScriptEngine explicitly checks for TemplateScript.  Rather than complicating the implementation there by
+    // creating a new Script class (as would be customary), this context is used to avoid the default rate limit.
+    public static final ScriptContext<Factory> INGEST_CONTEXT = new ScriptContext<>("ingest_template", Factory.class,
+            200, TimeValue.timeValueMillis(0), ScriptCache.UNLIMITED_COMPILATION_RATE.asTuple());
 }