From a04d79c5760eda51e410dd3289d9eb5c352f8efa Mon Sep 17 00:00:00 2001
From: Jay Anslow <jay@anslow.me.uk>
Date: Mon, 19 Jun 2017 16:53:07 +0100
Subject: [PATCH] Cache compiled scripts in JSPolicyProviderFactory

---
 .../policy/provider/js/JSPolicyProvider.java  | 33 +++++------------
 .../provider/js/JSPolicyProviderFactory.java  | 37 +++++++++++++++++--
 2 files changed, 44 insertions(+), 26 deletions(-)

diff --git a/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/js/JSPolicyProvider.java b/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/js/JSPolicyProvider.java
index 47992865ae..944ae02a90 100644
--- a/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/js/JSPolicyProvider.java
+++ b/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/js/JSPolicyProvider.java
@@ -17,27 +17,30 @@
  */
 package org.keycloak.authorization.policy.provider.js;
 
+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.models.RealmModel;
-import org.keycloak.models.ScriptModel;
 import org.keycloak.scripting.EvaluatableScriptAdapter;
-import org.keycloak.scripting.ScriptingProvider;
 
 /**
  * @author <a href="mailto:psilva@redhat.com">Pedro Igor</a>
  */
-public class JSPolicyProvider implements PolicyProvider {
+class JSPolicyProvider implements PolicyProvider {
+
+    private final BiFunction<AuthorizationProvider, Policy, EvaluatableScriptAdapter> evaluatableScript;
+
+    JSPolicyProvider(final BiFunction<AuthorizationProvider, Policy, EvaluatableScriptAdapter> evaluatableScript) {
+        this.evaluatableScript = evaluatableScript;
+    }
 
     @Override
     public void evaluate(Evaluation evaluation) {
         Policy policy = evaluation.getPolicy();
-
         AuthorizationProvider authorization = evaluation.getAuthorizationProvider();
-        ScriptModel script = getScriptModel(policy, authorization);
-        final EvaluatableScriptAdapter adapter = getScriptingProvider(authorization).prepareEvaluatableScript(script);
+        final EvaluatableScriptAdapter adapter = evaluatableScript.apply(authorization, policy);
 
         try {
             //how to deal with long running scripts -> timeout?
@@ -53,21 +56,5 @@ public class JSPolicyProvider implements PolicyProvider {
 
     @Override
     public void close() {
-
-    }
-
-    private ScriptModel getScriptModel(final Policy policy, final AuthorizationProvider authorization) {
-        String scriptName = policy.getName();
-        String scriptCode = policy.getConfig().get("code");
-        String scriptDescription = policy.getDescription();
-
-        RealmModel realm = authorization.getRealm();
-
-        //TODO lookup script by scriptId instead of creating it every time
-        return getScriptingProvider(authorization).createScript(realm.getId(), ScriptModel.TEXT_JAVASCRIPT, scriptName, scriptCode, scriptDescription);
-    }
-
-    private ScriptingProvider getScriptingProvider(final AuthorizationProvider authorization) {
-        return authorization.getKeycloakSession().getProvider(ScriptingProvider.class);
     }
 }
diff --git a/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/js/JSPolicyProviderFactory.java b/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/js/JSPolicyProviderFactory.java
index a26129654a..18dae2a2a0 100644
--- a/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/js/JSPolicyProviderFactory.java
+++ b/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/js/JSPolicyProviderFactory.java
@@ -1,5 +1,9 @@
 package org.keycloak.authorization.policy.provider.js;
 
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
 import org.keycloak.Config;
 import org.keycloak.authorization.AuthorizationProvider;
 import org.keycloak.authorization.model.Policy;
@@ -7,15 +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 JSPolicyProvider provider = new JSPolicyProvider();
+    private final JSPolicyProvider provider = new JSPolicyProvider(this::getEvaluatableScript);
+    private final Map<String, EvaluatableScriptAdapter> scripts = Collections.synchronizedMap(new HashMap<>());
 
     @Override
     public String getName() {
@@ -63,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
@@ -86,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);
+    }
 }