Merge pull request #679 from PBH-BTN/scriptengine-refactor

抽离脚本引擎 + 添加 BTN 的脚本支持
PBH-BTN · Nov 3, 2024 · 26ae4e0 · 26ae4e0
2 parents 28cb4b2 + 4735648
commit 26ae4e0
Show file tree

Hide file tree

Showing 14 changed files with 417 additions and 184 deletions.
diff --git a/src/main/java/com/ghostchu/peerbanhelper/PeerBanHelperServer.java b/src/main/java/com/ghostchu/peerbanhelper/PeerBanHelperServer.java
@@ -27,12 +27,15 @@
 import com.ghostchu.peerbanhelper.module.impl.rule.*;
 import com.ghostchu.peerbanhelper.module.impl.webapi.*;
 import com.ghostchu.peerbanhelper.peer.Peer;
+import com.ghostchu.peerbanhelper.scriptengine.ScriptEngine;
 import com.ghostchu.peerbanhelper.text.Lang;
 import com.ghostchu.peerbanhelper.text.TranslationComponent;
 import com.ghostchu.peerbanhelper.torrent.Torrent;
 import com.ghostchu.peerbanhelper.util.*;
+import com.ghostchu.peerbanhelper.util.json.JsonUtil;
 import com.ghostchu.peerbanhelper.util.rule.ModuleMatchCache;
 import com.ghostchu.peerbanhelper.util.time.ExceptedTime;
+import com.ghostchu.peerbanhelper.util.time.InfoHashUtil;
 import com.ghostchu.peerbanhelper.util.time.TimeoutProtect;
 import com.ghostchu.peerbanhelper.web.JavalinWebContainer;
 import com.ghostchu.peerbanhelper.wrapper.BanMetadata;
@@ -44,10 +47,15 @@
 import com.google.common.cache.Cache;
 import com.google.common.cache.CacheBuilder;
 import com.google.gson.JsonObject;
+import com.googlecode.aviator.AviatorEvaluator;
+import com.googlecode.aviator.EvalMode;
+import com.googlecode.aviator.Options;
+import com.googlecode.aviator.runtime.JavaMethodReflectionFunctionMissing;
 import inet.ipaddr.IPAddress;
 import io.javalin.util.JavalinBindException;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.StringUtils;
 import org.bspfsystems.yamlconfiguration.configuration.ConfigurationSection;
 import org.bspfsystems.yamlconfiguration.configuration.MemoryConfiguration;
 import org.bspfsystems.yamlconfiguration.file.YamlConfiguration;
@@ -62,6 +70,7 @@
 import java.lang.management.ManagementFactory;
 import java.lang.management.ThreadInfo;
 import java.lang.management.ThreadMXBean;
+import java.math.MathContext;
 import java.sql.SQLException;
 import java.util.*;
 import java.util.concurrent.*;
@@ -128,6 +137,8 @@ public class PeerBanHelperServer implements Reloadable {
     private AlertManager alertManager;
     @Autowired
     private BanListDao banListDao;
+    @Getter
+    private ScriptEngine scriptEngine;
 
     public PeerBanHelperServer() {
         reloadConfig();
@@ -171,6 +182,7 @@ public void start() throws SQLException {
         log.info(tlUI(Lang.MOTD, Main.getMeta().getVersion()));
         loadDownloaders();
         registerBanListInvokers();
+        setupScriptEngine();
         registerModules();
         registerHttpServer();
         setupIPDB();
@@ -190,6 +202,46 @@ public void start() throws SQLException {
 
     }
 
+    private void setupScriptEngine() {
+        AviatorEvaluator.getInstance().setCachedExpressionByDefault(true);
+        // ASM 性能优先
+        AviatorEvaluator.getInstance().setOption(Options.EVAL_MODE, EvalMode.ASM);
+        // EVAL 性能优先
+        AviatorEvaluator.getInstance().setOption(Options.OPTIMIZE_LEVEL, AviatorEvaluator.EVAL);
+        // 降低浮点计算精度
+        AviatorEvaluator.getInstance().setOption(Options.MATH_CONTEXT, MathContext.DECIMAL32);
+        // 启用变量语法糖
+        AviatorEvaluator.getInstance().setOption(Options.ENABLE_PROPERTY_SYNTAX_SUGAR, true);
+//        // 表达式允许序列化和反序列化
+//        AviatorEvaluator.getInstance().setOption(Options.SERIALIZABLE, true);
+        // 用户规则写糊保护
+        AviatorEvaluator.getInstance().setOption(Options.MAX_LOOP_COUNT, 5000);
+        // 启用反射方法查找
+        AviatorEvaluator.getInstance().setFunctionMissing(JavaMethodReflectionFunctionMissing.getInstance());
+        // 注册反射调用
+        registerFunctions(IPAddressUtil.class);
+        registerFunctions(HTTPUtil.class);
+        registerFunctions(JsonUtil.class);
+        registerFunctions(Lang.class);
+        registerFunctions(StrUtil.class);
+        registerFunctions(PeerBanHelperServer.class);
+        registerFunctions(InfoHashUtil.class);
+        registerFunctions(Main.class);
+    }
+
+    private void registerFunctions(Class<?> clazz) {
+        try {
+            AviatorEvaluator.addInstanceFunctions(StringUtils.uncapitalize(clazz.getSimpleName()), clazz);
+        } catch (IllegalAccessException | NoSuchMethodException e) {
+            log.error("Internal error: failed on register instance functions: {}", clazz.getName(), e);
+        }
+        try {
+            AviatorEvaluator.addStaticFunctions(StringUtils.capitalize(clazz.getSimpleName()), clazz);
+        } catch (IllegalAccessException | NoSuchMethodException e) {
+            log.error("Internal error: failed on register static functions: {}", clazz.getName(), e);
+        }
+    }
+
     private void sendSnapshotAlert() {
         if (Main.getMeta().isSnapshotOrBeta()) {
             alertManager.publishAlert(false, AlertLevel.INFO, "unstable-alert", new TranslationComponent(Lang.ALERT_SNAPSHOT), new TranslationComponent(Lang.ALERT_SNAPSHOT_DESCRIPTION));

diff --git a/src/main/java/com/ghostchu/peerbanhelper/btn/BtnConfig.java b/src/main/java/com/ghostchu/peerbanhelper/btn/BtnConfig.java
@@ -1,6 +1,7 @@
 package com.ghostchu.peerbanhelper.btn;
 
 import com.ghostchu.peerbanhelper.PeerBanHelperServer;
+import com.ghostchu.peerbanhelper.scriptengine.ScriptEngine;
 import com.ghostchu.peerbanhelper.text.Lang;
 import com.ghostchu.peerbanhelper.util.rule.ModuleMatchCache;
 import lombok.extern.slf4j.Slf4j;
@@ -22,7 +23,8 @@ public class BtnConfig {
     private String userAgent;
     @Autowired
     private ModuleMatchCache matchCache;
-
+    @Autowired
+    private ScriptEngine scriptEngine;
     @Bean
     public BtnNetwork btnNetwork() {
         ConfigurationSection section = server.getMainConfig().getConfigurationSection("btn");
@@ -32,7 +34,7 @@ public BtnNetwork btnNetwork() {
             var submit = server.getMainConfig().getBoolean("btn.submit");
             var appId = server.getMainConfig().getString("btn.app-id");
             var appSecret = server.getMainConfig().getString("btn.app-secret");
-            BtnNetwork btnNetwork = new BtnNetwork(server, userAgent, configUrl, submit, appId, appSecret, matchCache);
+            BtnNetwork btnNetwork = new BtnNetwork(server, scriptEngine, userAgent, configUrl, submit, appId, appSecret, matchCache);
             log.info(tlUI(Lang.BTN_NETWORK_ENABLED));
             return btnNetwork;
         } else {

diff --git a/src/main/java/com/ghostchu/peerbanhelper/btn/BtnNetwork.java b/src/main/java/com/ghostchu/peerbanhelper/btn/BtnNetwork.java
@@ -3,6 +3,7 @@
 import com.ghostchu.peerbanhelper.PeerBanHelperServer;
 import com.ghostchu.peerbanhelper.btn.ability.*;
 import com.ghostchu.peerbanhelper.database.dao.impl.PeerRecordDao;
+import com.ghostchu.peerbanhelper.scriptengine.ScriptEngine;
 import com.ghostchu.peerbanhelper.text.Lang;
 import com.ghostchu.peerbanhelper.util.HTTPUtil;
 import com.ghostchu.peerbanhelper.util.rule.ModuleMatchCache;
@@ -35,6 +36,7 @@ public class BtnNetwork {
     private static final int PBH_BTN_PROTOCOL_IMPL_VERSION = 8;
     @Getter
     private final Map<Class<? extends BtnAbility>, BtnAbility> abilities = new HashMap<>();
+    private final ScriptEngine scriptEngine;
     @Getter
     private ScheduledExecutorService executeService = null;
     private String configUrl;
@@ -52,8 +54,9 @@ public class BtnNetwork {
     private PeerRecordDao peerRecordDao;
     private ModuleMatchCache moduleMatchCache;
 
-    public BtnNetwork(PeerBanHelperServer server, String userAgent, String configUrl, boolean submit, String appId, String appSecret, ModuleMatchCache moduleMatchCache) {
+    public BtnNetwork(PeerBanHelperServer server, ScriptEngine scriptEngine, String userAgent, String configUrl, boolean submit, String appId, String appSecret, ModuleMatchCache moduleMatchCache) {
         this.server = server;
+        this.scriptEngine = scriptEngine;
         this.userAgent = userAgent;
         this.configUrl = configUrl;
         this.submit = submit;
@@ -109,7 +112,7 @@ public void configBtnNetwork() {
 //                abilities.put(BtnAbilitySubmitRulesHitRate.class, new BtnAbilitySubmitRulesHitRate(this, ability.get("submit_hitrate").getAsJsonObject()));
 //            }
             if (ability.has("rules")) {
-                abilities.put(BtnAbilityRules.class, new BtnAbilityRules(this, ability.get("rules").getAsJsonObject()));
+                abilities.put(BtnAbilityRules.class, new BtnAbilityRules(this, scriptEngine, ability.get("rules").getAsJsonObject()));
             }
             if (ability.has("reconfigure")) {
                 abilities.put(BtnAbilityReconfigure.class, new BtnAbilityReconfigure(this, ability.get("reconfigure").getAsJsonObject()));

diff --git a/src/main/java/com/ghostchu/peerbanhelper/btn/BtnRule.java b/src/main/java/com/ghostchu/peerbanhelper/btn/BtnRule.java
@@ -22,4 +22,6 @@ public class BtnRule {
     private Map<String, List<String>> ipRules;
     @SerializedName("port")
     private Map<String, List<Integer>> portRules;
+    @SerializedName("script")
+    private Map<String, String> scriptRules;
 }
diff --git a/src/main/java/com/ghostchu/peerbanhelper/btn/BtnRuleParsed.java b/src/main/java/com/ghostchu/peerbanhelper/btn/BtnRuleParsed.java
@@ -1,5 +1,7 @@
 package com.ghostchu.peerbanhelper.btn;
 
+import com.ghostchu.peerbanhelper.scriptengine.CompiledScript;
+import com.ghostchu.peerbanhelper.scriptengine.ScriptEngine;
 import com.ghostchu.peerbanhelper.text.Lang;
 import com.ghostchu.peerbanhelper.text.TranslationComponent;
 import com.ghostchu.peerbanhelper.util.IPAddressUtil;
@@ -10,27 +12,53 @@
 import com.ghostchu.peerbanhelper.util.rule.matcher.IPMatcher;
 import inet.ipaddr.IPAddress;
 import lombok.Data;
+import lombok.extern.slf4j.Slf4j;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+import static com.ghostchu.peerbanhelper.text.TextManager.tlUI;
+
 @Data
+@Slf4j
 public class BtnRuleParsed {
+    private final ScriptEngine scriptEngine;
     private String version;
     private Map<String, List<Rule>> peerIdRules;
     private Map<String, List<Rule>> clientNameRules;
     private Map<String, List<Rule>> ipRules;
     private Map<String, List<Rule>> portRules;
+    private Map<String, CompiledScript> scriptRules;
 
-    public BtnRuleParsed(BtnRule btnRule) {
+    public BtnRuleParsed(ScriptEngine scriptEngine, BtnRule btnRule) {
+        this.scriptEngine = scriptEngine;
         this.version = btnRule.getVersion();
         this.ipRules = parseIPRule(btnRule.getIpRules());
         this.portRules = parsePortRule(btnRule.getPortRules());
         this.peerIdRules = parseRule(btnRule.getPeerIdRules());
         this.clientNameRules = parseRule(btnRule.getClientNameRules());
+        this.scriptRules = compileScripts(btnRule.getScriptRules());
+    }
+
+    private Map<String, CompiledScript> compileScripts(Map<String, String> scriptRules) {
+        Map<String, CompiledScript> scripts = new HashMap<>();
+        log.info(tlUI(Lang.BTN_RULES_SCRIPT_COMPILING, scriptRules.size()));
+        long startAt = System.currentTimeMillis();
+        scriptRules.forEach((name, content) -> {
+            try {
+                var script = scriptEngine.compileScript(null, name, content);
+                if (script != null) {
+                    scripts.put(name, script);
+                }
+            } catch (Exception e) {
+                log.error("Unable to load BTN script {}", name, e);
+            }
+        });
+        log.info(tlUI(Lang.BTN_RULES_SCRIPT_COMPILED, scripts.size(), System.currentTimeMillis() - startAt));
+        return scripts;
     }
 
     private Map<String, List<Rule>> parsePortRule(Map<String, List<Integer>> portRules) {
@@ -66,6 +94,7 @@ public String matcherIdentifier() {
         return rules;
     }
 
+
     public Map<String, List<Rule>> parseIPRule(Map<String, List<String>> raw) {
         Map<String, List<Rule>> rules = new HashMap<>();
         raw.forEach((k, v) -> rules.put(k, List.of(new BtnRuleIpMatcher(version, k, k, v.stream().map(IPAddressUtil::getIPAddress).toList()))));

diff --git a/src/main/java/com/ghostchu/peerbanhelper/btn/ability/BtnAbilityRules.java b/src/main/java/com/ghostchu/peerbanhelper/btn/ability/BtnAbilityRules.java
@@ -5,6 +5,7 @@
 import com.ghostchu.peerbanhelper.btn.BtnRule;
 import com.ghostchu.peerbanhelper.btn.BtnRuleParsed;
 import com.ghostchu.peerbanhelper.event.BtnRuleUpdateEvent;
+import com.ghostchu.peerbanhelper.scriptengine.ScriptEngine;
 import com.ghostchu.peerbanhelper.text.Lang;
 import com.ghostchu.peerbanhelper.text.TranslationComponent;
 import com.ghostchu.peerbanhelper.util.HTTPUtil;
@@ -34,12 +35,14 @@ public class BtnAbilityRules extends AbstractBtnAbility {
     private final String endpoint;
     private final long randomInitialDelay;
     private final File btnCacheFile = new File(Main.getDataDirectory(), "btn.cache");
+    private final ScriptEngine scriptEngine;
     @Getter
     private BtnRuleParsed btnRule;
 
 
-    public BtnAbilityRules(BtnNetwork btnNetwork, JsonObject ability) {
+    public BtnAbilityRules(BtnNetwork btnNetwork, ScriptEngine scriptEngine, JsonObject ability) {
         this.btnNetwork = btnNetwork;
+        this.scriptEngine = scriptEngine;
         this.interval = ability.get("interval").getAsLong();
         this.endpoint = ability.get("endpoint").getAsString();
         this.randomInitialDelay = ability.get("random_initial_delay").getAsLong();
@@ -55,7 +58,7 @@ private void loadCacheFile() throws IOException {
         } else {
             try {
                 BtnRule btnRule = JsonUtil.getGson().fromJson(Files.readString(btnCacheFile.toPath()), BtnRule.class);
-                this.btnRule = new BtnRuleParsed(btnRule);
+                this.btnRule = new BtnRuleParsed(scriptEngine, btnRule);
             } catch (Throwable ignored) {
             }
         }
@@ -110,7 +113,7 @@ private void updateRule() {
                     } else {
                         try {
                             BtnRule btr = JsonUtil.getGson().fromJson(r.body(), BtnRule.class);
-                            this.btnRule = new BtnRuleParsed(btr);
+                            this.btnRule = new BtnRuleParsed(scriptEngine, btr);
                             Main.getEventBus().post(new BtnRuleUpdateEvent());
                             try {
                                 Files.writeString(btnCacheFile.toPath(), r.body(), StandardCharsets.UTF_8);

diff --git a/src/main/java/com/ghostchu/peerbanhelper/module/RuleFeatureModule.java b/src/main/java/com/ghostchu/peerbanhelper/module/RuleFeatureModule.java
@@ -5,19 +5,20 @@
 import com.ghostchu.peerbanhelper.torrent.Torrent;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.List;
 import java.util.concurrent.ExecutorService;
 
 public interface RuleFeatureModule extends FeatureModule {
     /**
      * 检查一个特定的 Torrent 和 Peer 是否应该封禁
      *
      * @param torrent             Torrent
-     * @param peer                Peer
+     * @param peers                Peers
      * @param ruleExecuteExecutor 如果需要并发执行任务，请在给定的执行器中执行，以接受线程池的约束避免资源消耗失控
      * @return 规则检查结果
      */
     @NotNull
-    CheckResult shouldBanPeer(@NotNull Torrent torrent, @NotNull Peer peer, @NotNull Downloader downloader, @NotNull ExecutorService ruleExecuteExecutor);
+    CheckResult shouldBanPeer(@NotNull Torrent torrent, @NotNull Peer peers, @NotNull Downloader downloader, @NotNull ExecutorService ruleExecuteExecutor);
 
     /**
      * 指示模块的内部处理逻辑是否是线程安全的，如果线程不安全，PeerBanHelper 将在同步块中执行不安全的模块