fix: add generate exception global handler

Author: ReaJason

boot/src/main/java/com/reajason/javaweb/boot/api/GlobalExceptionHandler.java

@@ -1,5 +1,6 @@
 package com.reajason.javaweb.boot.api;
 
+import com.reajason.javaweb.GenerationException;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseEntity;
@@ -22,10 +23,16 @@ public ResponseEntity<?> handleNoResourceException() {
         return ResponseEntity.status(HttpStatus.NOT_FOUND).build();
     }
 
+    @ResponseStatus(HttpStatus.INTERNAL_SERVER_ERROR)
+    @ExceptionHandler(GenerationException.class)
+    public ErrorResponse handleGenerationException(GenerationException exception) {
+        return new ErrorResponse(exception.getMessage());
+    }
+
     @ResponseStatus(HttpStatus.INTERNAL_SERVER_ERROR)
     @ExceptionHandler(Throwable.class)
     public ErrorResponse handleThrowable(Throwable throwable) {
-        log.error("请求出错", throwable);
+        log.error("Internal Exception", throwable);
         return new ErrorResponse(throwable.getMessage());
     }
 }

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

@@ -1,5 +1,6 @@
 package com.reajason.javaweb.memshell;
 
+import com.reajason.javaweb.GenerationException;
 import com.reajason.javaweb.memshell.config.InjectorConfig;
 import com.reajason.javaweb.memshell.config.ShellConfig;
 import com.reajason.javaweb.memshell.config.ShellToolConfig;
@@ -21,7 +22,7 @@ public static MemShellResult generate(ShellConfig shellConfig, InjectorConfig in
         String serverName = shellConfig.getServer();
         AbstractServer server = ServerFactory.getServer(serverName);
         if (server == null) {
-            throw new IllegalArgumentException("Unsupported server: " + serverName);
+            throw new GenerationException("Unsupported server: " + serverName);
         }
 
         if (StringUtils.isBlank(shellToolConfig.getShellClassName())) {
@@ -39,7 +40,7 @@ public static MemShellResult generate(ShellConfig shellConfig, InjectorConfig in
         } else {
             Pair<Class<?>, Class<?>> shellInjectorPair = server.getShellInjectorPair(shellConfig.getShellTool(), shellConfig.getShellType());
             if (shellInjectorPair == null) {
-                throw new UnsupportedOperationException(serverName + " unsupported shell type: " + shellConfig.getShellType() + " for tool: " + shellConfig.getShellTool());
+                throw new GenerationException(serverName + " unsupported shell type: " + shellConfig.getShellType() + " for tool: " + shellConfig.getShellTool());
             }
             Class<?> shellClass = shellInjectorPair.getLeft();
             injectorClass = shellInjectorPair.getRight();

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

@@ -1,6 +1,7 @@
 package com.reajason.javaweb.memshell.generator;
 
 import com.reajason.javaweb.ClassBytesShrink;
+import com.reajason.javaweb.GenerationException;
 import com.reajason.javaweb.asm.ClassRenameUtils;
 import com.reajason.javaweb.memshell.config.CustomConfig;
 import com.reajason.javaweb.memshell.config.ShellConfig;
@@ -23,7 +24,7 @@ public byte[] getBytes() {
         String shellClassBase64 = shellToolConfig.getShellClassBase64();
 
         if (StringUtils.isBlank(shellClassBase64)) {
-            throw new IllegalArgumentException("Custom shell class is empty");
+            throw new GenerationException("Custom shell class is empty");
         }
         byte[] classBytes = Base64.getDecoder().decode(shellClassBase64);
         byte[] bytes = ClassRenameUtils.renameClass(classBytes, shellToolConfig.getShellClassName());

generator/src/main/java/com/reajason/javaweb/probe/ProbeMethod.java

@@ -5,6 +5,7 @@
 import com.reajason.javaweb.probe.generator.DnsLogGenerator;
 import com.reajason.javaweb.probe.generator.SleepGenerator;
 import com.reajason.javaweb.probe.generator.response.ResponseBodyGenerator;
+import lombok.SneakyThrows;
 
 import java.lang.reflect.Constructor;
 
@@ -25,14 +26,11 @@ public enum ProbeMethod {
         this.configClass = configClass;
     }
 
+    @SneakyThrows
     public byte[] generateBytes(ProbeConfig probeConfig, ProbeContentConfig probeContentConfig) {
-        try {
-            Constructor<? extends ShellGenerator> constructor =
-                    generatorClass.getConstructor(ProbeConfig.class, configClass);
-            ShellGenerator generator = constructor.newInstance(probeConfig, configClass.cast(probeContentConfig));
-            return generator.getBytes();
-        } catch (Exception e) {
-            throw new RuntimeException("shell generate failed: " + e.getMessage(), e);
-        }
+        Constructor<? extends ShellGenerator> constructor =
+                generatorClass.getConstructor(ProbeConfig.class, configClass);
+        ShellGenerator generator = constructor.newInstance(probeConfig, configClass.cast(probeContentConfig));
+        return generator.getBytes();
     }
 }

generator/src/main/java/com/reajason/javaweb/probe/generator/DnsLogGenerator.java

@@ -1,5 +1,6 @@
 package com.reajason.javaweb.probe.generator;
 
+import com.reajason.javaweb.GenerationException;
 import com.reajason.javaweb.buddy.TargetJreVersionVisitorWrapper;
 import com.reajason.javaweb.probe.ProbeContent;
 import com.reajason.javaweb.probe.config.DnsLogConfig;
@@ -12,6 +13,7 @@
 import net.bytebuddy.ByteBuddy;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.dynamic.DynamicType;
+import org.apache.commons.lang3.StringUtils;
 
 import static net.bytebuddy.matcher.ElementMatchers.named;
 
@@ -27,6 +29,9 @@ public DnsLogGenerator(ProbeConfig probeConfig, DnsLogConfig probeContentConfig)
 
     @Override
     protected DynamicType.Builder<?> build(ByteBuddy buddy) {
+        if (StringUtils.isBlank(probeContentConfig.getHost())) {
+            throw new GenerationException("DNSLog probeShell host must be specified");
+        }
         ProbeContent detectContent = probeConfig.getProbeContent();
         switch (detectContent) {
             case Server:
@@ -42,7 +47,7 @@ protected DynamicType.Builder<?> build(ByteBuddy buddy) {
                         .field(named("host")).value(probeContentConfig.getHost())
                         .visit(Advice.to(JdkProbe.class).on(named("getJdk")));
             default:
-                throw new UnsupportedOperationException(detectContent + " not supported");
+                throw new GenerationException(detectContent + " not supported");
         }
     }
 }

generator/src/main/java/com/reajason/javaweb/probe/generator/SleepGenerator.java

@@ -1,5 +1,6 @@
 package com.reajason.javaweb.probe.generator;
 
+import com.reajason.javaweb.GenerationException;
 import com.reajason.javaweb.buddy.TargetJreVersionVisitorWrapper;
 import com.reajason.javaweb.probe.ProbeContent;
 import com.reajason.javaweb.probe.config.ProbeConfig;
@@ -10,6 +11,7 @@
 import net.bytebuddy.ByteBuddy;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.dynamic.DynamicType;
+import org.apache.commons.lang3.StringUtils;
 
 import static net.bytebuddy.matcher.ElementMatchers.named;
 
@@ -28,13 +30,19 @@ public SleepGenerator(ProbeConfig probeConfig, SleepConfig probeContentConfig) {
     protected DynamicType.Builder<?> build(ByteBuddy buddy) {
         ProbeContent detectContent = probeConfig.getProbeContent();
         if (ProbeContent.Server.equals(detectContent)) {
+            if (probeContentConfig.getSeconds() <= 0) {
+                throw new GenerationException("sleepProbe seconds must be greater than 0");
+            }
+            if (StringUtils.isEmpty(probeContentConfig.getServer())) {
+                throw new GenerationException("sleepProbe server must be specified");
+            }
             return buddy.redefine(SleepServer.class)
                     .name(CommonUtil.generateShellClassName())
                     .visit(TargetJreVersionVisitorWrapper.DEFAULT)
                     .field(named("server")).value(probeContentConfig.getServer())
                     .field(named("seconds")).value(probeContentConfig.getSeconds())
                     .visit(Advice.to(ServerProbe.class).on(named("getServer")));
         }
-        throw new UnsupportedOperationException("Sleep Probe not supported for " + detectContent);
+        throw new GenerationException("sleepProbe not supported for " + detectContent);
     }
 }

generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java

@@ -1,5 +1,6 @@
 package com.reajason.javaweb.probe.generator.response;
 
+import com.reajason.javaweb.GenerationException;
 import com.reajason.javaweb.Server;
 import com.reajason.javaweb.buddy.MethodCallReplaceVisitorWrapper;
 import com.reajason.javaweb.buddy.TargetJreVersionVisitorWrapper;
@@ -35,9 +36,11 @@ protected DynamicType.Builder<?> build(ByteBuddy buddy) {
         if (probeContentConfig.getReqParamName() != null) {
             name = probeContentConfig.getReqParamName();
             getDataFromReqInterceptor = getDataFromReqParamInterceptor.class;
-        } else {
+        } else if(probeContentConfig.getReqHeaderName() != null) {
             name = probeContentConfig.getReqHeaderName();
             getDataFromReqInterceptor = getDataFromReqHeaderInterceptor.class;
+        }else{
+            throw new GenerationException("responseBody probeShell must set headerName or paramName");
         }
         Class<?> writerClass = getWriterClass();
         Class<?> runnerClass = getRunnerClass();
@@ -58,7 +61,7 @@ private Class<?> getRunnerClass() {
             case Bytecode:
                 return ByteCodeProbe.class;
             default:
-                throw new IllegalArgumentException("responseBody not supported for probe content: " + probeConfig.getProbeContent());
+                throw new GenerationException("responseBody not supported for probe content: " + probeConfig.getProbeContent());
         }
     }
 
@@ -86,7 +89,7 @@ private Class<?> getWriterClass() {
             case Server.Apusic:
                 return ApusicWriter.class;
             default:
-                throw new IllegalArgumentException("responseBody now supported for server: " + probeContentConfig.getServer());
+                throw new GenerationException("responseBody not supported for server: " + probeContentConfig.getServer());
         }
     }