Merge pull request #4236 from CoreFiling/js-policy-performance

[KEYCLOAK-5072] - Improve performance of JSPolicyProvider

authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/js/JSPolicyProvider.java

@@ -17,43 +17,44 @@
  */
 package org.keycloak.authorization.policy.provider.js;
 
-import java.util.function.Supplier;
-
-import javax.script.ScriptEngine;
-import javax.script.ScriptException;
+import java.util.function.BiFunction;
 
+import org.keycloak.authorization.AuthorizationProvider;
 import org.keycloak.authorization.model.Policy;
 import org.keycloak.authorization.policy.evaluation.Evaluation;
 import org.keycloak.authorization.policy.provider.PolicyProvider;
+import org.keycloak.scripting.EvaluatableScriptAdapter;
 
 /**
  * @author <a href="mailto:psilva@redhat.com">Pedro Igor</a>
  */
-public class JSPolicyProvider implements PolicyProvider {
+class JSPolicyProvider implements PolicyProvider {
 
-    private Supplier<ScriptEngine> engineProvider;
+    private final BiFunction<AuthorizationProvider, Policy, EvaluatableScriptAdapter> evaluatableScript;
 
-    public JSPolicyProvider(Supplier<ScriptEngine> engineProvider) {
-        this.engineProvider = engineProvider;
+    JSPolicyProvider(final BiFunction<AuthorizationProvider, Policy, EvaluatableScriptAdapter> evaluatableScript) {
+        this.evaluatableScript = evaluatableScript;
     }
 
     @Override
     public void evaluate(Evaluation evaluation) {
-        ScriptEngine engine = engineProvider.get();
-
-        engine.put("$evaluation", evaluation);
-
         Policy policy = evaluation.getPolicy();
+        AuthorizationProvider authorization = evaluation.getAuthorizationProvider();
+        final EvaluatableScriptAdapter adapter = evaluatableScript.apply(authorization, policy);
 
         try {
-            engine.eval(policy.getConfig().get("code"));
-        } catch (ScriptException e) {
+            //how to deal with long running scripts -> timeout?
+            adapter.eval(bindings -> {
+                bindings.put("script", adapter.getScriptModel());
+                bindings.put("$evaluation", evaluation);
+            });
+        }
+        catch (Exception e) {
             throw new RuntimeException("Error evaluating JS Policy [" + policy.getName() + "].", e);
         }
     }
 
     @Override
     public void close() {
-
     }
 }

authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/js/JSPolicyProviderFactory.java

@@ -1,9 +1,9 @@
 package org.keycloak.authorization.policy.provider.js;
 
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 
-import javax.script.ScriptEngineManager;
-
 import org.keycloak.Config;
 import org.keycloak.authorization.AuthorizationProvider;
 import org.keycloak.authorization.model.Policy;
@@ -11,17 +11,20 @@ import org.keycloak.authorization.policy.provider.PolicyProvider;
 import org.keycloak.authorization.policy.provider.PolicyProviderFactory;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.RealmModel;
+import org.keycloak.models.ScriptModel;
 import org.keycloak.representations.idm.authorization.JSPolicyRepresentation;
 import org.keycloak.representations.idm.authorization.PolicyRepresentation;
+import org.keycloak.scripting.EvaluatableScriptAdapter;
+import org.keycloak.scripting.ScriptingProvider;
 
 /**
  * @author <a href="mailto:psilva@redhat.com">Pedro Igor</a>
  */
 public class JSPolicyProviderFactory implements PolicyProviderFactory<JSPolicyRepresentation> {
 
-    private static final String ENGINE = "nashorn";
-
-    private JSPolicyProvider provider = new JSPolicyProvider(() -> new ScriptEngineManager().getEngineByName(ENGINE));
+    private final JSPolicyProvider provider = new JSPolicyProvider(this::getEvaluatableScript);
+    private final Map<String, EvaluatableScriptAdapter> scripts = Collections.synchronizedMap(new HashMap<>());
 
     @Override
     public String getName() {
@@ -69,8 +72,9 @@ public class JSPolicyProviderFactory implements PolicyProviderFactory<JSPolicyRe
         updatePolicy(policy, representation.getConfig().get("code"));
     }
 
-    private void updatePolicy(Policy policy, String code) {
-        policy.putConfig("code", code);
+    @Override
+    public void onRemove(final Policy policy, final AuthorizationProvider authorization) {
+        scripts.remove(policy.getId());
     }
 
     @Override
@@ -92,4 +96,25 @@ public class JSPolicyProviderFactory implements PolicyProviderFactory<JSPolicyRe
     public String getId() {
         return "js";
     }
+
+    private EvaluatableScriptAdapter getEvaluatableScript(final AuthorizationProvider authz, final Policy policy) {
+        return scripts.computeIfAbsent(policy.getId(), id -> {
+            final ScriptingProvider scripting = authz.getKeycloakSession().getProvider(ScriptingProvider.class);
+            ScriptModel script = getScriptModel(policy, authz.getRealm(), scripting);
+            return scripting.prepareEvaluatableScript(script);
+        });
+    }
+
+    private ScriptModel getScriptModel(final Policy policy, final RealmModel realm, final ScriptingProvider scripting) {
+        String scriptName = policy.getName();
+        String scriptCode = policy.getConfig().get("code");
+        String scriptDescription = policy.getDescription();
+
+        //TODO lookup script by scriptId instead of creating it every time
+        return scripting.createScript(realm.getId(), ScriptModel.TEXT_JAVASCRIPT, scriptName, scriptCode, scriptDescription);
+    }
+
+    private void updatePolicy(Policy policy, String code) {
+        policy.putConfig("code", code);
+    }
 }

server-spi-private/src/main/java/org/keycloak/scripting/EvaluatableScriptAdapter.java

@@ -0,0 +1,14 @@
+package org.keycloak.scripting;
+
+import org.keycloak.models.ScriptModel;
+
+/**
+ * Wraps a {@link ScriptModel} so it can be evaluated with custom bindings.
+ *
+ * @author <a href="mailto:jay@anslow.me.uk">Jay Anslow</a>
+ */
+public interface EvaluatableScriptAdapter {
+    ScriptModel getScriptModel();
+
+    Object eval(ScriptBindingsConfigurer bindingsConfigurer) throws ScriptExecutionException;
+}

server-spi-private/src/main/java/org/keycloak/scripting/InvocableScriptAdapter.java

@@ -56,7 +56,7 @@ public class InvocableScriptAdapter implements Invocable {
         }
 
         this.scriptModel = scriptModel;
-        this.scriptEngine = loadScriptIntoEngine(scriptModel, scriptEngine);
+        this.scriptEngine = scriptEngine;
     }
 
     @Override
@@ -101,17 +101,6 @@ public class InvocableScriptAdapter implements Invocable {
         return candidate != null;
     }
 
-    private ScriptEngine loadScriptIntoEngine(ScriptModel script, ScriptEngine engine) {
-
-        try {
-            engine.eval(script.getCode());
-        } catch (ScriptException se) {
-            throw new ScriptExecutionException(script, se);
-        }
-
-        return engine;
-    }
-
     private Invocable getInvocableEngine() {
         return (Invocable) scriptEngine;
     }

server-spi-private/src/main/java/org/keycloak/scripting/ScriptingProvider.java

@@ -38,6 +38,14 @@ public interface ScriptingProvider extends Provider {
      */
     InvocableScriptAdapter prepareInvocableScript(ScriptModel scriptModel, ScriptBindingsConfigurer bindingsConfigurer);
 
+    /**
+     * Returns an {@link EvaluatableScriptAdapter} based on the given {@link ScriptModel}.
+     * <p>The {@code EvaluatableScriptAdapter} wraps a dedicated {@link ScriptEngine} that was populated with empty bindings.</p>
+     *
+     * @param scriptModel the scriptModel to wrap
+     */
+    EvaluatableScriptAdapter prepareEvaluatableScript(ScriptModel scriptModel);
+
     /**
      * Creates a new {@link ScriptModel} instance.
      *

services/src/main/java/org/keycloak/authorization/admin/ResourceSetService.java

@@ -86,6 +86,7 @@ public class ResourceSetService {
     }
 
     @POST
+    @NoCache
     @Consumes("application/json")
     @Produces("application/json")
     public Response create(@Context UriInfo uriInfo, ResourceRepresentation resource) {
@@ -288,8 +289,8 @@ public class ResourceSetService {
 
     @Path("/search")
     @GET
-    @Produces("application/json")
     @NoCache
+    @Produces("application/json")
     public Response find(@QueryParam("name") String name) {
         this.auth.requireView();
         StoreFactory storeFactory = authorization.getStoreFactory();

services/src/main/java/org/keycloak/authorization/admin/ScopeService.java

@@ -77,6 +77,7 @@ public class ScopeService {
     }
 
     @POST
+    @NoCache
     @Consumes(MediaType.APPLICATION_JSON)
     @Produces(MediaType.APPLICATION_JSON)
     public Response create(@Context UriInfo uriInfo,  ScopeRepresentation scope) {
@@ -150,6 +151,7 @@ public class ScopeService {
 
     @Path("{id}")
     @GET
+    @NoCache
     @Produces(MediaType.APPLICATION_JSON)
     public Response findById(@PathParam("id") String id) {
         this.auth.requireView();
@@ -164,6 +166,7 @@ public class ScopeService {
 
     @Path("{id}/resources")
     @GET
+    @NoCache
     @Produces(MediaType.APPLICATION_JSON)
     public Response getResources(@PathParam("id") String id) {
         this.auth.requireView();
@@ -186,6 +189,7 @@ public class ScopeService {
 
     @Path("{id}/permissions")
     @GET
+    @NoCache
     @Produces(MediaType.APPLICATION_JSON)
     public Response getPermissions(@PathParam("id") String id) {
         this.auth.requireView();
@@ -231,6 +235,7 @@ public class ScopeService {
     }
 
     @GET
+    @NoCache
     @Produces("application/json")
     public Response findAll(@QueryParam("scopeId") String id,
                             @QueryParam("name") String name,

services/src/main/java/org/keycloak/scripting/AbstractEvaluatableScriptAdapter.java

@@ -0,0 +1,76 @@
+package org.keycloak.scripting;
+
+import javax.script.Bindings;
+import javax.script.ScriptContext;
+import javax.script.ScriptEngine;
+import javax.script.ScriptException;
+
+import org.keycloak.models.ScriptModel;
+
+/**
+ * Abstract class for wrapping a {@link ScriptModel} to make it evaluatable.
+ *
+ * @author <a href="mailto:jay@anslow.me.uk">Jay Anslow</a>
+ */
+abstract class AbstractEvaluatableScriptAdapter implements EvaluatableScriptAdapter {
+    /**
+     * Holds the {@link ScriptModel}.
+     */
+    private final ScriptModel scriptModel;
+
+    AbstractEvaluatableScriptAdapter(final ScriptModel scriptModel) {
+        if (scriptModel == null) {
+            throw new IllegalArgumentException("scriptModel must not be null");
+        }
+        this.scriptModel = scriptModel;
+    }
+
+    @Override
+    public Object eval(final ScriptBindingsConfigurer bindingsConfigurer) throws ScriptExecutionException {
+        return evalUnchecked(createBindings(bindingsConfigurer));
+    }
+
+    @Override
+    public ScriptModel getScriptModel() {
+        return scriptModel;
+    }
+
+    /**
+     * Note, calling this method modifies the underlying {@link ScriptEngine},
+     * preventing concurrent use of the ScriptEngine (Nashorn's {@link ScriptEngine} and
+     * {@link javax.script.CompiledScript} is thread-safe, but {@link Bindings} isn't).
+     */
+    InvocableScriptAdapter prepareInvokableScript(final ScriptBindingsConfigurer bindingsConfigurer) {
+        final Bindings bindings = createBindings(bindingsConfigurer);
+        evalUnchecked(bindings);
+        final ScriptEngine engine = getEngine();
+        engine.setBindings(bindings, ScriptContext.ENGINE_SCOPE);
+        return new InvocableScriptAdapter(scriptModel, engine);
+    }
+
+    protected String getCode() {
+        return scriptModel.getCode();
+    }
+
+    protected abstract ScriptEngine getEngine();
+
+    protected abstract Object eval(Bindings bindings) throws ScriptException;
+
+    private Object evalUnchecked(final Bindings bindings) {
+        try {
+            return eval(bindings);
+        }
+        catch (ScriptException e) {
+            throw new ScriptExecutionException(scriptModel, e);
+        }
+    }
+
+    private Bindings createBindings(final ScriptBindingsConfigurer bindingsConfigurer) {
+        if (bindingsConfigurer == null) {
+            throw new IllegalArgumentException("bindingsConfigurer must not be null");
+        }
+        final Bindings bindings = getEngine().createBindings();
+        bindingsConfigurer.configureBindings(bindings);
+        return bindings;
+    }
+}

services/src/main/java/org/keycloak/scripting/CompiledEvaluatableScriptAdapter.java

@@ -0,0 +1,40 @@
+package org.keycloak.scripting;
+
+import javax.script.Bindings;
+import javax.script.CompiledScript;
+import javax.script.ScriptEngine;
+import javax.script.ScriptException;
+
+import org.keycloak.models.ScriptModel;
+
+/**
+ * Wraps a compiled {@link ScriptModel} so it can be evaluated.
+ *
+ * @author <a href="mailto:jay@anslow.me.uk">Jay Anslow</a>
+ */
+class CompiledEvaluatableScriptAdapter extends AbstractEvaluatableScriptAdapter {
+    /**
+     * Holds the {@link CompiledScript} for the {@link ScriptModel}.
+     */
+    private final CompiledScript compiledScript;
+
+    CompiledEvaluatableScriptAdapter(final ScriptModel scriptModel, final CompiledScript compiledScript) {
+        super(scriptModel);
+
+        if (compiledScript == null) {
+            throw new IllegalArgumentException("compiledScript must not be null");
+        }
+
+        this.compiledScript = compiledScript;
+    }
+
+    @Override
+    protected ScriptEngine getEngine() {
+        return compiledScript.getEngine();
+    }
+
+    @Override
+    protected Object eval(final Bindings bindings) throws ScriptException {
+        return compiledScript.eval(bindings);
+    }
+}

services/src/main/java/org/keycloak/scripting/DefaultScriptingProvider.java

@@ -16,12 +16,14 @@
  */
 package org.keycloak.scripting;
 
-import org.keycloak.models.ScriptModel;
-
 import javax.script.Bindings;
-import javax.script.ScriptContext;
+import javax.script.Compilable;
+import javax.script.CompiledScript;
 import javax.script.ScriptEngine;
 import javax.script.ScriptEngineManager;
+import javax.script.ScriptException;
+
+import org.keycloak.models.ScriptModel;
 
 /**
  * A {@link ScriptingProvider} that uses a {@link ScriptEngineManager} to evaluate scripts with a {@link ScriptEngine}.
@@ -32,8 +34,7 @@ public class DefaultScriptingProvider implements ScriptingProvider {
 
     private final ScriptEngineManager scriptEngineManager;
 
-    public DefaultScriptingProvider(ScriptEngineManager scriptEngineManager) {
-
+    DefaultScriptingProvider(ScriptEngineManager scriptEngineManager) {
         if (scriptEngineManager == null) {
             throw new IllegalStateException("scriptEngineManager must not be null!");
         }
@@ -44,13 +45,22 @@ public class DefaultScriptingProvider implements ScriptingProvider {
     /**
      * Wraps the provided {@link ScriptModel} in a {@link javax.script.Invocable} instance with bindings configured through the {@link ScriptBindingsConfigurer}.
      *
-     * @param scriptModel  must not be {@literal null}
+     * @param scriptModel        must not be {@literal null}
      * @param bindingsConfigurer must not be {@literal null}
-     * @return
      */
     @Override
     public InvocableScriptAdapter prepareInvocableScript(ScriptModel scriptModel, ScriptBindingsConfigurer bindingsConfigurer) {
+        final AbstractEvaluatableScriptAdapter evaluatable = prepareEvaluatableScript(scriptModel);
+        return evaluatable.prepareInvokableScript(bindingsConfigurer);
+    }
 
+    /**
+     * Wraps the provided {@link ScriptModel} in a {@link javax.script.Invocable} instance with bindings configured through the {@link ScriptBindingsConfigurer}.
+     *
+     * @param scriptModel must not be {@literal null}
+     */
+    @Override
+    public AbstractEvaluatableScriptAdapter prepareEvaluatableScript(ScriptModel scriptModel) {
         if (scriptModel == null) {
             throw new IllegalArgumentException("script must not be null");
         }
@@ -59,13 +69,18 @@ public class DefaultScriptingProvider implements ScriptingProvider {
             throw new IllegalArgumentException("script must not be null or empty");
         }
 
-        if (bindingsConfigurer == null) {
-            throw new IllegalArgumentException("bindingsConfigurer must not be null");
+        ScriptEngine engine = createPreparedScriptEngine(scriptModel);
+
+        if (engine instanceof Compilable) {
+            try {
+                final CompiledScript compiledScript = ((Compilable) engine).compile(scriptModel.getCode());
+                return new CompiledEvaluatableScriptAdapter(scriptModel, compiledScript);
+            }
+            catch (ScriptException e) {
+                throw new ScriptExecutionException(scriptModel, e);
+            }
         }
-
-        ScriptEngine engine = createPreparedScriptEngine(scriptModel, bindingsConfigurer);
-
-        return new InvocableScriptAdapter(scriptModel, engine);
+        return new UncompiledEvaluatableScriptAdapter(scriptModel, engine);
     }
 
     //TODO allow scripts to be maintained independently of other components, e.g. with dedicated persistence
@@ -74,38 +89,27 @@ public class DefaultScriptingProvider implements ScriptingProvider {
 
     @Override
     public ScriptModel createScript(String realmId, String mimeType, String scriptName, String scriptCode, String scriptDescription) {
+        return new Script(null /* scriptId */, realmId, scriptName, mimeType, scriptCode, scriptDescription);
+    }
 
-        ScriptModel script = new Script(null /* scriptId */, realmId, scriptName, mimeType, scriptCode, scriptDescription);
-        return script;
+    @Override
+    public void close() {
+        //NOOP
     }
 
     /**
      * Looks-up a {@link ScriptEngine} with prepared {@link Bindings} for the given {@link ScriptModel Script}.
-     *
-     * @param script
-     * @param bindingsConfigurer
-     * @return
      */
-    private ScriptEngine createPreparedScriptEngine(ScriptModel script, ScriptBindingsConfigurer bindingsConfigurer) {
-
+    private ScriptEngine createPreparedScriptEngine(ScriptModel script) {
         ScriptEngine scriptEngine = lookupScriptEngineFor(script);
 
         if (scriptEngine == null) {
             throw new IllegalStateException("Could not find ScriptEngine for script: " + script);
         }
 
-        configureBindings(bindingsConfigurer, scriptEngine);
-
         return scriptEngine;
     }
 
-    private void configureBindings(ScriptBindingsConfigurer bindingsConfigurer, ScriptEngine engine) {
-
-        Bindings bindings = engine.createBindings();
-        bindingsConfigurer.configureBindings(bindings);
-        engine.setBindings(bindings, ScriptContext.ENGINE_SCOPE);
-    }
-
     /**
      * Looks-up a {@link ScriptEngine} based on the MIME-type provided by the given {@link Script}.
      */
@@ -114,13 +118,9 @@ public class DefaultScriptingProvider implements ScriptingProvider {
         try {
             Thread.currentThread().setContextClassLoader(DefaultScriptingProvider.class.getClassLoader());
             return scriptEngineManager.getEngineByMimeType(script.getMimeType());
-        } finally {
+        }
+        finally {
             Thread.currentThread().setContextClassLoader(cl);
         }
     }
-
-    @Override
-    public void close() {
-        //NOOP
-    }
 }

services/src/main/java/org/keycloak/scripting/UncompiledEvaluatableScriptAdapter.java

@@ -0,0 +1,39 @@
+package org.keycloak.scripting;
+
+import javax.script.Bindings;
+import javax.script.ScriptEngine;
+import javax.script.ScriptException;
+
+import org.keycloak.models.ScriptModel;
+
+/**
+ * Wraps an uncompiled {@link ScriptModel} so it can be evaluated.
+ *
+ * @author <a href="mailto:jay@anslow.me.uk">Jay Anslow</a>
+ */
+class UncompiledEvaluatableScriptAdapter extends AbstractEvaluatableScriptAdapter {
+    /**
+     * Holds the {@link ScriptEngine} instance.
+     */
+    private final ScriptEngine scriptEngine;
+
+    UncompiledEvaluatableScriptAdapter(final ScriptModel scriptModel, final ScriptEngine scriptEngine) {
+        super(scriptModel);
+        if (scriptEngine == null) {
+            throw new IllegalArgumentException("scriptEngine must not be null");
+        }
+
+        this.scriptEngine = scriptEngine;
+    }
+
+    @Override
+    protected ScriptEngine getEngine() {
+        return scriptEngine;
+    }
+
+    @Override
+    protected Object eval(final Bindings bindings) throws ScriptException {
+        return getEngine().eval(getCode(), bindings);
+    }
+
+}