Handle MPQ maps without HM3W headers (#1258)

Author: Frotty

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstio/languageserver/ProjectConfigBuilder.java

@@ -107,7 +107,7 @@ public static MapRequest.CompilationResult apply(WurstProjectConfigData projectC
         w3I.write(result.w3i);
 
         // Apply map header (this is cheap, so we always do it)
-        applyMapHeader(projectConfig, targetMap);
+        applyMapHeader(projectConfig, targetMap, w3I.getPlayers().size(), w3I.getMapName(), w3I.getFlags().toInt());
 
         // Update the manifest with new config hash (must open writable to insert)
         try (MpqEditor mpq = MpqEditorFactory.getEditor(Optional.of(targetMap), false)) {
@@ -356,20 +356,50 @@ private static void setVolatilePlayerConfig(WurstProjectBuildPlayer wplayer, W3I
         }
     }
 
-    private static void applyMapHeader(WurstProjectConfigData projectConfig, File targetMap) throws IOException {
+    private static void applyMapHeader(WurstProjectConfigData projectConfig, File targetMap,
+                                       int existingPlayerCount, String existingMapName,
+                                       int existingMapFlags) throws IOException {
         boolean shouldWrite = false;
-        MapHeader mapHeader = MapHeader.ofFile(targetMap);
-        if (!projectConfig.buildMapData().players().isEmpty()) {
-            mapHeader.setMaxPlayersCount(projectConfig.buildMapData().players().size());
+        WurstProjectBuildMapData buildMapData = projectConfig.buildMapData();
+        if (buildMapData.players().isEmpty() && StringUtils.isBlank(buildMapData.name())) {
+            return;
+        }
+
+        // A Warcraft III map may omit the optional 512-byte HM3W prefix and start
+        // directly with its MPQ archive. MapHeader.ofFile only reads the prefix,
+        // so use a new header in that case; writeToMapFile will insert it before
+        // the archive.
+        boolean hasNoMapHeader = startsWithMpqArchive(targetMap);
+        MapHeader mapHeader = hasNoMapHeader
+            ? new MapHeader()
+            : MapHeader.ofFile(targetMap);
+        if (!buildMapData.players().isEmpty()) {
+            mapHeader.setMaxPlayersCount(buildMapData.players().size());
             shouldWrite = true;
+        } else if (hasNoMapHeader) {
+            mapHeader.setMaxPlayersCount(existingPlayerCount);
+        }
+        if (hasNoMapHeader && StringUtils.isBlank(buildMapData.name())) {
+            mapHeader.setMapName(existingMapName);
         }
-        if (StringUtils.isNotBlank(projectConfig.buildMapData().name())) {
-            mapHeader.setMapName(projectConfig.buildMapData().name());
+        if (hasNoMapHeader) {
+            mapHeader.setFlags(existingMapFlags);
+        }
+        if (StringUtils.isNotBlank(buildMapData.name())) {
+            mapHeader.setMapName(buildMapData.name());
             shouldWrite = true;
         }
         if (shouldWrite) {
             WLogger.info("Applying map header");
             mapHeader.writeToMapFile(targetMap);
         }
     }
+
+    private static boolean startsWithMpqArchive(File targetMap) throws IOException {
+        try (InputStream input = new FileInputStream(targetMap)) {
+            byte[] startToken = input.readNBytes(4);
+            return startToken.length == 4
+                && new String(startToken, StandardCharsets.US_ASCII).startsWith("MPQ");
+        }
+    }
 }

de.peeeq.wurstscript/src/test/java/tests/wurstscript/tests/WurstBuildConfigTests.java

@@ -1,16 +1,20 @@
 package tests.wurstscript.tests;
 
 import org.wurstscript.projectconfig.WurstProjectConfigData;
+import org.wurstscript.projectconfig.WurstProjectBuildMapData;
+import org.wurstscript.projectconfig.WurstProjectBuildPlayer;
 import de.peeeq.wurstio.languageserver.WFile;
 import de.peeeq.wurstio.languageserver.ProjectConfigBuilder;
 import de.peeeq.wurstio.languageserver.WurstBuildConfig;
 import de.peeeq.wurstio.languageserver.WurstCommands;
 import de.peeeq.wurstio.utils.W3InstallationData;
+import net.moonlightflower.wc3libs.bin.app.MapHeader;
 import net.moonlightflower.wc3libs.port.GameVersion;
 import org.testng.annotations.Test;
 
 import java.io.File;
 import java.lang.reflect.Method;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.List;
@@ -223,6 +227,100 @@ public void configInjectionPrefersPinnedPatchOverDetectedInstallVersion() throws
         );
     }
 
+    @Test
+    public void mapHeaderConfigAcceptsAnArchiveWithoutAnHm3wPrefix() throws Exception {
+        Path mapWithoutHeader = Files.createTempFile("wurst-map-without-header", ".w3x");
+        byte[] original = new byte[1024];
+        original[0] = 'M';
+        original[1] = 'P';
+        original[2] = 'Q';
+        original[3] = 0x1a;
+        Files.write(mapWithoutHeader, original);
+
+        Method applyMapHeader = ProjectConfigBuilder.class.getDeclaredMethod(
+            "applyMapHeader", WurstProjectConfigData.class, File.class, int.class, String.class, int.class
+        );
+        applyMapHeader.setAccessible(true);
+
+        applyMapHeader.invoke(
+            null,
+            new WurstProjectConfigData(
+                "Test",
+                List.of(),
+                new WurstProjectBuildMapData("Configured map", null, null, null, null, List.of(), List.of()),
+                null,
+                null
+            ),
+            mapWithoutHeader.toFile(),
+            3,
+            "Existing map",
+            0x1234
+        );
+
+        assertEquals(Files.readAllBytes(mapWithoutHeader)[0], (byte) 'H');
+        assertEquals(MapHeader.ofFile(mapWithoutHeader.toFile()).getMaxPlayersCount(), 3);
+        assertEquals(Files.size(mapWithoutHeader), original.length + 512);
+    }
+
+    @Test
+    public void mapHeaderConfigDoesNotReadAnArchiveWhenNothingNeedsChanging() throws Exception {
+        Path mapWithoutHeader = Files.createTempFile("wurst-map-without-header-noop", ".w3x");
+        byte[] original = new byte[1024];
+        original[0] = 'M';
+        original[1] = 'P';
+        original[2] = 'Q';
+        original[3] = 0x1a;
+        Files.write(mapWithoutHeader, original);
+
+        Method applyMapHeader = ProjectConfigBuilder.class.getDeclaredMethod(
+            "applyMapHeader", WurstProjectConfigData.class, File.class, int.class, String.class, int.class
+        );
+        applyMapHeader.setAccessible(true);
+        applyMapHeader.invoke(null, WurstProjectConfigData.empty(), mapWithoutHeader.toFile(), 0, null, 0);
+
+        assertEquals(Files.readAllBytes(mapWithoutHeader), original);
+    }
+
+    @Test
+    public void mapHeaderConfigPreservesExistingMapNameWhenOnlyPlayersChange() throws Exception {
+        Path mapWithoutHeader = Files.createTempFile("wurst-map-without-header-name", ".w3x");
+        byte[] original = new byte[1024];
+        original[0] = 'M';
+        original[1] = 'P';
+        original[2] = 'Q';
+        original[3] = 0x1a;
+        Files.write(mapWithoutHeader, original);
+
+        Method applyMapHeader = ProjectConfigBuilder.class.getDeclaredMethod(
+            "applyMapHeader", WurstProjectConfigData.class, File.class, int.class, String.class, int.class
+        );
+        applyMapHeader.setAccessible(true);
+
+        applyMapHeader.invoke(
+            null,
+            new WurstProjectConfigData(
+                "Test",
+                List.of(),
+                new WurstProjectBuildMapData(
+                    "", null, null, null, null,
+                    List.of(new WurstProjectBuildPlayer(0, null, null, null, null)),
+                    List.of()
+                ),
+                null,
+                null
+            ),
+            mapWithoutHeader.toFile(),
+            4,
+            "Existing map",
+            0x1234
+        );
+
+        byte[] result = Files.readAllBytes(mapWithoutHeader);
+        assertEquals(MapHeader.ofFile(mapWithoutHeader.toFile()).getMaxPlayersCount(), 1);
+        assertEquals(MapHeader.ofFile(mapWithoutHeader.toFile()).getFlags(), 0x1234);
+        assertTrue(new String(result, StandardCharsets.UTF_8).contains("Existing map"));
+    }
+
     private static String calculateProjectConfigHash(File buildDir) throws Exception {
         Method method = ProjectConfigBuilder.class.getDeclaredMethod(
             "calculateProjectConfigHash",