From 63310043e6319995abc6b89c279db1219bf6ee48 Mon Sep 17 00:00:00 2001 From: ReaJason Date: Wed, 13 Aug 2025 22:19:10 +0800 Subject: [PATCH] fix: add generate exception global handler --- .../javaweb/boot/api/GlobalExceptionHandler.java | 9 ++++++++- .../javaweb/memshell/MemShellGenerator.java | 5 +++-- .../memshell/generator/CustomShellGenerator.java | 3 ++- .../com/reajason/javaweb/probe/ProbeMethod.java | 14 ++++++-------- .../javaweb/probe/generator/DnsLogGenerator.java | 7 ++++++- .../javaweb/probe/generator/SleepGenerator.java | 10 +++++++++- .../generator/response/ResponseBodyGenerator.java | 9 ++++++--- 7 files changed, 40 insertions(+), 17 deletions(-) diff --git a/boot/src/main/java/com/reajason/javaweb/boot/api/GlobalExceptionHandler.java b/boot/src/main/java/com/reajason/javaweb/boot/api/GlobalExceptionHandler.java index 8433174a..f7a772e4 100644 --- a/boot/src/main/java/com/reajason/javaweb/boot/api/GlobalExceptionHandler.java +++ b/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 class GlobalExceptionHandler { 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()); } } 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 ceb2903a..6f0da681 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java +++ b/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 class MemShellGenerator { 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 class MemShellGenerator { } else { Pair, 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(); 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 1037d4f5..8266fa53 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 @@ -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 class CustomShellGenerator extends ASMShellGenerator { 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()); diff --git a/generator/src/main/java/com/reajason/javaweb/probe/ProbeMethod.java b/generator/src/main/java/com/reajason/javaweb/probe/ProbeMethod.java index f4c102fa..ae37b0bb 100644 --- a/generator/src/main/java/com/reajason/javaweb/probe/ProbeMethod.java +++ b/generator/src/main/java/com/reajason/javaweb/probe/ProbeMethod.java @@ -5,6 +5,7 @@ import com.reajason.javaweb.probe.config.*; 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 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 constructor = + generatorClass.getConstructor(ProbeConfig.class, configClass); + ShellGenerator generator = constructor.newInstance(probeConfig, configClass.cast(probeContentConfig)); + return generator.getBytes(); } } diff --git a/generator/src/main/java/com/reajason/javaweb/probe/generator/DnsLogGenerator.java b/generator/src/main/java/com/reajason/javaweb/probe/generator/DnsLogGenerator.java index e59b3c74..74fd73b6 100644 --- a/generator/src/main/java/com/reajason/javaweb/probe/generator/DnsLogGenerator.java +++ b/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 com.reajason.javaweb.utils.CommonUtil; 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 class DnsLogGenerator extends ByteBuddyShellGenerator { @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 @@ public class DnsLogGenerator extends ByteBuddyShellGenerator { .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"); } } } diff --git a/generator/src/main/java/com/reajason/javaweb/probe/generator/SleepGenerator.java b/generator/src/main/java/com/reajason/javaweb/probe/generator/SleepGenerator.java index 859c37a6..e38d85dc 100644 --- a/generator/src/main/java/com/reajason/javaweb/probe/generator/SleepGenerator.java +++ b/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 com.reajason.javaweb.utils.CommonUtil; 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,6 +30,12 @@ public class SleepGenerator extends ByteBuddyShellGenerator { 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) @@ -35,6 +43,6 @@ public class SleepGenerator extends ByteBuddyShellGenerator { .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); } } \ No newline at end of file diff --git a/generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java b/generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java index 88922670..03b480d0 100644 --- a/generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java +++ b/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 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator writerClass = getWriterClass(); Class runnerClass = getRunnerClass(); @@ -58,7 +61,7 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator