diff --git a/rasp/jvm/JVMAgent/src/main/java/com/security/smithloader/SmithAgent.java b/rasp/jvm/JVMAgent/src/main/java/com/security/smithloader/SmithAgent.java index 64be383be..7a225d973 100644 --- a/rasp/jvm/JVMAgent/src/main/java/com/security/smithloader/SmithAgent.java +++ b/rasp/jvm/JVMAgent/src/main/java/com/security/smithloader/SmithAgent.java @@ -25,6 +25,8 @@ public class SmithAgent { private static String checksumStr = null; private static String proberPath = null; private static Instrumentation instrumentation = null; + // just fo benchmark test + private static boolean isBenchMark = false; public static InheritableThreadLocal checkRecursive = new InheritableThreadLocal() { @Override protected Boolean initialValue() { @@ -100,24 +102,25 @@ public static Object ExceptionProxy(Object MethodNameObj,int classID, int method } public static void RecordProxy(int classID, int methodID, Long t1, Long t2) { - try { - if (checkRecursive != null && checkRecursive.get() == true) { - return; - } - if (checkRecursive != null && checkRecursive.get() == false) { - checkRecursive.set(true); - } - if(SmithProberObj != null) { - Class[] argType = new Class[]{int.class,int.class, Long.class, Long.class}; - Reflection.invokeMethod(SmithProberObj,"record",argType,classID,methodID,t1,t2); - } - if (checkRecursive != null && checkRecursive.get() == true) { - checkRecursive.set(false); + if (isBenchMark) { + try { + if (checkRecursive != null && checkRecursive.get() == true) { + return; + } + if (checkRecursive != null && checkRecursive.get() == false) { + checkRecursive.set(true); + } + if(SmithProberObj != null) { + Class[] argType = new Class[]{int.class,int.class, Long.class, Long.class}; + Reflection.invokeMethod(SmithProberObj,"record",argType,classID,methodID,t1,t2); + } + if (checkRecursive != null && checkRecursive.get() == true) { + checkRecursive.set(false); + } + } catch (Throwable e) { + SmithAgentLogger.exception(e); } - } catch (Throwable e) { - SmithAgentLogger.exception(e); } - } private static boolean loadSmithProber(String proberPath, Instrumentation inst) { 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 c3cd983b7..aa112e5ac 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 @@ -115,6 +115,28 @@ public boolean cancel() { } } +class BenchMarkTimerTask extends TimerTask { + private boolean isCancel = false; + private SmithProbe smithProbe = null; + + public void setSmithProbe(SmithProbe smithProbe) { + this.smithProbe = smithProbe; + } + + @Override + public void run() { + if(!isCancel) { + smithProbe.show(); + } + } + + @Override + public boolean cancel() { + isCancel = true; + return super.cancel(); + } +} + class MatchRulePredicate implements Predicate { private final Trace trace; @@ -162,12 +184,16 @@ public class SmithProbe implements ClassFileTransformer, MessageHandler, EventHa private Rule_Mgr rulemgr; private Rule_Config ruleconfig; private Timer detectTimer; + private Timer benchMarkTimer; private Timer smithproxyTimer; private DetectTimerTask detectTimerTask; private SmithproxyTimerTask smithproxyTimerTask; + private BenchMarkTimerTask benchMarkTimerTask; private String proberVersion; private String proberPath; private JsRuleEngine jsRuleEngine; + // just for benchmark test + private boolean isBenchMark; public SmithProbe() { disable = false; @@ -176,6 +202,7 @@ public SmithProbe() { recordsTotal = new HashMap<>(); hooktimeRecords = new HashMap<>(); runtimeRecords = new HashMap<>(); + isBenchMark = false; } public void setInst(Instrumentation inst) { @@ -399,6 +426,18 @@ public void start() { TimeUnit.MINUTES.toMillis(1) ); + if (isBenchMark) { + benchMarkTimerTask = new BenchMarkTimerTask(); + benchMarkTimerTask.setSmithProbe(this); + + benchMarkTimer = new Timer(true); + benchMarkTimer.schedule( + benchMarkTimerTask, + TimeUnit.SECONDS.toMillis(5), + TimeUnit.SECONDS.toMillis(10) + ); + } + smithproxyTimerTask = new SmithproxyTimerTask(); smithproxyTimerTask.setSmithProxy(smithProxy); @@ -423,16 +462,7 @@ public void start() { reloadClasses(); SmithLogger.logger.info("probe start leave"); - new Timer(true).schedule( - new TimerTask() { - @Override - public void run() { - show(); - } - }, - TimeUnit.SECONDS.toMillis(5), - TimeUnit.SECONDS.toMillis(10) - ); + } public void stop() { @@ -452,6 +482,11 @@ public void 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"); @@ -465,6 +500,9 @@ public void stop() { detectTimerTask = null; detectTimer =null; + benchMarkTimerTask = null; + benchMarkTimer = null; + smithproxyTimerTask = null; smithproxyTimer = null; @@ -608,7 +646,7 @@ private Long tp(List times, double percent) { return times.get((int)(percent / 100 * times.size() - 1)); } - private void show() { + public void show() { synchronized (records) { SmithLogger.logger.info("=================== statistics ==================="); diff --git a/rasp/jvm/JVMProbe/src/main/java/com/security/smith/asm/SmithMethodVisitor.java b/rasp/jvm/JVMProbe/src/main/java/com/security/smith/asm/SmithMethodVisitor.java index 7116386a0..549f99fbd 100644 --- a/rasp/jvm/JVMProbe/src/main/java/com/security/smith/asm/SmithMethodVisitor.java +++ b/rasp/jvm/JVMProbe/src/main/java/com/security/smith/asm/SmithMethodVisitor.java @@ -23,8 +23,8 @@ public class SmithMethodVisitor extends AdviceAdapter { private final boolean canBlock; private final boolean isStatic; private final boolean isConstructor; - private final int stopWatchVariable; - private final int stopWatchTotalVariable; + private int stopWatchVariable; + private int stopWatchTotalVariable; private final int argumentsVariable; private final int returnVariable; private final Label start; @@ -34,6 +34,7 @@ public class SmithMethodVisitor extends AdviceAdapter { private String postHook; private String exceptionHook; private String xHook; + // just fo benchmark test private final boolean isBenchMark; private final Map> smithProcessors = new HashMap>() {{ @@ -61,15 +62,18 @@ protected SmithMethodVisitor(int api, Type classType, int classID, int methodID, this.preHook = pre_hook; this.postHook = post_hook; this.exceptionHook = exception_hook; - this.isBenchMark = true; + this.isBenchMark = false; start = new Label(); end = new Label(); handler = new Label(); argumentsVariable = newLocal(Type.getType(Object[].class)); - stopWatchTotalVariable = newLocal(Type.getType(StopWatch.class)); - stopWatchVariable = newLocal(Type.getType(StopWatch.class)); + if (isBenchMark) { + stopWatchTotalVariable = newLocal(Type.getType(StopWatch.class)); + stopWatchVariable = newLocal(Type.getType(StopWatch.class)); + } + returnVariable = newLocal(Type.getType(Object.class)); isConstructor = name.equals(""); isStatic = (access & Opcodes.ACC_STATIC) != 0;