From 798e447ab62c1c2472c1becdee46540798118ef7 Mon Sep 17 00:00:00 2001
From: ReaJason <reajason1225@gmail.com>
Date: Fri, 19 Sep 2025 21:44:52 +0800
Subject: [PATCH] fix: custom shell classname use itself

---
 .../javaweb/memshell/MemShellGenerator.java      | 16 +++++++---------
 .../generator/ByteBuddyShellGenerator.java       |  2 +-
 .../memshell/generator/CustomShellGenerator.java |  4 ++++
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java
index 102e80e9..9fdea05f 100644
--- a/generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java
+++ b/generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java
@@ -24,15 +24,6 @@ public class MemShellGenerator {
         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 class MemShellGenerator {
             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);
diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/generator/ByteBuddyShellGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/generator/ByteBuddyShellGenerator.java
index f1d8eddb..296e2568 100644
--- a/generator/src/main/java/com/reajason/javaweb/memshell/generator/ByteBuddyShellGenerator.java
+++ b/generator/src/main/java/com/reajason/javaweb/memshell/generator/ByteBuddyShellGenerator.java
@@ -31,8 +31,8 @@ public abstract class ByteBuddyShellGenerator<T extends ShellToolConfig> impleme
 
     @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));
diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java
index d0f57a93..8c0bc76e 100644
--- a/generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java
+++ b/generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java
@@ -8,6 +8,7 @@ import net.bytebuddy.dynamic.ClassFileLocator;
 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 @@ public class CustomShellGenerator extends ByteBuddyShellGenerator<CustomConfig>
         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,