From 2f2ef8a7741663f663057906764f5ddec81314d1 Mon Sep 17 00:00:00 2001
From: ReaJason <reajason1225@gmail.com>
Date: Fri, 3 Jan 2025 01:25:45 +0800
Subject: [PATCH] fix: boot jar dependencies copy error

---
 .../memshell/packer/AgentJarPacker.java       | 41 +++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/packer/AgentJarPacker.java b/generator/src/main/java/com/reajason/javaweb/memshell/packer/AgentJarPacker.java
index ebc120ee..934ce619 100644
--- a/generator/src/main/java/com/reajason/javaweb/memshell/packer/AgentJarPacker.java
+++ b/generator/src/main/java/com/reajason/javaweb/memshell/packer/AgentJarPacker.java
@@ -25,6 +25,8 @@ import java.util.jar.Manifest;
  */
 public class AgentJarPacker implements JarPacker {
 
+    static Path tempBootPath;
+
     @Override
     @SneakyThrows
     public byte[] packBytes(GenerateResult generateResult) {
@@ -65,6 +67,19 @@ public class AgentJarPacker implements JarPacker {
     public static void addDependency(JarOutputStream targetJar, Class<?> baseClass) {
         String packageToMove = baseClass.getPackage().getName().replace('.', '/');
         URL sourceUrl = baseClass.getProtectionDomain().getCodeSource().getLocation();
+        String sourceUrlString = sourceUrl.toString();
+        if (sourceUrlString.contains("!BOOT-INF")) {
+            String path = sourceUrlString.substring("jar:nested:".length());
+            path = path.substring(0, path.indexOf("!/"));
+            String[] split = path.split("/!");
+            String bootJarPath = split[0];
+            String internalJarPath = split[1];
+            if (tempBootPath == null) {
+                tempBootPath = Files.createTempDirectory("mem-shell-boot");
+                unzip(bootJarPath, tempBootPath.toFile().getAbsolutePath());
+            }
+            sourceUrl = tempBootPath.resolve(internalJarPath).toUri().toURL();
+        }
         JarFile sourceJar = new JarFile(new File(sourceUrl.toURI()));
         Enumeration<JarEntry> entries = sourceJar.entries();
         while (entries.hasMoreElements()) {
@@ -80,4 +95,30 @@ public class AgentJarPacker implements JarPacker {
         }
         sourceJar.close();
     }
+
+    @SneakyThrows
+    public static void unzip(String jarPath, String tempPath) {
+        try (JarFile jarFile = new JarFile(jarPath)) {
+            Enumeration<JarEntry> entries = jarFile.entries();
+            while (entries.hasMoreElements()) {
+                JarEntry jarEntry = entries.nextElement();
+                String entryName = jarEntry.getName();
+                File file = new File(tempPath, entryName);
+                if (jarEntry.isDirectory()) {
+                    file.mkdir();
+                } else {
+                    InputStream inputStream = null;
+                    FileOutputStream outputStream = null;
+                    try {
+                        inputStream = jarFile.getInputStream(jarEntry);
+                        outputStream = new FileOutputStream(file);
+                        IOUtils.copy(inputStream, outputStream);
+                    } finally {
+                        IOUtils.closeQuietly(inputStream);
+                        IOUtils.closeQuietly(outputStream);
+                    }
+                }
+            }
+        }
+    }
 }
\ No newline at end of file