Use ArchiveStreamReader in controller-server

7 files changed, 137 insertions, 169 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ApplicationPackage.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ApplicationPackage.java
index 5c7ae5041fe..258884a4d11 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ApplicationPackage.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ApplicationPackage.java
@@ -4,6 +4,9 @@ package com.yahoo.vespa.hosted.controller.application.pkg;
 import com.google.common.hash.Funnel;
 import com.google.common.hash.Hashing;
 import com.yahoo.component.Version;
+import com.yahoo.compress.ArchiveStreamReader;
+import com.yahoo.compress.ArchiveStreamReader.ArchiveFile;
+import com.yahoo.compress.ArchiveStreamReader.Options;
 import com.yahoo.config.application.FileSystemWrapper;
 import com.yahoo.config.application.FileSystemWrapper.FileWrapper;
 import com.yahoo.config.application.XmlPreProcessor;
@@ -24,6 +27,7 @@ import com.yahoo.yolean.Exceptions;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.InputStreamReader;
+import java.io.OutputStream;
 import java.nio.file.NoSuchFileException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -40,6 +44,7 @@ import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.SortedMap;
+import java.util.TreeMap;
 import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.function.Function;
 import java.util.function.Predicate;
@@ -108,7 +113,7 @@ public class ApplicationPackage {
 
         this.trustedCertificates = files.get(trustedCertificatesFile).map(bytes -> X509CertificateUtils.certificateListFromPem(new String(bytes, UTF_8))).orElse(List.of());
 
-        this.bundleHash = calculateBundleHash();
+        this.bundleHash = calculateBundleHash(zippedContent);
 
         preProcessAndPopulateCache();
     }
@@ -120,7 +125,7 @@ public class ApplicationPackage {
         byte[] certificatesBytes = X509CertificateUtils.toPem(trustedCertificates).getBytes(UTF_8);
 
         ByteArrayOutputStream modified = new ByteArrayOutputStream(zippedContent.length + certificatesBytes.length);
-        ZipStreamReader.transferAndWrite(modified, new ByteArrayInputStream(zippedContent), trustedCertificatesFile, certificatesBytes);
+        ZipEntries.transferAndWrite(modified, new ByteArrayInputStream(zippedContent), trustedCertificatesFile, certificatesBytes);
         return new ApplicationPackage(modified.toByteArray());
     }
 
@@ -227,15 +232,23 @@ public class ApplicationPackage {
     }
 
     // Hashes all files and settings that require a deployment to be forwarded to configservers
-    private String calculateBundleHash() {
+    private String calculateBundleHash(byte[] zippedContent) {
         Predicate<String> entryMatcher = name -> ! name.endsWith(deploymentFile) && ! name.endsWith(buildMetaFile);
-        SortedMap<String, Long> entryCRCs = ZipStreamReader.getEntryCRCs(new ByteArrayInputStream(zippedContent), entryMatcher);
-        Funnel<SortedMap<String, Long>> funnel = (from, into) -> from.entrySet().forEach(entry -> {
-            into.putBytes(entry.getKey().getBytes());
-            into.putLong(entry.getValue());
+        SortedMap<String, Long> crcByEntry = new TreeMap<>();
+        Options options = Options.standard().pathPredicate(entryMatcher);
+        ArchiveFile file;
+        try (ArchiveStreamReader reader = ArchiveStreamReader.ofZip(new ByteArrayInputStream(zippedContent), options)) {
+            OutputStream discard = OutputStream.nullOutputStream();
+            while ((file = reader.readNextTo(discard)) != null) {
+                crcByEntry.put(file.path().toString(), file.crc32().orElse(-1));
+            }
+        }
+        Funnel<SortedMap<String, Long>> funnel = (from, into) -> from.forEach((key, value) -> {
+            into.putBytes(key.getBytes());
+            into.putLong(value);
         });
         return Hashing.sha1().newHasher()
-                      .putObject(entryCRCs, funnel)
+                      .putObject(crcByEntry, funnel)
                       .putInt(deploymentSpec.deployableHashCode())
                       .hash().toString();
     }
@@ -285,13 +298,13 @@ public class ApplicationPackage {
         }
 
         private Map<Path, Optional<byte[]>> read(Collection<String> names) {
-            var entries = new ZipStreamReader(new ByteArrayInputStream(zip),
-                                              name -> names.contains(withoutLegacyDir(name)),
-                                              maxSize,
-                                              true)
-                    .entries().stream()
-                    .collect(toMap(entry -> Paths.get(withoutLegacyDir(entry.zipEntry().getName())).normalize(),
-                             ZipStreamReader.ZipEntryWithContent::content));
+            var entries = ZipEntries.from(zip,
+                                          name -> names.contains(withoutLegacyDir(name)),
+                                          maxSize,
+                                          true)
+                                    .asList().stream()
+                                    .collect(toMap(entry -> Paths.get(withoutLegacyDir(entry.name())).normalize(),
+                                                   ZipEntries.ZipEntryWithContent::content));
             names.stream().map(Paths::get).forEach(path -> entries.putIfAbsent(path.normalize(), Optional.empty()));
             return entries;
         }
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ApplicationPackageDiff.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ApplicationPackageDiff.java
index 97810b9de80..e18f6247cb1 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ApplicationPackageDiff.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ApplicationPackageDiff.java
@@ -15,7 +15,7 @@ import java.util.Optional;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
-import static com.yahoo.vespa.hosted.controller.application.pkg.ZipStreamReader.ZipEntryWithContent;
+import static com.yahoo.vespa.hosted.controller.application.pkg.ZipEntries.ZipEntryWithContent;
 
 /**
  * @author freva
@@ -75,8 +75,8 @@ public class ApplicationPackageDiff {
     }
 
     private static Map<String, ZipEntryWithContent> readContents(ApplicationPackage app, int maxFileSizeToDiff) {
-        return new ZipStreamReader(new ByteArrayInputStream(app.zippedContent()), entry -> true, maxFileSizeToDiff, false).entries().stream()
-                .collect(Collectors.toMap(entry -> entry.zipEntry().getName(), e -> e));
+        return ZipEntries.from(app.zippedContent(), entry -> true, maxFileSizeToDiff, false).asList().stream()
+                         .collect(Collectors.toMap(ZipEntryWithContent::name, e -> e));
     }
 
     private static List<String> lines(byte[] data) {
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ZipEntries.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ZipEntries.java
new file mode 100644
index 00000000000..e35cb55b87a
--- /dev/null
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ZipEntries.java
@@ -0,0 +1,98 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.controller.application.pkg;
+
+import com.yahoo.compress.ArchiveStreamReader;
+import com.yahoo.compress.ArchiveStreamReader.ArchiveFile;
+import com.yahoo.compress.ArchiveStreamReader.Options;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.io.UncheckedIOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.function.Predicate;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipInputStream;
+import java.util.zip.ZipOutputStream;
+
+/**
+ * A list of entries read from a ZIP archive, and their contents.
+ *
+ * @author bratseth
+ */
+public class ZipEntries {
+
+    private final List<ZipEntryWithContent> entries;
+
+    private ZipEntries(List<ZipEntryWithContent> entries) {
+        this.entries = List.copyOf(Objects.requireNonNull(entries));
+    }
+
+    /** Copies the zipped content from in to out, adding/overwriting an entry with the given name and content. */
+    public static void transferAndWrite(OutputStream out, InputStream in, String name, byte[] content) {
+        try (ZipOutputStream zipOut = new ZipOutputStream(out);
+             ZipInputStream zipIn = new ZipInputStream(in)) {
+            for (ZipEntry entry = zipIn.getNextEntry(); entry != null; entry = zipIn.getNextEntry()) {
+                if (entry.getName().equals(name))
+                    continue;
+
+                zipOut.putNextEntry(new ZipEntry(entry.getName()));
+                zipIn.transferTo(zipOut);
+                zipOut.closeEntry();
+            }
+            zipOut.putNextEntry(new ZipEntry(name));
+            zipOut.write(content);
+            zipOut.closeEntry();
+        }
+        catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    /** Read ZIP entries from inputStream */
+    public static ZipEntries from(byte[] zip, Predicate<String> entryNameMatcher, int maxEntrySizeInBytes, boolean throwIfEntryExceedsMaxSize) {
+        Options options = Options.standard()
+                                 .pathPredicate(entryNameMatcher)
+                                 .sizeLimit(2 * (long) Math.pow(1024, 3)) // 2 GB
+                                 .entrySizeLimit(maxEntrySizeInBytes)
+                                 .truncateEntry(!throwIfEntryExceedsMaxSize);
+        List<ZipEntryWithContent> entries = new ArrayList<>();
+        try (ArchiveStreamReader reader = ArchiveStreamReader.ofZip(new ByteArrayInputStream(zip), options)) {
+            ArchiveFile file;
+            ByteArrayOutputStream baos = new ByteArrayOutputStream();
+            while ((file = reader.readNextTo(baos)) != null) {
+                entries.add(new ZipEntryWithContent(file.path().toString(),
+                                                    Optional.of(baos.toByteArray()).filter(b -> b.length > 0),
+                                                    file.size()));
+                baos.reset();
+            }
+        }
+        return new ZipEntries(entries);
+    }
+
+    public List<ZipEntryWithContent> asList() { return entries; }
+
+    public static class ZipEntryWithContent {
+
+        private final String name;
+        private final Optional<byte[]> content;
+        private final long size;
+
+        public ZipEntryWithContent(String name, Optional<byte[]> content, long size) {
+            this.name = name;
+            this.content = content;
+            this.size = size;
+        }
+
+        public String name() { return name; }
+        public byte[] contentOrThrow() { return content.orElseThrow(); }
+        public Optional<byte[]> content() { return content; }
+        public long size() { return size; }
+    }
+
+}
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ZipStreamReader.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ZipStreamReader.java
deleted file mode 100644
index 174ac4cb8b0..00000000000
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/pkg/ZipStreamReader.java
+++ /dev/null
@@ -1,138 +0,0 @@
-// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.vespa.hosted.controller.application.pkg;
-
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-import java.io.UncheckedIOException;
-import java.nio.file.Path;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.SortedMap;
-import java.util.TreeMap;
-import java.util.function.Predicate;
-import java.util.zip.ZipEntry;
-import java.util.zip.ZipInputStream;
-import java.util.zip.ZipOutputStream;
-
-/**
- * @author bratseth
- */
-public class ZipStreamReader {
-
-    private final List<ZipEntryWithContent> entries = new ArrayList<>();
-    private final int maxEntrySizeInBytes;
-
-    public ZipStreamReader(InputStream input, Predicate<String> entryNameMatcher, int maxEntrySizeInBytes, boolean throwIfEntryExceedsMaxSize) {
-        this.maxEntrySizeInBytes = maxEntrySizeInBytes;
-        try (ZipInputStream zipInput = new ZipInputStream(input)) {
-            ZipEntry zipEntry;
-
-            while (null != (zipEntry = zipInput.getNextEntry())) {
-                if (!entryNameMatcher.test(requireName(zipEntry.getName()))) continue;
-                entries.add(readContent(zipEntry, zipInput, throwIfEntryExceedsMaxSize));
-            }
-        } catch (IOException e) {
-            throw new UncheckedIOException("IO error reading zip content", e);
-        }
-    }
-
-    /** Copies the zipped content from in to out, adding/overwriting an entry with the given name and content. */
-    public static void transferAndWrite(OutputStream out, InputStream in, String name, byte[] content) {
-        try (ZipOutputStream zipOut = new ZipOutputStream(out);
-             ZipInputStream zipIn = new ZipInputStream(in)) {
-            for (ZipEntry entry = zipIn.getNextEntry(); entry != null; entry = zipIn.getNextEntry()) {
-                if (entry.getName().equals(name))
-                    continue;
-
-                zipOut.putNextEntry(new ZipEntry(entry.getName()));
-                zipIn.transferTo(zipOut);
-                zipOut.closeEntry();
-            }
-            zipOut.putNextEntry(new ZipEntry(name));
-            zipOut.write(content);
-            zipOut.closeEntry();
-        }
-        catch (IOException e) {
-            throw new UncheckedIOException(e);
-        }
-    }
-
-    public static SortedMap<String, Long> getEntryCRCs(InputStream in, Predicate<String> entryNameMatcher) {
-        SortedMap<String, Long> entryCRCs = new TreeMap<>();
-        byte[] buffer = new byte[2048];
-        try (ZipInputStream zipIn = new ZipInputStream(in)) {
-            for (ZipEntry entry = zipIn.getNextEntry(); entry != null; entry = zipIn.getNextEntry()) {
-                if ( ! entryNameMatcher.test(entry.getName()))
-                    continue;
-                // CRC is not set until entry is read
-                while ( -1 != zipIn.read(buffer)){}
-                entryCRCs.put(entry.getName(), entry.getCrc());
-            }
-        } catch (IOException e) {
-            throw new UncheckedIOException(e);
-        }
-        return entryCRCs;
-    }
-
-    private ZipEntryWithContent readContent(ZipEntry zipEntry, ZipInputStream zipInput, boolean throwIfEntryExceedsMaxSize) {
-        try (ByteArrayOutputStream bis = new ByteArrayOutputStream()) {
-            byte[] buffer = new byte[2048];
-            int read;
-            long size = 0;
-            while ( -1 != (read = zipInput.read(buffer))) {
-                size += read;
-                if (size > maxEntrySizeInBytes) {
-                    if (throwIfEntryExceedsMaxSize) throw new IllegalArgumentException(
-                            "Entry in zip content exceeded size limit of " + maxEntrySizeInBytes + " bytes");
-                } else bis.write(buffer, 0, read);
-            }
-
-            boolean hasContent = size <= maxEntrySizeInBytes;
-            return new ZipEntryWithContent(zipEntry,
-                    Optional.of(bis).filter(__ -> hasContent).map(ByteArrayOutputStream::toByteArray),
-                    size);
-        } catch (IOException e) {
-            throw new UncheckedIOException("Failed reading from zipped content", e);
-        }
-    }
-
-    public List<ZipEntryWithContent> entries() { return Collections.unmodifiableList(entries); }
-
-    private static String requireName(String name) {
-        if (List.of(name.split("/")).contains("..") ||
-            !trimTrailingSlash(name).equals(Path.of(name).normalize().toString())) {
-            throw new IllegalArgumentException("Unexpected non-normalized path found in zip content: '" + name + "'");
-        }
-        return name;
-    }
-
-    private static String trimTrailingSlash(String name) {
-        if (name.endsWith("/")) return name.substring(0, name.length() - 1);
-        return name;
-    }
-
-    public static class ZipEntryWithContent {
-
-        private final ZipEntry zipEntry;
-        private final Optional<byte[]> content;
-        private final long size;
-
-        public ZipEntryWithContent(ZipEntry zipEntry, Optional<byte[]> content, long size) {
-            this.zipEntry = zipEntry;
-            this.content = content;
-            this.size = size;
-        }
-
-        public ZipEntry zipEntry() { return zipEntry; }
-        public byte[] contentOrThrow() { return content.orElseThrow(); }
-        public Optional<byte[]> content() { return content; }
-        public long size() { return size; }
-    }
-
-}
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/application/pkg/ApplicationPackageTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/application/pkg/ApplicationPackageTest.java
index 709a7967a5e..4e155e937b9 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/application/pkg/ApplicationPackageTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/application/pkg/ApplicationPackageTest.java
@@ -5,7 +5,6 @@ import com.yahoo.config.application.api.DeploymentSpec;
 import com.yahoo.config.application.api.ValidationId;
 import org.junit.Test;
 
-import java.io.ByteArrayInputStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.time.Instant;
@@ -154,9 +153,9 @@ public class ApplicationPackageTest {
     }
 
     private static Map<String, String> unzip(byte[] zip) {
-        return new ZipStreamReader(new ByteArrayInputStream(zip), __ -> true, 1 << 10, true)
-                .entries().stream()
-                .collect(Collectors.toMap(entry -> entry.zipEntry().getName(),
+        return ZipEntries.from(zip, __ -> true, 1 << 10, true)
+                         .asList().stream()
+                         .collect(Collectors.toMap(ZipEntries.ZipEntryWithContent::name,
                                           entry -> new String(entry.contentOrThrow(), UTF_8)));
     }
 
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/application/pkg/ZipStreamReaderTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/application/pkg/ZipEntriesTest.java
index 33c18d123d2..5777e3e246c 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/application/pkg/ZipStreamReaderTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/application/pkg/ZipEntriesTest.java
@@ -8,7 +8,6 @@ import com.yahoo.security.X509CertificateBuilder;
 import org.junit.Test;
 
 import javax.security.auth.x500.X500Principal;
-import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.UncheckedIOException;
@@ -32,21 +31,21 @@ import static org.junit.Assert.fail;
 /**
  * @author mpolden
  */
-public class ZipStreamReaderTest {
+public class ZipEntriesTest {
 
     @Test
     public void test_size_limit() {
         Map<String, String> entries = Map.of("foo.xml", "foobar");
         try {
-            new ZipStreamReader(new ByteArrayInputStream(zip(entries)), "foo.xml"::equals, 1, true);
+            ZipEntries.from(zip(entries), "foo.xml"::equals, 1, true);
             fail("Expected exception");
         } catch (IllegalArgumentException ignored) {}
 
         entries = Map.of("foo.xml", "foobar",
                          "foo.jar", "0".repeat(100) // File not extracted and thus not subject to size limit
         );
-        ZipStreamReader reader = new ZipStreamReader(new ByteArrayInputStream(zip(entries)), "foo.xml"::equals, 10, true);
-        byte[] extracted = reader.entries().get(0).contentOrThrow();
+        ZipEntries reader = ZipEntries.from(zip(entries), "foo.xml"::equals, 10, true);
+        byte[] extracted = reader.asList().get(0).contentOrThrow();
         assertEquals("foobar", new String(extracted, StandardCharsets.UTF_8));
     }
 
@@ -65,7 +64,7 @@ public class ZipStreamReaderTest {
         );
         tests.forEach((name, expectException) -> {
             try {
-                new ZipStreamReader(new ByteArrayInputStream(zip(Map.of(name, "foo"))), name::equals, 1024, true);
+                ZipEntries.from(zip(Map.of(name, "foo")), name::equals, 1024, true);
                 assertFalse("Expected exception for '" + name + "'", expectException);
             } catch (IllegalArgumentException ignored) {
                 assertTrue("Unexpected exception for '" + name + "'", expectException);
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/DeploymentContext.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/DeploymentContext.java
index 989a7c31821..86c21839c96 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/DeploymentContext.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/DeploymentContext.java
@@ -32,7 +32,6 @@ import com.yahoo.vespa.hosted.controller.application.pkg.ApplicationPackage;
 import com.yahoo.vespa.hosted.controller.application.Deployment;
 import com.yahoo.vespa.hosted.controller.application.EndpointId;
 import com.yahoo.vespa.hosted.controller.application.TenantAndApplicationId;
-import com.yahoo.vespa.hosted.controller.application.pkg.ZipStreamReader;
 import com.yahoo.vespa.hosted.controller.integration.ConfigServerMock;
 import com.yahoo.vespa.hosted.controller.maintenance.JobRunner;
 import com.yahoo.vespa.hosted.controller.maintenance.NameServiceDispatcher;
@@ -41,8 +40,6 @@ import com.yahoo.vespa.hosted.controller.routing.RoutingPolicy;
 import com.yahoo.vespa.hosted.controller.routing.RoutingPolicyId;
 
 import javax.security.auth.x500.X500Principal;
-import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
 import java.math.BigInteger;
 import java.security.KeyPair;
 import java.security.cert.X509Certificate;