fix: custom shell classname use itself

Author: ReaJason

generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java

@@ -24,15 +24,6 @@ public static MemShellResult generate(ShellConfig shellConfig, InjectorConfig in
         if (server == null) {
             throw new GenerationException("Unsupported server: " + serverName);
         }
-
-        if (StringUtils.isBlank(shellToolConfig.getShellClassName())) {
-            shellToolConfig.setShellClassName(CommonUtil.generateShellClassName(serverName, shellConfig.getShellType()));
-        }
-
-        if (StringUtils.isBlank(injectorConfig.getInjectorClassName())) {
-            injectorConfig.setInjectorClassName(CommonUtil.generateInjectorClassName());
-        }
-
         Class<?> injectorClass = null;
 
         if (ShellTool.Custom.equals(shellConfig.getShellTool())) {
@@ -45,6 +36,13 @@ public static MemShellResult generate(ShellConfig shellConfig, InjectorConfig in
             Class<?> shellClass = shellInjectorPair.getLeft();
             injectorClass = shellInjectorPair.getRight();
             shellToolConfig.setShellClass(shellClass);
+            if (StringUtils.isBlank(shellToolConfig.getShellClassName())) {
+                shellToolConfig.setShellClassName(CommonUtil.generateShellClassName(serverName, shellConfig.getShellType()));
+            }
+        }
+
+        if (StringUtils.isBlank(injectorConfig.getInjectorClassName())) {
+            injectorConfig.setInjectorClassName(CommonUtil.generateInjectorClassName());
         }
 
         byte[] shellBytes = ShellToolFactory.generateBytes(shellConfig, shellToolConfig);

generator/src/main/java/com/reajason/javaweb/memshell/generator/ByteBuddyShellGenerator.java

@@ -31,8 +31,8 @@ protected ByteBuddyShellGenerator(ShellConfig shellConfig, T shellToolConfig) {
 
     @Override
     public byte[] getBytes() {
-        String shellClassName = shellToolConfig.getShellClassName();
         DynamicType.Builder<?> builder = getBuilder();
+        String shellClassName = shellToolConfig.getShellClassName();
         Class<?> shellClass = shellToolConfig.getShellClass();
         if (shellClass != null) {
             shellToolConfig.setShellTypeDescription(TypeDescription.ForLoadedType.of(shellClass));

generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java

@@ -8,6 +8,7 @@
 import net.bytebuddy.dynamic.DynamicType;
 import net.bytebuddy.jar.asm.ClassReader;
 import net.bytebuddy.pool.TypePool;
+import org.apache.commons.lang3.StringUtils;
 
 import java.util.Base64;
 
@@ -27,6 +28,9 @@ protected DynamicType.Builder<?> getBuilder() {
         byte[] classBytes = Base64.getDecoder().decode(shellClassBase64);
         ClassReader classReader = new ClassReader(classBytes);
         String className = classReader.getClassName().replace('/', '.');
+        if (StringUtils.isBlank(shellToolConfig.getShellClassName())) {
+            shellToolConfig.setShellClassName(className);
+        }
         ClassFileLocator classFileLocator = ClassFileLocator.Simple.of(className, classBytes);
         TypeDescription typeDescription = new TypePool.Default(
                 new TypePool.CacheProvider.Simple(), classFileLocator,