diff --git a/rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbe.java b/rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbe.java index 15660190a..669e69ad7 100644 --- a/rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbe.java +++ b/rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbe.java @@ -160,6 +160,8 @@ public class SmithProbe implements ClassFileTransformer, MessageHandler, EventHa private SmithProbe ourInstance = null; private SmithProbeProxy smithProxy = null; private int TRACE_BUFFER_SIZE = 1024; + private final int CLASS_MAX_ID = 50; + private final int METHOD_MAX_ID = 20; private Object xClassLoaderObj; private Boolean disable; @@ -177,7 +179,7 @@ public class SmithProbe implements ClassFileTransformer, MessageHandler, EventHa private final Map, List> recordsTotal; private final Map, Long> hooktimeRecords; private final Map, Long> runtimeRecords; - private Map> hookTypes; + private Set [][] hookTypes; private Disruptor disruptor; private Map switchConfig; @@ -244,10 +246,15 @@ public void init() { filters = new ConcurrentHashMap<>(); blocks = new ConcurrentHashMap<>(); limits = new ConcurrentHashMap<>(); - hookTypes = new ConcurrentHashMap<>(); switchConfig = new ConcurrentHashMap<>(); + hookTypes = new Set[CLASS_MAX_ID][METHOD_MAX_ID]; + for (int i = 0; i < CLASS_MAX_ID; i++) { + for (int j = 0; j < METHOD_MAX_ID; j++) { + hookTypes[i][j] = new HashSet<>(); + } + } MessageSerializer.initInstance(proberVersion); MessageEncoder.initInstance(); MessageDecoder.initInstance(); @@ -271,7 +278,7 @@ public Trace newInstance() { rulemgr = new Rule_Mgr(); ruleconfig = new Rule_Config(rulemgr); - smithProxy = new SmithProbeProxy(); + smithProxy = new SmithProbeProxy(CLASS_MAX_ID, METHOD_MAX_ID); InputStream inputStream = getResourceAsStream("class.yaml"); @@ -285,7 +292,7 @@ public Trace newInstance() { for (SmithMethod smithMethod : smithClass.getMethods()) { if (smithMethod.getTypes() != null && !smithMethod.getTypes().isEmpty()) - hookTypes.put(smithClass.getId() + "-" + smithMethod.getId(), smithMethod.getTypes()); + hookTypes[smithClass.getId()][smithMethod.getId()] = smithMethod.getTypes(); } smithClasses.put(smithClass.getName(), smithClass); } @@ -298,7 +305,7 @@ public Trace newInstance() { } try { - SmithLogger.logger.info("jsRuleEngine init"); + // SmithLogger.logger.info("jsRuleEngine init"); jsRuleEngine = JsRuleEngine.InitializeEngine(); if (jsRuleEngine != null) { SmithLogger.logger.info("jsRuleEngine init success"); @@ -314,16 +321,10 @@ public boolean addJsFile(Path scriptFilePath) { boolean ret = false; try { if (scriptFilePath != null && jsRuleEngine != null) { - SmithLogger.logger.info("add js rule enter"); int result = jsRuleEngine.addJsRule(scriptFilePath); if (result == 0) { - SmithLogger.logger.info("add js rule success"); ret = true; - } else { - SmithLogger.logger.info("add js rule failed, ret :" + result); } - } else { - SmithLogger.logger.info("not find js rule path: " + scriptFilePath); } } catch (Throwable e) { @@ -384,11 +385,10 @@ private boolean isBypassHookClass(String className) { public boolean isFunctionEnabled(int classId, int methodId) { - if (switchConfig == null || switchConfig.isEmpty()) { + if (switchConfig == null || switchConfig.isEmpty() || classId >= CLASS_MAX_ID || methodId >= METHOD_MAX_ID || hookTypes == null) { return true; } - String key = classId + "-" + methodId; - Set types = hookTypes.get(key); + Set types = hookTypes[classId][methodId]; if (types != null) { for (String type : types) { @@ -404,7 +404,6 @@ public void start() { SmithLogger.logger.info("probe start"); AttachInfo.info(); - SmithLogger.logger.info("init ClassUploadTransformer"); ClassUploadTransformer.getInstance().start(client, inst); @@ -451,6 +450,7 @@ public void start() { smithProxy.setDisruptor(disruptor); smithProxy.setProbe(this); + try { addJsRule(); } catch (Exception e) { @@ -470,32 +470,25 @@ public void stop() { inst.removeTransformer(this); reloadClasses(); - SmithLogger.logger.info("Transformer stop"); disable = true; scanswitch = false; ClassUploadTransformer.getInstance().stop(); - SmithLogger.logger.info("Upload Transformer stop"); detectTimer.cancel(); smithproxyTimer.cancel(); - SmithLogger.logger.info("detect Timer stop"); if (isBenchMark) { benchMarkTimer.cancel(); - SmithLogger.logger.info("benchMark Timer stop"); } client.stop(); - SmithLogger.logger.info("client stop"); ruleconfig.destry(); - SmithLogger.logger.info("ruleconfig stop"); rulemgr.destry(); - SmithLogger.logger.info("rulemgr stop"); detectTimerTask = null; detectTimer =null; @@ -537,11 +530,18 @@ public void uninit() { value.removeAll(); blocks.remove(key); } + for (int i = 0; i < hookTypes.length; i++) { + for (int j = 0; j < hookTypes[i].length; j++) { + hookTypes[i][j].clear(); + hookTypes[i][j] = null; + } + } + hookTypes = null; blocks.clear(); blocks = null; limits.clear(); limits = null; - SmithLogger.logger.info("probe uninit 0"); + disruptor = null; ruleconfig = null; @@ -1343,10 +1343,12 @@ public Disruptor getDisruptor() { public String getFuncTypes(int classId, int methodId) { String types = ""; + if (classId < 0 || methodId < 0 || classId >= CLASS_MAX_ID || methodId >= METHOD_MAX_ID ) { + return types; + } try { - - if (hookTypes.containsKey(classId + "-" + methodId)) { - for (String type: hookTypes.get(classId + "-" + methodId)) { + if (hookTypes[classId][methodId] != null) { + for (String type: hookTypes[classId][methodId]) { types += type + ","; } } diff --git a/rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbeProxy.java b/rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbeProxy.java index 43b6e745f..72819b371 100644 --- a/rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbeProxy.java +++ b/rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbeProxy.java @@ -26,8 +26,8 @@ import com.google.gson.JsonElement; import com.google.gson.GsonBuilder; public class SmithProbeProxy { - private final int CLASS_MAX_ID = 50; - private final int METHOD_MAX_ID = 20; + private final int CLASS_MAX_ID; + private final int METHOD_MAX_ID; private final int DEFAULT_QUOTA = 12000; private SmithProbe SmithProbeObj = null; @@ -117,9 +117,11 @@ private void RemoveThreadLocalVar() { } } - public SmithProbeProxy() { + public SmithProbeProxy(int classMaxID, int methodMaxID) { stopX = false; + CLASS_MAX_ID = classMaxID; + METHOD_MAX_ID = methodMaxID; quotas = new AtomicIntegerArray[CLASS_MAX_ID]; for (int i = 0; i < CLASS_MAX_ID; i++) { quotas[i] = new AtomicIntegerArray(METHOD_MAX_ID); @@ -280,7 +282,7 @@ public void sendMetadataClass(Class cla, int classID, int methodID) { JsRuleResult result = SmithProbeObj.getJsRuleEngine().detect(1,argsX); if(result != null) { - SmithLogger.logger.info("Js Rule Result +" + result.toString()); + // SmithLogger.logger.info("Js Rule Result +" + result.toString()); ClassFilter classFilter = new ClassFilter(); SmithHandler.queryClassFilter(cla, classFilter); classFilter.setTransId(); @@ -803,7 +805,7 @@ public void checkWildflyaddFilterPre(int classID, int methodID, Object[] args) { } public void handleReflectField(int classID, int methodID, Object[] args, Object ret, boolean blocked) { - if(stopX) { + if(stopX || SmithProbeObj.isFunctionEnabled(classID, methodID) == false) { return; } if (args.length < 2) {