introducing captureAndRestoreAsmTypeComparator

pfichtner · Oct 7, 2023 · 7a0f412 · 7a0f412
1 parent 645fcfa
commit 7a0f412
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 14 deletions.
diff --git a/src/test/java/com/github/pfichtner/log4shell/scanner/Log4ShellHunterTest.java b/src/test/java/com/github/pfichtner/log4shell/scanner/Log4ShellHunterTest.java
@@ -2,6 +2,7 @@
 
 import static com.github.pfichtner.log4shell.scanner.Detectors.allDetectors;
 import static com.github.pfichtner.log4shell.scanner.Scrubbers.basedirScrubber;
+import static com.github.pfichtner.log4shell.scanner.util.Util.captureAndRestoreAsmTypeComparator;
 import static com.github.stefanbirkner.systemlambda.SystemLambda.catchSystemExit;
 import static com.github.stefanbirkner.systemlambda.SystemLambda.tapSystemErr;
 import static com.github.stefanbirkner.systemlambda.SystemLambda.tapSystemOut;
@@ -141,13 +142,8 @@ private void verifyMain(String... args) throws Exception {
 
 	private String execMain(String... args) throws Exception {
 		Map<String, String> values = new HashMap<>();
-		AsmTypeComparator old = AsmTypeComparator.typeComparator();
-		try {
-			values.put(STDERR, tapSystemErr(() -> values.put(STDOUT, tapSystemOut(
-					() -> values.put(RC, String.valueOf(catchSystemExit(() -> Log4ShellHunter.main(args))))))));
-		} finally {
-			AsmTypeComparator.useTypeComparator(old);
-		}
+		captureAndRestoreAsmTypeComparator(() -> values.put(STDERR, tapSystemErr(() -> values.put(STDOUT, tapSystemOut(
+				() -> values.put(RC, String.valueOf(catchSystemExit(() -> Log4ShellHunter.main(args)))))))));
 		List<String> elements = asList( //
 				"stdout", values.getOrDefault(STDOUT, ""), //
 				"stderr", values.getOrDefault(STDERR, ""), //

diff --git a/src/test/java/com/github/pfichtner/log4shell/scanner/Log4jSamplesIT.java b/src/test/java/com/github/pfichtner/log4shell/scanner/Log4jSamplesIT.java
@@ -2,6 +2,7 @@
 
 import static com.github.pfichtner.log4shell.scanner.Scrubbers.basedirScrubber;
 import static com.github.pfichtner.log4shell.scanner.io.Files.isArchive;
+import static com.github.pfichtner.log4shell.scanner.util.Util.captureAndRestoreAsmTypeComparator;
 import static com.github.stefanbirkner.systemlambda.SystemLambda.tapSystemOut;
 import static java.nio.file.Files.walk;
 import static java.util.Arrays.asList;
@@ -65,14 +66,9 @@ private static Stream<AsmTypeComparator> allModes() {
 		return EnumSet.allOf(AsmTypeComparator.class).stream();
 	}
 
-	private static void doCheck(List<String> filenames, AsmTypeComparator typeComparator) throws IOException {
+	private static void doCheck(List<String> filenames, AsmTypeComparator typeComparator) throws Exception {
 		System.out.println("*** using " + typeComparator);
-		AsmTypeComparator old = AsmTypeComparator.typeComparator();
-		try {
-			Log4ShellHunter.main(args(filenames, typeComparator));
-		} finally {
-			AsmTypeComparator.useTypeComparator(old);
-		}
+		captureAndRestoreAsmTypeComparator(() -> Log4ShellHunter.main(args(filenames, typeComparator)));
 	}
 
 	private static String[] args(List<String> filenames, AsmTypeComparator typeComparator) {

diff --git a/src/test/java/com/github/pfichtner/log4shell/scanner/util/Util.java b/src/test/java/com/github/pfichtner/log4shell/scanner/util/Util.java
@@ -16,6 +16,7 @@
 import com.github.pfichtner.log4shell.scanner.DetectionCollector;
 import com.github.pfichtner.log4shell.scanner.DetectionCollector.Detection;
 import com.github.pfichtner.log4shell.scanner.detectors.AbstractDetector;
+import com.github.stefanbirkner.systemlambda.Statement;
 
 public final class Util {
 
@@ -53,4 +54,13 @@ public static Map<File, List<Detection>> analyse(Log4jJars log4jJars, AbstractDe
 		return results;
 	}
 
+	public static void captureAndRestoreAsmTypeComparator(Statement statement) throws Exception {
+		AsmTypeComparator old = AsmTypeComparator.typeComparator();
+		try {
+			statement.execute();
+		} finally {
+			AsmTypeComparator.useTypeComparator(old);
+		}
+	}
+
 }