Add a tool for zip injection

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')

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'
+        }
+    }
+}

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<File> baseO = parser.accepts("base", "The base zip to inject into").withRequiredArg().ofType(File.class).required();
+        OptionSpec<File> outputO = parser.accepts("output", "The location of the output zip").withRequiredArg().ofType(File.class).required();
+        OptionSpec<File> injectO = parser.accepts("inject", "A zip or directory to inject").withRequiredArg().ofType(File.class);
+        OptionSpec<String> sourcePrefixO = parser.accepts("path-prefix", "A prefix to strip from source file paths").withRequiredArg();
+        OptionSpec<String> 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<File> injects = options.valuesOf(injectO);
+            List<Path> 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<Path> 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<Path> 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<String> packagePrefixes) throws IOException {
+        Set<String> 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);
+    }
+}

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<String, String> expectedFiles = new LinkedHashMap<>();
+        for (int i = 0; i < entries.length; i += 2) {
+            expectedFiles.put(entries[i], entries[i + 1]);
+        }
+
+        Map<String, String> actualFiles = new LinkedHashMap<>();
+        Enumeration<? extends ZipEntry> 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);
+    }
+}