Fix Nashorn out of memory issue in adaptive script execution

components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/config/model/graph/BaseThreadLocalScriptEngineHolder.java

@@ -0,0 +1,102 @@
+/*
+ * Copyright (c) 2023, WSO2 LLC. (http://www.wso2.com).
+ *
+ * WSO2 LLC. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.wso2.carbon.identity.application.authentication.framework.config.model.graph;
+
+import java.util.Optional;
+
+import javax.script.Bindings;
+import javax.script.ScriptContext;
+import javax.script.ScriptEngine;
+import javax.script.SimpleScriptContext;
+
+/**
+ * Base class for ThreadLocal Implementation of the Script Engine.
+ * This holds different script engine instances between different
+ * threads. Thread safety is achieved by binding the script
+ * engine to the current running Thread.
+ */
+public abstract class BaseThreadLocalScriptEngineHolder {
+
+    private static final ThreadLocal<ScriptEngine> THREAD_LOCAL_SCRIPT_ENGINE = new ThreadLocal<>();
+    protected static final String[] NASHORN_ARGS = {"--no-java"};
+
+    public BaseThreadLocalScriptEngineHolder() {
+
+        init();
+    }
+
+    protected void init() {
+
+        getScriptEngine()
+                .ifPresent(this::clearScriptEngineBindings);
+
+        ScriptEngine scriptEngine = getScriptEngine()
+                .orElseGet(this::createScriptEngine);
+
+        initializeScriptContext(scriptEngine);
+
+        // Set the threadLocal Script Engine.
+        setScriptEngine(scriptEngine);
+    }
+
+    /**
+     * Get the thread local script engine.
+     *
+     * @return ScriptEngine
+     */
+    public Optional<ScriptEngine> getScriptEngine() {
+
+        ScriptEngine scriptEngine = THREAD_LOCAL_SCRIPT_ENGINE.get();
+        if (scriptEngine != null) {
+            return Optional.of(scriptEngine);
+        }
+        return Optional.empty();
+    }
+
+    private void setScriptEngine(ScriptEngine scriptEngine) {
+
+        THREAD_LOCAL_SCRIPT_ENGINE.set(scriptEngine);
+    }
+
+    private void initializeScriptContext(ScriptEngine scriptEngine) {
+
+        ScriptContext scriptContext = new SimpleScriptContext();
+
+        // Set the engine scope bindings.
+        Bindings engineScopeBindings = scriptEngine.createBindings();
+        scriptContext.setBindings(engineScopeBindings, ScriptContext.ENGINE_SCOPE);
+
+        // Set the global scope bindings.
+        Bindings globalScopeBindings = scriptEngine.createBindings();
+        scriptContext.setBindings(globalScopeBindings, ScriptContext.GLOBAL_SCOPE);
+
+        scriptEngine.setContext(scriptContext);
+    }
+
+    private void clearScriptEngineBindings(ScriptEngine scriptEngine) {
+
+        // Clear the existing bindings.
+        scriptEngine.getBindings(ScriptContext.ENGINE_SCOPE).clear();
+        scriptEngine.getBindings(ScriptContext.GLOBAL_SCOPE).clear();
+    }
+
+    protected abstract ClassLoader getClassLoader();
+
+    protected abstract ScriptEngine createScriptEngine();
+}

components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/config/model/graph/JsGraphBuilderFactory.java

@@ -30,9 +30,12 @@
 import org.wso2.carbon.identity.application.authentication.framework.handler.sequence.impl.SelectAcrFromFunction;
 import org.wso2.carbon.identity.application.authentication.framework.handler.sequence.impl.SelectOneFunction;
 import org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants;
+import org.wso2.carbon.identity.base.IdentityRuntimeException;
+import org.wso2.carbon.identity.core.util.IdentityUtil;
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 
 import javax.script.Bindings;
 import javax.script.ScriptContext;
@@ -54,11 +57,13 @@ public class JsGraphBuilderFactory implements JsBaseGraphBuilderFactory {
     @SuppressWarnings("removal")
     private NashornScriptEngineFactory factory;
 
+    private static boolean useThreadLocalScriptEngine = false;
 
     public void init() {
 
         factory = new NashornScriptEngineFactory();
         classFilter = new RestrictedClassFilter();
+        setUseThreadLocalScriptEngine();
     }
 
     public static void restoreCurrentContext(AuthenticationContext context, ScriptEngine engine)
@@ -88,17 +93,26 @@ public static void persistCurrentContext(AuthenticationContext context, ScriptEn
 
     public ScriptEngine createEngine(AuthenticationContext authenticationContext) {
 
-        ScriptEngine engine = factory.getScriptEngine(NASHORN_ARGS, getClassLoader(), classFilter);
-        Bindings bindings = engine.createBindings();
-        engine.setBindings(bindings, ScriptContext.GLOBAL_SCOPE);
-        engine.setBindings(engine.createBindings(), ScriptContext.ENGINE_SCOPE);
+        ScriptEngine engine;
+        Bindings globalBindings;
+        if (useThreadLocalScriptEngine) {
+            Optional<ScriptEngine> optionalScriptEngine = new ThreadLocalScriptEngineHolder().getScriptEngine();
+            engine = optionalScriptEngine.orElseThrow(
+                    () -> new IdentityRuntimeException("Failed to create a script engine"));
+            globalBindings = engine.getBindings(ScriptContext.GLOBAL_SCOPE);
+        } else {
+            engine = factory.getScriptEngine(NASHORN_ARGS, getClassLoader(), classFilter);
+            globalBindings = engine.createBindings();
+            engine.setBindings(globalBindings, ScriptContext.GLOBAL_SCOPE);
+            engine.setBindings(engine.createBindings(), ScriptContext.ENGINE_SCOPE);
+        }
         SelectAcrFromFunction selectAcrFromFunction = new SelectAcrFromFunction();
 //        todo move to functions registry
-        bindings.put(FrameworkConstants.JSAttributes.JS_FUNC_SELECT_ACR_FROM,
-            (SelectOneFunction) selectAcrFromFunction::evaluate);
+        globalBindings.put(FrameworkConstants.JSAttributes.JS_FUNC_SELECT_ACR_FROM,
+                (SelectOneFunction) selectAcrFromFunction::evaluate);
 
         JsLogger jsLogger = new JsLogger();
-        bindings.put(FrameworkConstants.JSAttributes.JS_LOG, jsLogger);
+        globalBindings.put(FrameworkConstants.JSAttributes.JS_LOG, jsLogger);
         return engine;
     }
 
@@ -132,4 +146,11 @@ public JsNashornGraphBuilder createBuilder(AuthenticationContext authenticationC
         return new JsNashornGraphBuilder(authenticationContext, stepConfigMap,
                 createEngine(authenticationContext), currentNode);
     }
+
+    private void setUseThreadLocalScriptEngine() {
+
+        useThreadLocalScriptEngine =
+                Boolean.parseBoolean(IdentityUtil.getProperty(FrameworkConstants.THREAD_LOCAL_SCRIPT_ENGINE_CONFIG));
+    }
+
 }

components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/config/model/graph/ThreadLocalScriptEngineHolder.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) 2023, WSO2 LLC. (http://www.wso2.com).
+ *
+ * WSO2 LLC. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.wso2.carbon.identity.application.authentication.framework.config.model.graph;
+
+import jdk.nashorn.api.scripting.ClassFilter;
+import jdk.nashorn.api.scripting.NashornScriptEngineFactory;
+
+import javax.script.ScriptEngine;
+
+/**
+ * ThreadLocal Implementation of the deprecated Nashorn Script Engine.
+ * This holds different script engine instances for different
+ * threads. Thread safety is achieved by binding the script
+ * engine to the current running Thread.
+ */
+public class ThreadLocalScriptEngineHolder extends BaseThreadLocalScriptEngineHolder {
+
+    protected ClassLoader getClassLoader() {
+
+        ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
+        return classLoader == null ? NashornScriptEngineFactory.class.getClassLoader() : classLoader;
+    }
+
+    protected ScriptEngine createScriptEngine() {
+
+        NashornScriptEngineFactory factory = new NashornScriptEngineFactory();
+        ClassFilter classFilter = new RestrictedClassFilter();
+        return factory.getScriptEngine(NASHORN_ARGS, getClassLoader(), classFilter);
+    }
+
+}

components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/config/model/graph/openjdk/nashorn/JsOpenJdkNashornGraphBuilderFactory.java

@@ -34,9 +34,12 @@
 import org.wso2.carbon.identity.application.authentication.framework.handler.sequence.impl.OpenJdkSelectAcrFromFunction;
 import org.wso2.carbon.identity.application.authentication.framework.handler.sequence.impl.SelectOneFunction;
 import org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants;
+import org.wso2.carbon.identity.base.IdentityRuntimeException;
+import org.wso2.carbon.identity.core.util.IdentityUtil;
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 
 import javax.script.Bindings;
 import javax.script.ScriptContext;
@@ -58,16 +61,17 @@ public class JsOpenJdkNashornGraphBuilderFactory implements JsBaseGraphBuilderFa
     // Suppress the Nashorn deprecation warnings in jdk 11
     @SuppressWarnings("removal")
     private NashornScriptEngineFactory factory;
-
+    private boolean useThreadLocalScriptEngine = false;
 
     public void init() {
 
         factory = new NashornScriptEngineFactory();
         classFilter = new OpenJdkNashornRestrictedClassFilter();
+        setUseThreadLocalScriptEngine();
     }
 
     public static void restoreCurrentContext(AuthenticationContext context, ScriptEngine engine)
-        throws FrameworkException {
+            throws FrameworkException {
 
         Map<String, Object> map = (Map<String, Object>) context.getProperty(JS_BINDING_CURRENT_CONTEXT);
         Bindings bindings = engine.getBindings(ScriptContext.ENGINE_SCOPE);
@@ -96,17 +100,27 @@ public static void persistCurrentContext(AuthenticationContext context, ScriptEn
 
     public ScriptEngine createEngine(AuthenticationContext authenticationContext) {
 
-        ScriptEngine engine = factory.getScriptEngine(NASHORN_ARGS, getClassLoader(), classFilter);
-        Bindings bindings = engine.createBindings();
-        engine.setBindings(bindings, ScriptContext.GLOBAL_SCOPE);
-        engine.setBindings(engine.createBindings(), ScriptContext.ENGINE_SCOPE);
+        ScriptEngine engine;
+        Bindings globalBindings;
+        if (useThreadLocalScriptEngine) {
+            Optional<ScriptEngine> optionalScriptEngine =
+                    new OpenJdkNashornThreadLocalScriptEngineHolder().getScriptEngine();
+            engine = optionalScriptEngine.orElseThrow(
+                    () -> new IdentityRuntimeException("Script engine is not available"));
+            globalBindings = engine.getBindings(ScriptContext.GLOBAL_SCOPE);
+        } else {
+            engine = factory.getScriptEngine(NASHORN_ARGS, getClassLoader(), classFilter);
+            globalBindings = engine.createBindings();
+            engine.setBindings(globalBindings, ScriptContext.GLOBAL_SCOPE);
+            engine.setBindings(engine.createBindings(), ScriptContext.ENGINE_SCOPE);
+        }
         OpenJdkSelectAcrFromFunction selectAcrFromFunction = new OpenJdkSelectAcrFromFunction();
 //        todo move to functions registry
-        bindings.put(FrameworkConstants.JSAttributes.JS_FUNC_SELECT_ACR_FROM,
-            (SelectOneFunction) selectAcrFromFunction::evaluate);
+        globalBindings.put(FrameworkConstants.JSAttributes.JS_FUNC_SELECT_ACR_FROM,
+                (SelectOneFunction) selectAcrFromFunction::evaluate);
 
         JsLogger jsLogger = new JsLogger();
-        bindings.put(FrameworkConstants.JSAttributes.JS_LOG, jsLogger);
+        globalBindings.put(FrameworkConstants.JSAttributes.JS_LOG, jsLogger);
         return engine;
     }
 
@@ -141,4 +155,10 @@ public JsBaseGraphBuilder getCurrentBuilder() {
 
         return JsOpenJdkNashornGraphBuilder.getCurrentBuilder();
     }
+
+    private void setUseThreadLocalScriptEngine() {
+
+        useThreadLocalScriptEngine =
+                Boolean.parseBoolean(IdentityUtil.getProperty(FrameworkConstants.THREAD_LOCAL_SCRIPT_ENGINE_CONFIG));
+    }
 }

components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/config/model/graph/openjdk/nashorn/OpenJdkNashornThreadLocalScriptEngineHolder.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) 2023, WSO2 LLC. (http://www.wso2.com).
+ *
+ * WSO2 LLC. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.wso2.carbon.identity.application.authentication.framework.config.model.graph.openjdk.nashorn;
+
+import org.openjdk.nashorn.api.scripting.ClassFilter;
+import org.openjdk.nashorn.api.scripting.NashornScriptEngineFactory;
+import org.wso2.carbon.identity.application.authentication.framework.config.model.graph.BaseThreadLocalScriptEngineHolder;
+
+import javax.script.ScriptEngine;
+
+/**
+ * ThreadLocal Implementation of the OpenJdk Nashorn Script Engine.
+ * This holds different script engine instances between different
+ * threads. Thread safety is achieved by binding the script
+ * engine to the current running Thread.
+ */
+public class OpenJdkNashornThreadLocalScriptEngineHolder extends BaseThreadLocalScriptEngineHolder {
+
+    protected ClassLoader getClassLoader() {
+
+        ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
+        return classLoader == null ? NashornScriptEngineFactory.class.getClassLoader() : classLoader;
+    }
+
+    protected ScriptEngine createScriptEngine() {
+
+        NashornScriptEngineFactory factory = new NashornScriptEngineFactory();
+        ClassFilter classFilter = new OpenJdkNashornRestrictedClassFilter();
+        return factory.getScriptEngine(NASHORN_ARGS, getClassLoader(), classFilter);
+    }
+
+}

components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkConstants.java

@@ -217,6 +217,7 @@ public abstract class FrameworkConstants {
     public static final String ORACLE = "Oracle";
 
     public static final String SCRIPT_ENGINE_CONFIG = "AdaptiveAuth.ScriptEngine";
+    public static final String THREAD_LOCAL_SCRIPT_ENGINE_CONFIG = "AdaptiveAuth.LimitScriptEngineCreation";
     public static final String OPENJDK_NASHORN = "openjdkNashorn";
     public static final String NASHORN = "nashorn";
 

features/identity-core/org.wso2.carbon.identity.core.server.feature/resources/identity.xml.j2

@@ -1157,6 +1157,7 @@
     -->
     <AdaptiveAuth>
    	    <ScriptEngine>{{AdaptiveAuth.ScriptEngine}}</ScriptEngine>
+   	    <LimitScriptEngineCreation>{{AdaptiveAuth.LimitScriptEngineCreation}}</LimitScriptEngineCreation>
     </AdaptiveAuth>
 
     <MultifactorAuthentication>