Add the ability to do a filtered copy of a ProtoApk

RELNOTES: None
PiperOrigin-RevId: 206949407

1 files changed, 157 insertions, 15 deletions

diff --git a/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoApk.java b/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoApk.java
index 5a13ff8e0a..a5fe4cd854 100644
--- a/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoApk.java
+++ b/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoApk.java
@@ -20,13 +20,13 @@ import com.android.aapt.Resources.CompoundValue;
 import com.android.aapt.Resources.ConfigValue;
 import com.android.aapt.Resources.Entry;
 import com.android.aapt.Resources.FileReference;
-import com.android.aapt.Resources.FileReference.Type;
 import com.android.aapt.Resources.Item;
 import com.android.aapt.Resources.Package;
 import com.android.aapt.Resources.Plural;
 import com.android.aapt.Resources.Reference;
 import com.android.aapt.Resources.ResourceTable;
 import com.android.aapt.Resources.Style;
+import com.android.aapt.Resources.Type;
 import com.android.aapt.Resources.Value;
 import com.android.aapt.Resources.XmlAttribute;
 import com.android.aapt.Resources.XmlElement;
@@ -35,45 +35,144 @@ import com.android.resources.ResourceType;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import java.io.Closeable;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.OutputStream;
 import java.io.UnsupportedEncodingException;
 import java.net.URI;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.nio.file.FileSystem;
 import java.nio.file.FileSystems;
+import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.StandardOpenOption;
+import java.nio.file.attribute.BasicFileAttributes;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.BiPredicate;
 
-/** Provides an interface to an apk in proto format. */
-public class ProtoApk {
+/**
+ * Provides an interface to an apk in proto format. Since the apk is backed by a zip, it is
+ * important to close the ProtoApk when done.
+ */
+public class ProtoApk implements Closeable {
 
   private static final String RESOURCE_TABLE = "resources.pb";
   private static final String MANIFEST = "AndroidManifest.xml";
+  private static final String RES_DIRECTORY = "res";
+  private final URI uri;
   private final FileSystem apkFileSystem;
 
-  private ProtoApk(FileSystem apkFileSystem) {
+  private ProtoApk(URI uri, FileSystem apkFileSystem) {
+    this.uri = uri;
     this.apkFileSystem = apkFileSystem;
   }
 
   /** Reads a ProtoApk from a path and verifies that it is in the expected format. */
   public static ProtoApk readFrom(Path apkPath) throws IOException {
-    final FileSystem apkFileSystem =
-        FileSystems.newFileSystem(URI.create("jar:" + apkPath.toUri()), ImmutableMap.of());
+    final URI uri = URI.create("jar:" + apkPath.toUri());
+    return readFrom(uri);
+  }
 
+  private static ProtoApk readFrom(URI uri) throws IOException {
+    final FileSystem apkFileSystem = FileSystems.newFileSystem(uri, ImmutableMap.of());
     Preconditions.checkArgument(Files.exists(apkFileSystem.getPath(RESOURCE_TABLE)));
     Preconditions.checkArgument(Files.exists(apkFileSystem.getPath(MANIFEST)));
-    return new ProtoApk(apkFileSystem);
+    return new ProtoApk(URI.create(uri.getSchemeSpecificPart()), apkFileSystem);
   }
 
-  /** Visits all resource declarations and references using the {@link ResourceVisitor}. */
-  public <T extends ResourceVisitor<T>> T visitResources(T sink) throws IOException {
+  /**
+   * Creates a copy of the current apk.
+   *
+   * @param destination Path to the new apk destination.
+   * @param resourceFilter A filter for determining whether a given resource will be included in the
+   *     copy.
+   * @return The new ProtoApk.
+   * @throws IOException when there are issues reading the apk.
+   */
+  public ProtoApk copy(Path destination, BiPredicate<ResourceType, String> resourceFilter)
+      throws IOException {
+
+    final URI dstZipUri = URI.create("jar:" + destination.toUri());
+    try (FileSystem dstZip =
+        FileSystems.newFileSystem(dstZipUri, ImmutableMap.of("create", "true"))) {
+
+      final ResourceTable.Builder dstTableBuilder = ResourceTable.newBuilder();
+      final ResourceTable resourceTable =
+          ResourceTable.parseFrom(Files.newInputStream(apkFileSystem.getPath(RESOURCE_TABLE)));
+      dstTableBuilder.setSourcePool(resourceTable.getSourcePool());
+      for (Package pkg : resourceTable.getPackageList()) {
+        Package dstPkg = copyPackage(resourceFilter, dstZip, pkg);
+        if (!dstPkg.getTypeList().isEmpty()) {
+          dstTableBuilder.addPackage(dstPkg);
+        }
+      }
+      try (OutputStream output =
+          Files.newOutputStream(dstZip.getPath(RESOURCE_TABLE), StandardOpenOption.CREATE_NEW)) {
+        dstTableBuilder.build().writeTo(output);
+      }
+
+      Files.walkFileTree(apkFileSystem.getPath("/"), new CopyingFileVisitor(dstZip));
+    }
+
+    return readFrom(dstZipUri);
+  }
+
+  private Package copyPackage(
+      BiPredicate<ResourceType, String> resourceFilter, FileSystem dstZip, Package pkg)
+      throws IOException {
+    Package.Builder dstPkgBuilder = Package.newBuilder(pkg);
+    dstPkgBuilder.clearType();
+    for (Resources.Type type : pkg.getTypeList()) {
+      copyResourceType(resourceFilter, dstZip, dstPkgBuilder, type);
+    }
+    return dstPkgBuilder.build();
+  }
+
+  private void copyResourceType(
+      BiPredicate<ResourceType, String> resourceFilter,
+      FileSystem dstZip,
+      Package.Builder dstPkgBuilder,
+      Resources.Type type)
+      throws IOException {
+    Type.Builder dstTypeBuilder = Resources.Type.newBuilder(type);
+    dstTypeBuilder.clearEntry();
+
+    ResourceType resourceType = ResourceType.getEnum(type.getName());
+    for (Entry entry : type.getEntryList()) {
+      if (resourceFilter.test(resourceType, entry.getName())) {
+        copyEntry(dstZip, dstTypeBuilder, entry);
+      }
+    }
+    final Resources.Type dstType = dstTypeBuilder.build();
+    if (!dstType.getEntryList().isEmpty()) {
+      dstPkgBuilder.addType(dstType);
+    }
+  }
+
+  private void copyEntry(FileSystem dstZip, Type.Builder dstTypeBuilder, Entry entry)
+      throws IOException {
+    dstTypeBuilder.addEntry(Entry.newBuilder(entry));
+    for (ConfigValue configValue : entry.getConfigValueList()) {
+      if (configValue.hasValue()
+          && configValue.getValue().hasItem()
+          && configValue.getValue().getItem().hasFile()) {
+        final String path = configValue.getValue().getItem().getFile().getPath();
+        final Path resourcePath = dstZip.getPath(path);
+        Files.createDirectories(resourcePath.getParent());
+        Files.copy(apkFileSystem.getPath(path), resourcePath);
+      }
+    }
+  }
+
+  public <T extends ResourceVisitor> T visitResources(T visitor) throws IOException {
 
     // visit manifest
-    visitXmlResource(apkFileSystem.getPath(MANIFEST), sink.enteringManifest());
+    visitXmlResource(apkFileSystem.getPath(MANIFEST), visitor.enteringManifest());
 
     // visit resource table and associated files.
     final ResourceTable resourceTable =
@@ -85,7 +184,7 @@ public class ProtoApk {
             : ImmutableList.of();
 
     for (Package pkg : resourceTable.getPackageList()) {
-      ResourcePackageVisitor pkgVisitor = sink.enteringPackage(pkg.getPackageId().getId());
+      ResourcePackageVisitor pkgVisitor = visitor.enteringPackage(pkg.getPackageId().getId());
       for (Resources.Type type : pkg.getTypeList()) {
         ResourceTypeVisitor typeVisitor =
             pkgVisitor.enteringResourceType(
@@ -101,10 +200,15 @@ public class ProtoApk {
         }
       }
     }
-    return sink;
+    return visitor;
   }
 
-  // TODO(corysmith): Centralize duplicated code with AndroidCompiledDataDeserializer.
+  /** Return the underlying uri for this apk. */
+  public URI asApk() {
+    return uri.normalize();
+  }
+
+  // TODO(72324748): Centralize duplicated code with AndroidCompiledDataDeserializer.
   private static List<String> decodeSourcePool(byte[] bytes) throws UnsupportedEncodingException {
     ByteBuffer byteBuffer = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN);
 
@@ -280,9 +384,9 @@ public class ProtoApk {
 
   private void visitFile(ResourceValueVisitor entryVisitor, FileReference file) {
     final Path path = apkFileSystem.getPath(file.getPath());
-    if (file.getType() == Type.PROTO_XML) {
+    if (file.getType() == FileReference.Type.PROTO_XML) {
       visitXmlResource(path, entryVisitor.entering(path));
-    } else if (file.getType() != Type.PNG) {
+    } else if (file.getType() != FileReference.Type.PNG) {
       entryVisitor.acceptOpaqueFileType(path);
     }
   }
@@ -320,6 +424,11 @@ public class ProtoApk {
     }
   }
 
+  @Override
+  public void close() throws IOException {
+    apkFileSystem.close();
+  }
+
   /** Provides an entry point to recording declared and referenced resources in the apk. */
   public interface ResourceVisitor<T extends ResourceVisitor<T>> {
     /** Called when entering the manifest. */
@@ -372,4 +481,37 @@ public class ProtoApk {
       // pass
     }
   }
+
+  private static class CopyingFileVisitor extends SimpleFileVisitor<Path> {
+
+    private final FileSystem dstZip;
+
+    CopyingFileVisitor(FileSystem dstZip) {
+      this.dstZip = dstZip;
+    }
+
+    @Override
+    public FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs) {
+      // Skip the resources, they are copied above.
+      if (dir.endsWith(RES_DIRECTORY)) {
+        return FileVisitResult.SKIP_SUBTREE;
+      }
+      return FileVisitResult.CONTINUE;
+    }
+
+    @Override
+    @SuppressWarnings("JavaOptionalSuggestions")
+    // Not using Files.copy(Path, Path), as it has been shown to corrupt on certain OSs when copying
+    // between filesystems.
+    public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
+      if (!RESOURCE_TABLE.equals(file.getFileName().toString()) && !Files.isDirectory(file)) {
+        Path dest = dstZip.getPath(file.toString());
+        Files.createDirectories(dest.getParent());
+        try (InputStream in = Files.newInputStream(file)) {
+          Files.copy(in, dest);
+        }
+      }
+      return FileVisitResult.CONTINUE;
+    }
+  }
 }