From eda3bb267638743262eabce1bb6d293d029b198f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=83=AD=E8=89=B3=E5=90=9B?= Date: Mon, 19 Feb 2024 14:34:55 +0800 Subject: [PATCH 1/7] Add run time statistics code --- .../java/com/security/smith/SmithProbe.java | 89 ++++++++++ .../smith/asm/SmithMethodVisitor.java | 165 +++++++++++++++++- 2 files changed, 252 insertions(+), 2 deletions(-) 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 1cc303482..e355a36bf 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 @@ -151,6 +151,10 @@ public class SmithProbe implements ClassFileTransformer, MessageHandler, EventHa private Map, Filter> filters; private Map, Block> blocks; private Map, Integer> limits; + private final Map, List> records; + private final Map, List> recordsTotal; + private final Map, Long> hooktimeRecords; + private final Map, Long> runtimeRecords; private Map> hookTypes; private Disruptor disruptor; private Map switchConfig; @@ -211,6 +215,10 @@ public void init() { limits = new ConcurrentHashMap<>(); hookTypes = new ConcurrentHashMap<>(); switchConfig = new ConcurrentHashMap<>(); + records = new ConcurrentHashMap<>(); + recordsTotal = new ConcurrentHashMap<>(); + hooktimeRecords = new ConcurrentHashMap<>(); + runtimeRecords = new ConcurrentHashMap<>(); MessageSerializer.initInstance(proberVersion); MessageEncoder.initInstance(); @@ -507,6 +515,28 @@ public void uninit() { SmithLogger.loggerProberUnInit(); } + new Timer(true).schedule( + new TimerTask() { + @Override + public void run() { + onTimer(); + } + }, + 0, + TimeUnit.MINUTES.toMillis(1) + ); + + new Timer(true).schedule( + new TimerTask() { + @Override + public void run() { + show(); + } + }, + TimeUnit.SECONDS.toMillis(5), + TimeUnit.SECONDS.toMillis(10) + ); +} private void reloadClasses() { reloadClasses(smithClasses.keySet()); @@ -584,6 +614,65 @@ private void reloadClasses(Collection classes) { } } + private Long tp(List times, double percent) { + return times.get((int)(percent / 100 * times.size() - 1)); + } + + private void show() { + synchronized (records) { + SmithLogger.logger.info("=================== statistics ==================="); + + records.forEach((k, v) -> { + Collections.sort(v); + + List tv = recordsTotal.get(new ImmutablePair<>(k.getLeft(), k.getRight())); + + Collections.sort(tv); + + Long hooktime = hooktimeRecords.get(new ImmutablePair<>(k.getLeft(), k.getRight())); + Long runtime = runtimeRecords.get(new ImmutablePair<>(k.getLeft(), k.getRight())); + + SmithLogger.logger.info( + String.format( + "class: %d method: %d count: %d tp50: %d tp90: %d tp95: %d tp99: %d tp99.99: %d max: %d total-max:%d hooktime:%d runtime:%d", + k.getLeft(), + k.getRight(), + v.size(), + tp(v, 50), + tp(v, 90), + tp(v, 95), + tp(v, 99), + tp(v, 99.99), + v.get(v.size() - 1), + tv.get(tv.size() - 1), + hooktime, + runtime + ) + ); + }); + } + } + + public void record(int classID, int methodID, long time,long totaltime) { + synchronized (records) { + records.computeIfAbsent(new ImmutablePair<>(classID, methodID), k -> new ArrayList<>()).add(time); + } + + synchronized (recordsTotal) { + recordsTotal.computeIfAbsent(new ImmutablePair<>(classID, methodID), k -> new ArrayList<>()).add(totaltime); + } + + synchronized (hooktimeRecords) { + hooktimeRecords.computeIfAbsent(new ImmutablePair<>(classID, methodID), k -> time); + hooktimeRecords.computeIfPresent(new ImmutablePair<>(classID, methodID),(k,v) -> v+time); + } + + synchronized (runtimeRecords) { + runtimeRecords.computeIfAbsent(new ImmutablePair<>(classID, methodID), k -> totaltime); + runtimeRecords.computeIfPresent(new ImmutablePair<>(classID, methodID),(k,v) -> v+totaltime); + } + } + @Override public void onEvent(Trace trace, long sequence, boolean endOfBatch) { Filter filter = filters.get(new ImmutablePair<>(trace.getClassID(), trace.getMethodID())); 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 e5ed9da71..1e25e13dc 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 @@ -3,6 +3,7 @@ //import com.security.smith.SmithProbeProxy; import com.security.smith.processor.*; import org.apache.commons.lang3.ArrayUtils; +import org.apache.commons.lang3.time.StopWatch; import org.objectweb.asm.Label; import org.objectweb.asm.MethodVisitor; import org.objectweb.asm.Opcodes; @@ -22,8 +23,11 @@ public class SmithMethodVisitor extends AdviceAdapter { private final boolean canBlock; private final boolean isStatic; private final boolean isConstructor; - private final int returnVariable; + private final int stopWatchVariable; + private final int stopWatchTotalVariable; private final int argumentsVariable; + private final int traceVariable; + private final int returnVariable; private final Label start; private final Label end; private final Label handler; @@ -63,6 +67,10 @@ protected SmithMethodVisitor(int api, Type classType, int classID, int methodID, handler = new Label(); argumentsVariable = newLocal(Type.getType(Object[].class)); + stopWatchTotalVariable = newLocal(Type.getType(StopWatch.class)); + stopWatchVariable = newLocal(Type.getType(StopWatch.class)); + argumentsVariable = newLocal(Type.getType(Object[].class)); + traceVariable = newLocal(Type.getType(Trace.class)); returnVariable = newLocal(Type.getType(Object.class)); isConstructor = name.equals(""); @@ -101,6 +109,28 @@ protected void onMethodEnter() { visitTryCatchBlock(start, end, handler, Type.getInternalName(Exception.class)); + invokeStatic( + Type.getType(StopWatch.class), + new Method( + "createStarted", + Type.getType(StopWatch.class), + new Type[]{} + ) + ); + + storeLocal(stopWatchTotalVariable); + + invokeStatic( + Type.getType(StopWatch.class), + new Method( + "createStarted", + Type.getType(StopWatch.class), + new Type[]{} + ) + ); + + storeLocal(stopWatchVariable); + loadArgArray(); storeLocal(argumentsVariable); @@ -115,6 +145,28 @@ protected void onMethodEnter() { } else { preHook = "detect"; } + if (!canBlock) { + Label label = new Label(); + + surplus(); + ifZCmp(EQ, label); + + newTrace(); + storeLocal(traceVariable); + + loadLocal(stopWatchVariable); + + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "suspend", + Type.VOID_TYPE, + new Type[]{} + ) + ); + + mark(label); + return; } push(preHook); @@ -135,6 +187,27 @@ protected void onMethodEnter() { } ) ); + + Label label = new Label(); + + surplus(); + ifZCmp(NE, label); + + visitInsn(ACONST_NULL); + storeLocal(traceVariable); + + mark(label); + + loadLocal(stopWatchVariable); + + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "suspend", + Type.VOID_TYPE, + new Type[]{} + ) + ); } @Override @@ -143,7 +216,22 @@ protected void onMethodExit(int opcode) { if (opcode == ATHROW) { return; - } + + loadLocal(stopWatchVariable); + + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "resume", + Type.VOID_TYPE, + new Type[]{} + ) + ); + + Label label = new Label(); + + loadLocal(traceVariable); + ifNull(label); Type returnType = Type.getReturnType(methodDesc); @@ -195,6 +283,79 @@ protected void onMethodExit(int opcode) { } ) ); + + post(); + mark(label); + + loadLocal(stopWatchVariable); + + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "stop", + Type.VOID_TYPE, + new Type[]{} + ) + ); + + loadLocal(stopWatchTotalVariable); + + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "stop", + Type.VOID_TYPE, + new Type[]{} + ) + ); + + invokeStatic( + Type.getType(SmithProbe.class), + new Method( + "getInstance", + Type.getType(SmithProbe.class), + new Type[]{} + ) + ); + + push(classID); + push(methodID); + + loadLocal(stopWatchVariable); + + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "getNanoTime", + Type.LONG_TYPE, + new Type[]{} + ) + ); + + loadLocal(stopWatchTotalVariable); + + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "getNanoTime", + Type.LONG_TYPE, + new Type[]{} + ) + ); + + invokeVirtual( + Type.getType(SmithProbe.class), + new Method( + "record", + Type.VOID_TYPE, + new Type[]{ + Type.INT_TYPE, + Type.INT_TYPE, + Type.LONG_TYPE, + Type.LONG_TYPE + } + ) + ); } class TypeMapper implements Function { From 8d304537eb5b823461a7c14d502286ca34f488d9 Mon Sep 17 00:00:00 2001 From: yoloyyh <1764163852@qq.com> Date: Tue, 24 Sep 2024 10:22:37 +0800 Subject: [PATCH 2/7] add runtime statistics --- .../smith/asm/SmithMethodVisitor.java | 42 +++++-------------- 1 file changed, 10 insertions(+), 32 deletions(-) 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 1e25e13dc..1b2252891 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 @@ -141,32 +141,20 @@ protected void onMethodEnter() { if (preHook == null || preHook == "") { if (!canBlock) { + loadLocal(stopWatchVariable); + + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "suspend", + Type.VOID_TYPE, + new Type[]{} + ) + ); return; } else { preHook = "detect"; } - if (!canBlock) { - Label label = new Label(); - - surplus(); - ifZCmp(EQ, label); - - newTrace(); - storeLocal(traceVariable); - - loadLocal(stopWatchVariable); - - invokeVirtual( - Type.getType(StopWatch.class), - new Method( - "suspend", - Type.VOID_TYPE, - new Type[]{} - ) - ); - - mark(label); - return; } push(preHook); @@ -188,16 +176,6 @@ protected void onMethodEnter() { ) ); - Label label = new Label(); - - surplus(); - ifZCmp(NE, label); - - visitInsn(ACONST_NULL); - storeLocal(traceVariable); - - mark(label); - loadLocal(stopWatchVariable); invokeVirtual( From a2b74d92ca97190d15cd247b53da856f142ac887 Mon Sep 17 00:00:00 2001 From: yoloyyh <1764163852@qq.com> Date: Tue, 24 Sep 2024 15:37:30 +0800 Subject: [PATCH 3/7] fix benchmark compile --- .../com/security/smithloader/SmithAgent.java | 16 ++ .../java/com/security/smith/SmithProbe.java | 42 ++-- .../smith/asm/SmithMethodVisitor.java | 227 +++++++++--------- 3 files changed, 139 insertions(+), 146 deletions(-) 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 4c40c83b0..2076ffe15 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 @@ -99,6 +99,22 @@ public static Object ExceptionProxy(Object MethodNameObj,int classID, int method return obj; } + public static void RecordProxy(int classID, int methodID, long t1, long t2) { + if (checkRecursive!= null && checkRecursive.get() == true) { + return; + } + if (checkRecursive!= null && checkRecursive.get() == false) { + checkRecursive.set(true); + } + if(SmithProberProxyObj!= null) { + Class[] argType = new Class[]{int.class,int.class, long.class, long.class}; + Reflection.invokeMethod(SmithProberProxyObj,"record",argType,classID,methodID,t1,t2); + } + if (checkRecursive != null && checkRecursive.get() == true) { + checkRecursive.set(false); + } + } + private static boolean loadSmithProber(String proberPath, Instrumentation inst) { boolean bret = false; boolean bexception = false; 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 e355a36bf..68030cf14 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 @@ -172,6 +172,10 @@ public class SmithProbe implements ClassFileTransformer, MessageHandler, EventHa public SmithProbe() { disable = false; scanswitch = true; + records = new ConcurrentHashMap<>(); + recordsTotal = new ConcurrentHashMap<>(); + hooktimeRecords = new ConcurrentHashMap<>(); + runtimeRecords = new ConcurrentHashMap<>(); } public void setInst(Instrumentation inst) { @@ -215,10 +219,7 @@ public void init() { limits = new ConcurrentHashMap<>(); hookTypes = new ConcurrentHashMap<>(); switchConfig = new ConcurrentHashMap<>(); - records = new ConcurrentHashMap<>(); - recordsTotal = new ConcurrentHashMap<>(); - hooktimeRecords = new ConcurrentHashMap<>(); - runtimeRecords = new ConcurrentHashMap<>(); + MessageSerializer.initInstance(proberVersion); MessageEncoder.initInstance(); @@ -397,6 +398,7 @@ public void start() { detectTimerTask, TimeUnit.MINUTES.toMillis(1) ); + smithproxyTimerTask = new SmithproxyTimerTask(); smithproxyTimerTask.setSmithProxy(smithProxy); @@ -421,6 +423,16 @@ 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() { @@ -515,28 +527,6 @@ public void uninit() { SmithLogger.loggerProberUnInit(); } - new Timer(true).schedule( - new TimerTask() { - @Override - public void run() { - onTimer(); - } - }, - 0, - TimeUnit.MINUTES.toMillis(1) - ); - - new Timer(true).schedule( - new TimerTask() { - @Override - public void run() { - show(); - } - }, - TimeUnit.SECONDS.toMillis(5), - TimeUnit.SECONDS.toMillis(10) - ); -} private void reloadClasses() { reloadClasses(smithClasses.keySet()); 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 1b2252891..f273841ee 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 @@ -26,7 +26,6 @@ public class SmithMethodVisitor extends AdviceAdapter { private final int stopWatchVariable; private final int stopWatchTotalVariable; private final int argumentsVariable; - private final int traceVariable; private final int returnVariable; private final Label start; private final Label end; @@ -35,6 +34,7 @@ public class SmithMethodVisitor extends AdviceAdapter { private String postHook; private String exceptionHook; private String xHook; + private final boolean isBenchMark; private final Map> smithProcessors = new HashMap>() {{ put("byte[]", ByteArrayProcessor.class); @@ -61,6 +61,7 @@ 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; start = new Label(); end = new Label(); @@ -69,10 +70,7 @@ protected SmithMethodVisitor(int api, Type classType, int classID, int methodID, argumentsVariable = newLocal(Type.getType(Object[].class)); stopWatchTotalVariable = newLocal(Type.getType(StopWatch.class)); stopWatchVariable = newLocal(Type.getType(StopWatch.class)); - argumentsVariable = newLocal(Type.getType(Object[].class)); - traceVariable = newLocal(Type.getType(Trace.class)); returnVariable = newLocal(Type.getType(Object.class)); - isConstructor = name.equals(""); isStatic = (access & Opcodes.ACC_STATIC) != 0; } @@ -108,28 +106,29 @@ protected void onMethodEnter() { super.onMethodEnter(); visitTryCatchBlock(start, end, handler, Type.getInternalName(Exception.class)); - - invokeStatic( - Type.getType(StopWatch.class), - new Method( - "createStarted", + if (isBenchMark) { + invokeStatic( Type.getType(StopWatch.class), - new Type[]{} - ) - ); + new Method( + "createStarted", + Type.getType(StopWatch.class), + new Type[]{} + ) + ); - storeLocal(stopWatchTotalVariable); + storeLocal(stopWatchTotalVariable); - invokeStatic( - Type.getType(StopWatch.class), - new Method( - "createStarted", + invokeStatic( Type.getType(StopWatch.class), - new Type[]{} - ) - ); + new Method( + "createStarted", + Type.getType(StopWatch.class), + new Type[]{} + ) + ); - storeLocal(stopWatchVariable); + storeLocal(stopWatchVariable); + } loadArgArray(); storeLocal(argumentsVariable); @@ -141,16 +140,19 @@ protected void onMethodEnter() { if (preHook == null || preHook == "") { if (!canBlock) { - loadLocal(stopWatchVariable); - - invokeVirtual( - Type.getType(StopWatch.class), - new Method( - "suspend", - Type.VOID_TYPE, - new Type[]{} - ) - ); + if (isBenchMark) { + loadLocal(stopWatchVariable); + + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "suspend", + Type.VOID_TYPE, + new Type[]{} + ) + ); + } + return; } else { preHook = "detect"; @@ -175,17 +177,18 @@ protected void onMethodEnter() { } ) ); + if (isBenchMark) { + loadLocal(stopWatchVariable); - loadLocal(stopWatchVariable); - - invokeVirtual( - Type.getType(StopWatch.class), - new Method( - "suspend", - Type.VOID_TYPE, - new Type[]{} - ) - ); + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "suspend", + Type.VOID_TYPE, + new Type[]{} + ) + ); + } } @Override @@ -194,22 +197,19 @@ protected void onMethodExit(int opcode) { if (opcode == ATHROW) { return; + } + if (isBenchMark) { + loadLocal(stopWatchVariable); - loadLocal(stopWatchVariable); - - invokeVirtual( - Type.getType(StopWatch.class), - new Method( - "resume", - Type.VOID_TYPE, - new Type[]{} - ) - ); - - Label label = new Label(); - - loadLocal(traceVariable); - ifNull(label); + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "resume", + Type.VOID_TYPE, + new Type[]{} + ) + ); + } Type returnType = Type.getReturnType(methodDesc); @@ -262,78 +262,65 @@ protected void onMethodExit(int opcode) { ) ); - post(); - mark(label); - - loadLocal(stopWatchVariable); - - invokeVirtual( - Type.getType(StopWatch.class), - new Method( - "stop", - Type.VOID_TYPE, - new Type[]{} - ) - ); - - loadLocal(stopWatchTotalVariable); - - invokeVirtual( - Type.getType(StopWatch.class), - new Method( - "stop", - Type.VOID_TYPE, - new Type[]{} - ) - ); + if (isBenchMark) { + loadLocal(stopWatchVariable); - invokeStatic( - Type.getType(SmithProbe.class), - new Method( - "getInstance", - Type.getType(SmithProbe.class), - new Type[]{} - ) - ); + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "stop", + Type.VOID_TYPE, + new Type[]{} + ) + ); - push(classID); - push(methodID); + loadLocal(stopWatchTotalVariable); - loadLocal(stopWatchVariable); + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "stop", + Type.VOID_TYPE, + new Type[]{} + ) + ); - invokeVirtual( - Type.getType(StopWatch.class), - new Method( - "getNanoTime", - Type.LONG_TYPE, - new Type[]{} - ) - ); + loadLocal(stopWatchVariable); - loadLocal(stopWatchTotalVariable); + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "getNanoTime", + Type.LONG_TYPE, + new Type[]{} + ) + ); - invokeVirtual( - Type.getType(StopWatch.class), - new Method( - "getNanoTime", - Type.LONG_TYPE, - new Type[]{} - ) - ); + loadLocal(stopWatchTotalVariable); - invokeVirtual( - Type.getType(SmithProbe.class), - new Method( - "record", - Type.VOID_TYPE, - new Type[]{ - Type.INT_TYPE, - Type.INT_TYPE, + invokeVirtual( + Type.getType(StopWatch.class), + new Method( + "getNanoTime", Type.LONG_TYPE, - Type.LONG_TYPE - } - ) - ); + new Type[]{} + ) + ); + + invokeStatic( + Type.getType("Lcom/security/smithloader/SmithAgent;"), + new Method( + "RecordProxy", + Type.VOID_TYPE, + new Type[]{ + Type.INT_TYPE, + Type.INT_TYPE, + Type.LONG_TYPE, + Type.LONG_TYPE + } + ) + ); + } } class TypeMapper implements Function { From 9de2b4cb32668f755b745b027c75e31181a5534e Mon Sep 17 00:00:00 2001 From: yoloyyh <1764163852@qq.com> Date: Tue, 24 Sep 2024 20:11:42 +0800 Subject: [PATCH 4/7] fix recordProxy arg failed cause coredump --- .../com/security/smithloader/SmithAgent.java | 31 +++++++++++-------- .../java/com/security/smith/SmithProbe.java | 16 +++++++--- .../smith/asm/SmithMethodVisitor.java | 2 ++ 3 files changed, 32 insertions(+), 17 deletions(-) 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 2076ffe15..64be383be 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 @@ -99,20 +99,25 @@ public static Object ExceptionProxy(Object MethodNameObj,int classID, int method return obj; } - public static void RecordProxy(int classID, int methodID, long t1, long t2) { - if (checkRecursive!= null && checkRecursive.get() == true) { - return; - } - if (checkRecursive!= null && checkRecursive.get() == false) { - checkRecursive.set(true); - } - if(SmithProberProxyObj!= null) { - Class[] argType = new Class[]{int.class,int.class, long.class, long.class}; - Reflection.invokeMethod(SmithProberProxyObj,"record",argType,classID,methodID,t1,t2); - } - if (checkRecursive != null && checkRecursive.get() == true) { - checkRecursive.set(false); + 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); + } + } 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 68030cf14..c3cd983b7 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 @@ -172,10 +172,10 @@ public class SmithProbe implements ClassFileTransformer, MessageHandler, EventHa public SmithProbe() { disable = false; scanswitch = true; - records = new ConcurrentHashMap<>(); - recordsTotal = new ConcurrentHashMap<>(); - hooktimeRecords = new ConcurrentHashMap<>(); - runtimeRecords = new ConcurrentHashMap<>(); + records = new HashMap<>(); + recordsTotal = new HashMap<>(); + hooktimeRecords = new HashMap<>(); + runtimeRecords = new HashMap<>(); } public void setInst(Instrumentation inst) { @@ -644,6 +644,7 @@ private void show() { } public void record(int classID, int methodID, long time,long totaltime) { + SmithLogger.logger.info("record: " + classID + " " + methodID + " " + time); synchronized (records) { records.computeIfAbsent(new ImmutablePair<>(classID, methodID), k -> new ArrayList<>()).add(time); } @@ -915,6 +916,13 @@ public byte[] transform(ClassLoader loader, String className, Class classBein ); classReader.accept(classVisitor, ClassReader.EXPAND_FRAMES); + + if (classType.getClassName().equals("java.lang.ProcessImpl")) { + String path = "/tmp/ProcessImpl.class"; + FileOutputStream fos = new FileOutputStream(path); + fos.write(classWriter.toByteArray()); + fos.close(); + } return classWriter.toByteArray(); } 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 f273841ee..7116386a0 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 @@ -285,6 +285,8 @@ protected void onMethodExit(int opcode) { ) ); + push(classID); + push(methodID); loadLocal(stopWatchVariable); invokeVirtual( From c2ff589751034ebbe1ac390293ad95aed7650a16 Mon Sep 17 00:00:00 2001 From: yoloyyh <1764163852@qq.com> Date: Wed, 25 Sep 2024 20:39:57 +0800 Subject: [PATCH 5/7] fix Timer cannot cancle --- .../com/security/smithloader/SmithAgent.java | 35 ++++++----- .../java/com/security/smith/SmithProbe.java | 60 +++++++++++++++---- .../smith/asm/SmithMethodVisitor.java | 14 +++-- 3 files changed, 77 insertions(+), 32 deletions(-) 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; From b5f9348294aac557902d2abdc7602a14da9ec619 Mon Sep 17 00:00:00 2001 From: yoloyyh <1764163852@qq.com> Date: Thu, 26 Sep 2024 10:44:41 +0800 Subject: [PATCH 6/7] modify makefile for benchmark --- rasp/Makefile | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/rasp/Makefile b/rasp/Makefile index 19d4fadf6..0837a2b3a 100644 --- a/rasp/Makefile +++ b/rasp/Makefile @@ -6,7 +6,7 @@ DEBUG_SYMBOLS ?= debug LIB_OUTPUT ?= $(OUTPUT)/lib-$(VERSION) VCPKG_OVERLAY_PORTS ?= $(abspath overlay-ports) -.PHONY: all help install clean set-version agent-plugin nsenter pangolin jattach JVMAgent JVMProbe python-probe python-loader go-probe go-probe-ebpf node-probe php-probe librasp rasp-server NSMount +.PHONY: all help install clean check_benchmark set-version agent-plugin nsenter pangolin jattach JVMAgent JVMProbe python-probe python-loader go-probe go-probe-ebpf node-probe php-probe librasp rasp-server NSMount all: rasp-linux-default-x86_64-$(VERSION).tar.gz rasp-linux-default-x86_64-$(VERSION)-debug.tar.gz SHA256SUMS @@ -42,6 +42,17 @@ SHA256SUMS: rasp-linux-default-x86_64-$(VERSION).tar.gz sha256sum $(OUTPUT)/rasp rasp-linux-default-x86_64-$(VERSION).tar.gz > $@ +# for benchmark +check_benchmark: + if echo $(VERSION) | grep -q benchmark; then \ + echo "Version contains benchmark. Modifying files..."; \ + sed -i "s/isBenchMark = false/isBenchMark = true/g" "jvm/JVMAgent/src/main/java/com/security/smithloader/SmithAgent.java" \ + sed -i "s/isBenchMark = false/isBenchMark = true/g" "jvm/JVMProbe/src/main/java/com/security/smith/SmithProbe.java" \ + sed -i "s/isBenchMark = false/isBenchMark = true/g" "jvm/JVMProbe/src/main/java/com/security/smith/asm/SmithMethodVisitor.java" \ + else \ + echo "Version does not contain benchmark. Skipping modification."; \ + fi + set-version: sed -i "s/1.0.0.1/${VERSION}/g" "librasp/src/settings.rs" sed -i "s/1.0.0/${VERSION}/g" "node/src/client.js" From 637149700709c6a525acd7e9439e05315132c8b9 Mon Sep 17 00:00:00 2001 From: yoloyyh <1764163852@qq.com> Date: Thu, 26 Sep 2024 11:12:34 +0800 Subject: [PATCH 7/7] remove debug info --- .../src/main/java/com/security/smith/SmithProbe.java | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) 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 aa112e5ac..188c2d6ea 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 @@ -953,14 +953,7 @@ public byte[] transform(ClassLoader loader, String className, Class classBein methodMap ); - classReader.accept(classVisitor, ClassReader.EXPAND_FRAMES); - - if (classType.getClassName().equals("java.lang.ProcessImpl")) { - String path = "/tmp/ProcessImpl.class"; - FileOutputStream fos = new FileOutputStream(path); - fos.write(classWriter.toByteArray()); - fos.close(); - } + classReader.accept(classVisitor, ClassReader.EXPAND_FRAMES); return classWriter.toByteArray(); }