fix recordProxy arg failed cause coredump

bytedance · Sep 24, 2024 · 9de2b4c · 9de2b4c
1 parent a2b74d9
commit 9de2b4c
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 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
@@ -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
@@ -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
@@ -285,6 +285,8 @@ protected void onMethodExit(int opcode) {
                         )
                 );
 
+                push(classID);
+                push(methodID);
                 loadLocal(stopWatchVariable);
 
                 invokeVirtual(