From 070ff2f9d9bec4394dbe459c1336b491669cc431 Mon Sep 17 00:00:00 2001 From: Matyrobbrt Date: Sun, 23 Jun 2024 22:34:54 +0300 Subject: [PATCH] Add a tool for zip injection --- settings.gradle | 4 +- zipinject/build.gradle | 25 +++ .../net/neoforged/zipinject/ConsoleTool.java | 148 ++++++++++++++++++ .../neoforged/zipinject/ZipInjectTest.java | 128 +++++++++++++++ 4 files changed, 304 insertions(+), 1 deletion(-) create mode 100644 zipinject/build.gradle create mode 100644 zipinject/src/main/java/net/neoforged/zipinject/ConsoleTool.java create mode 100644 zipinject/src/test/java/net/neoforged/zipinject/ZipInjectTest.java diff --git a/settings.gradle b/settings.gradle index 65e4abf..c4aa9a2 100644 --- a/settings.gradle +++ b/settings.gradle @@ -5,7 +5,8 @@ dependencyResolutionManagement { library('junit-api', 'org.junit.jupiter', 'junit-jupiter-api').versionRef('junit') library('junit-engine', 'org.junit.jupiter', 'junit-jupiter-engine').versionRef('junit') library('junit-platform-launcher', 'org.junit.platform:junit-platform-launcher:1.10.0') - bundle('junit', ['junit-engine', 'junit-platform-launcher', 'junit-api']) + library('assert4j', 'org.assertj:assertj-core:3.25.1') + bundle('junit', ['junit-engine', 'junit-platform-launcher', 'junit-api', 'assert4j']) library('srgutils', 'net.neoforged:srgutils:1.0.0') library('jopt', 'net.sf.jopt-simple:jopt-simple:5.0.4') @@ -18,3 +19,4 @@ dependencyResolutionManagement { include(':cli-utils') include(':jarsplitter') include(':binarypatcher') +include(':zipinject') diff --git a/zipinject/build.gradle b/zipinject/build.gradle new file mode 100644 index 0000000..4b9ad64 --- /dev/null +++ b/zipinject/build.gradle @@ -0,0 +1,25 @@ +dependencies { + implementation(libs.jopt) + testImplementation(libs.bundles.junit) +} + +application { + mainClass = 'net.neoforged.zipinject.ConsoleTool' +} + +test { + useJUnitPlatform() +} + +publishing { + publications.register('mavenJava', MavenPublication) { + from components.java + + artifactId = 'zipinject' + + pom { + name = 'Zip Inject' + description = 'Injects a folder or another zip into a zip' + } + } +} diff --git a/zipinject/src/main/java/net/neoforged/zipinject/ConsoleTool.java b/zipinject/src/main/java/net/neoforged/zipinject/ConsoleTool.java new file mode 100644 index 0000000..a8c10d7 --- /dev/null +++ b/zipinject/src/main/java/net/neoforged/zipinject/ConsoleTool.java @@ -0,0 +1,148 @@ +package net.neoforged.zipinject; + +import joptsimple.OptionException; +import joptsimple.OptionParser; +import joptsimple.OptionSet; +import joptsimple.OptionSpec; + +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.io.UncheckedIOException; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.nio.file.FileSystem; +import java.nio.file.FileSystems; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.TimeZone; +import java.util.stream.Stream; +import java.util.zip.ZipEntry; +import java.util.zip.ZipInputStream; +import java.util.zip.ZipOutputStream; + +public class ConsoleTool { + public static void main(String[] args) throws Exception { + TimeZone.setDefault(TimeZone.getTimeZone("GMT")); + OptionParser parser = new OptionParser(); + OptionSpec baseO = parser.accepts("base", "The base zip to inject into").withRequiredArg().ofType(File.class).required(); + OptionSpec outputO = parser.accepts("output", "The location of the output zip").withRequiredArg().ofType(File.class).required(); + OptionSpec injectO = parser.accepts("inject", "A zip or directory to inject").withRequiredArg().ofType(File.class); + OptionSpec sourcePrefixO = parser.accepts("path-prefix", "A prefix to strip from source file paths").withRequiredArg(); + OptionSpec packageInfoPackagesO = parser.accepts("inject-package-info", "A prefix that packages that should contain a package-info have").withRequiredArg().ofType(String.class); + + try { + OptionSet options = parser.parse(args); + List injects = options.valuesOf(injectO); + List injectRoots = new ArrayList<>(); + for (File inject : injects) { + if (!inject.isDirectory()) { + if (!inject.exists()) { + throw new IllegalArgumentException("Injection path " + inject + " doesn't exist"); + } + try { + FileSystem fs = FileSystems.newFileSystem(URI.create("jar:" + inject.toURI()), Collections.emptyMap()); + injectRoots.add(fs.getRootDirectories().iterator().next()); + } catch (Exception exception) { + throw new IllegalArgumentException("Injection path " + inject + " is not a zip", exception); + } + } else { + injectRoots.add(inject.toPath().toAbsolutePath()); + } + } + + String packageInfoTemplate = null; + if (options.has(packageInfoPackagesO)) { + packageInfoTemplate = findPackageInfoTemplate(injectRoots); + } + + String sourcePrefix = options.has(sourcePrefixO) ? options.valueOf(sourcePrefixO) : null; + + try (ZipOutputStream zos = new ZipOutputStream(Files.newOutputStream(options.valueOf(outputO).toPath()))) { + copyInputZipContent(options.valueOf(baseO).toPath(), zos, packageInfoTemplate, options.valuesOf(packageInfoPackagesO)); + + for (Path folder : injectRoots) { + try (Stream stream = Files.walk(folder).sorted()) { + stream.filter(Files::isRegularFile).forEach(path -> { + String outputPath = folder.relativize(path).toString().replace('\\', '/'); + try { + if (sourcePrefix != null && outputPath.startsWith(sourcePrefix)) { + outputPath = outputPath.substring(sourcePrefix.length()); + if (outputPath.isEmpty()) return; + } else if (outputPath.equals("package-info-template.java")) { + // Don't include the template file in the output + return; + } + zos.putNextEntry(new ZipEntry(outputPath)); + Files.copy(path, zos); + zos.closeEntry(); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + }); + } + } + } + + // Now close the opened zip file systems + for (Path injectRoot : injectRoots) { + if (injectRoot.getFileSystem().provider().getScheme().equals("jar")) { + injectRoot.getFileSystem().close(); + } + } + } catch (OptionException e) { + parser.printHelpOn(System.out); + e.printStackTrace(); + } + } + + private static String findPackageInfoTemplate(List roots) throws IOException { + // Try to find a package-info-template.java + for (Path injectedSource : roots) { + Path subPath = injectedSource.resolve("package-info-template.java"); + if (Files.isRegularFile(subPath)) { + return new String(Files.readAllBytes(subPath), StandardCharsets.UTF_8); + } + } + return null; + } + + private static void copyInputZipContent(Path inputZipFile, ZipOutputStream zos, String packageInfoTemplateContent, List packagePrefixes) throws IOException { + Set visited = new HashSet<>(); + try (ZipInputStream zis = new ZipInputStream(Files.newInputStream(inputZipFile))) { + ZipEntry entry; + while ((entry = zis.getNextEntry()) != null) { + zos.putNextEntry(entry); + copy(zis, zos); + zos.closeEntry(); + + if (packageInfoTemplateContent != null) { + String pkg = entry.isDirectory() && !entry.getName().endsWith("/") ? entry.getName() : (entry.getName().indexOf('/') == -1 ? "" : entry.getName().substring(0, entry.getName().lastIndexOf('/'))); + if (visited.add(pkg)) { + for (String prefix : packagePrefixes) { + if (pkg.startsWith(prefix)) { + zos.putNextEntry(new ZipEntry(pkg + "/package-info.java")); + zos.write(packageInfoTemplateContent.replace("{PACKAGE}", pkg.replaceAll("/", ".")).getBytes(StandardCharsets.UTF_8)); + zos.closeEntry(); + break; + } + } + } + } + } + } + } + + public static void copy(InputStream input, OutputStream output) throws IOException { + byte[] buffer = new byte[1024]; + int read; + while ((read = input.read(buffer)) != -1) + output.write(buffer, 0, read); + } +} diff --git a/zipinject/src/test/java/net/neoforged/zipinject/ZipInjectTest.java b/zipinject/src/test/java/net/neoforged/zipinject/ZipInjectTest.java new file mode 100644 index 0000000..0128b84 --- /dev/null +++ b/zipinject/src/test/java/net/neoforged/zipinject/ZipInjectTest.java @@ -0,0 +1,128 @@ +package net.neoforged.zipinject; + +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.ByteArrayOutputStream; +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Enumeration; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.zip.ZipEntry; +import java.util.zip.ZipFile; +import java.util.zip.ZipOutputStream; + +public class ZipInjectTest { + @TempDir + File tempDir; + + @Test + void testInjectDir() throws Exception { + final Path injectDir = tempDir.toPath().resolve("injection"); + Files.createDirectories(injectDir); + + Files.write(injectDir.resolve("a.txt"), "file a".getBytes(StandardCharsets.UTF_8)); + Files.write(injectDir.resolve("b.txt"), "file b".getBytes(StandardCharsets.UTF_8)); + + final File base = new File(tempDir, "base.zip"); + writeJar(base, "c.txt", "file c", "d.txt", "file d"); + + final File out = new File(tempDir, "out.zip"); + invoke("--base", base, "--output", out, "--inject", injectDir); + + assertContents(new ZipFile(out), "c.txt", "file c", "d.txt", "file d", "a.txt", "file a", "b.txt", "file b"); + } + + @Test + void testInjectZip() throws Exception { + final File base = new File(tempDir, "base.zip"); + writeJar(base, "c.txt", "file c", "d.txt", "file d"); + + final File inject = new File(tempDir, "inject.zip"); + writeJar(inject, "e.txt", "yes it's e"); + + final File out = new File(tempDir, "out.zip"); + invoke("--base", base, "--output", out, "--inject", inject); + + assertContents(new ZipFile(out), "c.txt", "file c", "d.txt", "file d", "e.txt", "yes it's e"); + } + + @Test + void testInjectWithPrefix() throws Exception { + final File base = new File(tempDir, "base.zip"); + writeJar(base, "c.txt", "file c", "d.txt", "file d"); + + final File inject = new File(tempDir, "inject.zip"); + writeJar(inject, "correct/e.txt", "yes it's e", "correct/sub/f.txt", "sub file"); + + final File out = new File(tempDir, "out.zip"); + invoke("--base", base, "--output", out, "--inject", inject, "--path-prefix", "correct/"); + + assertContents(new ZipFile(out), "c.txt", "file c", "d.txt", "file d", "e.txt", "yes it's e", "sub/f.txt", "sub file"); + } + + @Test + void testInjectPackageInfo() throws Exception { + final File base = new File(tempDir, "base.zip"); + writeJar(base, "com/mojang/A.file", "a file", "com/notmojang/B.file", "b file"); + + final File inject = new File(tempDir, "inject.zip"); + writeJar(inject, "package-info-template.java", "package {PACKAGE};"); + + final File out = new File(tempDir, "out.zip"); + invoke("--base", base, "--output", out, "--inject", inject, "--inject-package-info", "com/mojang"); + + assertContents(new ZipFile(out), "com/mojang/A.file", "a file", "com/mojang/package-info.java", "package com.mojang;", "com/notmojang/B.file", "b file"); + } + + private static void assertContents(ZipFile file, String... entries) throws IOException { + Map expectedFiles = new LinkedHashMap<>(); + for (int i = 0; i < entries.length; i += 2) { + expectedFiles.put(entries[i], entries[i + 1]); + } + + Map actualFiles = new LinkedHashMap<>(); + Enumeration contained = file.entries(); + while (contained.hasMoreElements()) { + ZipEntry entry = contained.nextElement(); + if (entry.isDirectory()) continue; + ByteArrayOutputStream bos = new ByteArrayOutputStream(); + ConsoleTool.copy(file.getInputStream(entry), bos); + actualFiles.put(entry.getName(), new String(bos.toByteArray(), StandardCharsets.UTF_8)); + } + file.close(); + + Assertions.assertThat(actualFiles) + .containsExactlyEntriesOf(expectedFiles); + } + + private static void writeJar(File location, String... entries) throws IOException { + try (final ZipOutputStream zos = new ZipOutputStream(new FileOutputStream(location))) { + for (int i = 0; i < entries.length; i += 2) { + zos.putNextEntry(new ZipEntry(entries[i])); + zos.write(entries[i + 1].getBytes(StandardCharsets.UTF_8)); + zos.closeEntry(); + } + } + } + + private static void invoke(Object... args) throws Exception { + final String[] finalArgs = new String[args.length]; + for (int i = 0; i < args.length; i++) { + if (args[i] instanceof File) { + finalArgs[i] = ((File) args[i]).getAbsolutePath(); + } else if (args[i] instanceof Path) { + finalArgs[i] = ((Path) args[i]).toAbsolutePath().toString(); + } else { + finalArgs[i] = args[i].toString(); + } + } + ConsoleTool.main(finalArgs); + } +}