Merge pull request #655 from bytedance/feat-memshell-upload

add memshell upload
bytedance · Jul 15, 2024 · cc4149b · cc4149b
2 parents eea3e5e + c2a12f4
commit cc4149b
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 9 deletions.
diff --git a/rasp/jvm/JVMProbe/src/main/java/com/security/smith/client/ClassUploadTransformer.java b/rasp/jvm/JVMProbe/src/main/java/com/security/smith/client/ClassUploadTransformer.java
@@ -116,14 +116,13 @@ public void run() {
                     this.clazzToUpload = info.clazz;
                     this.transId = info.transId;
 
-                    // not upload. just send
-                    // if (inst.isModifiableClass(info.clazz) && !info.clazz.getName().startsWith("java.lang.invoke.LambdaForm")) {
-                    //     try {
-                    //         inst.retransformClasses(info.clazz);
-                    //     } catch (Exception e) {
-                    //         SmithLogger.exception(e);
-                    //     }
-                    // }
+                    if (inst.isModifiableClass(info.clazz) && !info.clazz.getName().startsWith("java.lang.invoke.LambdaForm")) {
+                        try {
+                            inst.retransformClasses(info.clazz);
+                        } catch (Exception e) {
+                            SmithLogger.exception(e);
+                        }
+                    }
                 }
                 finally {
                     this.clazzToUpload = null;
@@ -369,8 +368,8 @@ private void sendClass(Class<?> clazz, byte[] data) {
                 classUpload.setClassData(data);
 
                 if (client != null) {
-                    client.write(Operate.CLASSUPLOAD, classUpload);
                     SmithLogger.logger.info("send classdata: " + classUpload.toString());
+                    client.write(Operate.CLASSUPLOAD, classUpload);
                 }
             }
 

diff --git a/rasp/jvm/JVMProbe/src/main/java/com/security/smith/client/Client.java b/rasp/jvm/JVMProbe/src/main/java/com/security/smith/client/Client.java
@@ -81,8 +81,17 @@ public void write(int operate, Object object) {
         if (channel == null || !channel.isActive() || !channel.isWritable())
             return;
 
+        try {
+            //TODO must send metadata first
+            if (operate == Operate.CLASSUPLOAD) {
+                Thread.sleep(1000);
+            }
+        } catch (Exception e) {
+
+        }
         ObjectMapper objectMapper = new ObjectMapper()
                 .setPropertyNamingStrategy(PropertyNamingStrategies.SNAKE_CASE);
+
 
         Message message = new Message();