| /* |
| * Copyright 2013-present Facebook, Inc. |
| * |
| * Licensed under the Apache License, Version 2.0 (the "License"); you may |
| * not use this file except in compliance with the License. You may obtain |
| * a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT |
| * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the |
| * License for the specific language governing permissions and limitations |
| * under the License. |
| */ |
| |
| package com.facebook.buck.zip; |
| |
| import com.facebook.buck.timing.Clock; |
| import com.google.common.base.Preconditions; |
| import com.google.common.collect.Maps; |
| import com.google.common.hash.Hashing; |
| |
| import java.io.BufferedOutputStream; |
| import java.io.File; |
| import java.io.FileOutputStream; |
| import java.io.IOException; |
| import java.io.OutputStream; |
| import java.nio.file.FileVisitResult; |
| import java.nio.file.Files; |
| import java.nio.file.Path; |
| import java.nio.file.SimpleFileVisitor; |
| import java.nio.file.attribute.BasicFileAttributes; |
| import java.util.Map; |
| import java.util.zip.ZipEntry; |
| import java.util.zip.ZipException; |
| |
| /** |
| * An implementation of an {@link OutputStream} for zip files that allows newer entries to overwrite |
| * or refresh previously written entries. |
| * <p> |
| * This class works by spooling the bytes of each entry to a temporary holding file named after the |
| * name of the {@link ZipEntry} being stored. Once the stream is closed, these files are spooled |
| * off disk and written to the OutputStream given to the constructor. |
| */ |
| public class OverwritingZipOutputStream extends CustomZipOutputStream { |
| // Attempt to maintain ordering of files that are added. |
| private final Map<File, EntryAccounting> entries = Maps.newLinkedHashMap(); |
| private final File scratchDir; |
| private final Clock clock; |
| private EntryAccounting currentEntry; |
| /** Place-holder for bytes. */ |
| private OutputStream currentOutput; |
| |
| public OverwritingZipOutputStream(Clock clock, OutputStream out) { |
| super(out); |
| this.clock = Preconditions.checkNotNull(clock); |
| |
| try { |
| scratchDir = Files.createTempDirectory("overwritingzip").toFile(); |
| // Reading the source, it seems like the temp dir isn't scheduled for deletion. We will delete |
| // the directory when we close the stream, but if that method is never called, we'd leave |
| // cruft on the FS. It's not foolproof, but try and avoid that. |
| scratchDir.deleteOnExit(); |
| } catch (IOException e) { |
| throw new RuntimeException(e); |
| } |
| } |
| |
| @Override |
| protected void actuallyPutNextEntry(ZipEntry entry) throws IOException { |
| // We calculate the actual offset when closing the stream, so 0 is fine. |
| currentEntry = new EntryAccounting(clock, entry, /* currentOffset */ 0); |
| |
| long md5 = Hashing.md5().hashUnencodedChars(entry.getName()).asLong(); |
| String name = String.valueOf(md5); |
| |
| File file = new File(scratchDir, name); |
| entries.put(file, currentEntry); |
| if (file.exists() && !file.delete()) { |
| throw new ZipException("Unable to delete existing file: " + entry.getName()); |
| } |
| currentOutput = new BufferedOutputStream(new FileOutputStream(file)); |
| } |
| |
| @Override |
| protected void actuallyCloseEntry() throws IOException { |
| // We'll close the entry once we have the ultimate output stream and know the entry's location |
| // within the generated zip. |
| currentOutput.close(); |
| currentOutput = null; |
| currentEntry = null; |
| } |
| |
| @Override |
| protected void actuallyWrite(byte[] b, int off, int len) throws IOException { |
| currentEntry.write(currentOutput, b, off, len); |
| } |
| |
| @Override |
| protected void actuallyClose() throws IOException { |
| long currentOffset = 0; |
| |
| for (Map.Entry<File, EntryAccounting> mapEntry : entries.entrySet()) { |
| EntryAccounting entry = mapEntry.getValue(); |
| entry.setOffset(currentOffset); |
| currentOffset += entry.writeLocalFileHeader(delegate); |
| currentOffset += Files.copy(mapEntry.getKey().toPath(), delegate); |
| currentOffset += entry.close(delegate); |
| } |
| |
| new CentralDirectory().writeCentralDirectory(delegate, currentOffset, entries.values()); |
| |
| delegate.close(); |
| |
| // Ideally we'd just do this, but that introduces some nasty circular references. *sigh* Instead |
| // we'll do this the tedious way by hand. |
| // MoreFiles.deleteRecursively(scratchDir.toPath()); |
| |
| SimpleFileVisitor<Path> visitor = new SimpleFileVisitor<Path>() { |
| @Override |
| public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException { |
| Files.delete(file); |
| return FileVisitResult.CONTINUE; |
| } |
| |
| @Override |
| public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException { |
| if (exc == null) { |
| Files.delete(dir); |
| return FileVisitResult.CONTINUE; |
| } |
| throw exc; |
| } |
| }; |
| Files.walkFileTree(scratchDir.toPath(), visitor); |
| } |
| } |