Skip to content

Commit

Permalink
fix Timer cannot cancle
Browse files Browse the repository at this point in the history
  • Loading branch information
@yoloyyh
yoloyyh committed Sep 25, 2024
1 parent 9de2b4c commit c2ff589
Show file tree
Hide file tree
Showing 3 changed files with 77 additions and 32 deletions.
35 changes: 19 additions & 16 deletions rasp/jvm/JVMAgent/src/main/java/com/security/smithloader/SmithAgent.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<Boolean> checkRecursive = new InheritableThreadLocal<Boolean>() {
@Override
protected Boolean initialValue() {
Expand Down Expand Up @@ -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) {
Expand Down
60 changes: 49 additions & 11 deletions rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbe.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<MatchRule> {
private final Trace trace;

Expand Down Expand Up @@ -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;
Expand All @@ -176,6 +202,7 @@ public SmithProbe() {
recordsTotal = new HashMap<>();
hooktimeRecords = new HashMap<>();
runtimeRecords = new HashMap<>();
isBenchMark = false;
}

public void setInst(Instrumentation inst) {
Expand Down Expand Up @@ -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);

Expand All @@ -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() {
Expand All @@ -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");
Expand All @@ -465,6 +500,9 @@ public void stop() {
detectTimerTask = null;
detectTimer =null;

benchMarkTimerTask = null;
benchMarkTimer = null;

smithproxyTimerTask = null;
smithproxyTimer = null;

Expand Down Expand Up @@ -608,7 +646,7 @@ private Long tp(List<Long> times, double percent) {
return times.get((int)(percent / 100 * times.size() - 1));
}

private void show() {
public void show() {
synchronized (records) {
SmithLogger.logger.info("=================== statistics ===================");

Expand Down
14 changes: 9 additions & 5 deletions rasp/jvm/JVMProbe/src/main/java/com/security/smith/asm/SmithMethodVisitor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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<String, Class<?>> smithProcessors = new HashMap<String, Class<?>>() {{
Expand Down Expand Up @@ -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("<init>");
isStatic = (access & Opcodes.ACC_STATIC) != 0;
Expand Down

0 comments on commit c2ff589

Please sign in to comment.