fix benchmark compile

bytedance · Sep 24, 2024 · a2b74d9 · a2b74d9
1 parent 8d30453
commit a2b74d9
Show file tree

Hide file tree

Showing 3 changed files with 139 additions and 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
@@ -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
@@ -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());