Add --verify-dexes

Summary: As title. Support running a command over produced dexes.

Reviewed By: beicy

Differential Revision: D48088407

fbshipit-source-id: 4eeb716b63a37aef41bb9bded7a6797591b9f6a0

pyredex/utils.py

@@ -13,12 +13,14 @@
 import glob
 import json
 import logging
+import multiprocessing
 import os
 import re
 import shutil
 import subprocess
 import sys
 import tempfile
+import timeit
 import typing
 import zipfile
 from os.path import basename, dirname, isfile, join
@@ -494,3 +496,41 @@ def ensure_libs_dir(libs_dir: str, sub_dir: str) -> str:
 
 def get_file_ext(file_name: str) -> str:
     return os.path.splitext(file_name)[1]
+
+
+def _verify_dex(dex_file: str, cmd: str) -> bool:
+    logging.debug("Verifying %s...", dex_file)
+
+    res = subprocess.run(
+        f"{cmd} '{dex_file}'", shell=True, text=True, capture_output=True
+    )
+    if res.returncode == 0:
+        return True
+
+    logging.error(
+        "Failed verification for %s:\n%s\n---\n%s", dex_file, res.stdout, res.stderr
+    )
+    return False
+
+
+def verify_dexes(dex_dir: str, cmd: str) -> None:
+    timer = timeit.default_timer
+    start = timer()
+
+    dex_files = glob.glob(join(join(dex_dir, "**"), "*.dex"), recursive=True)
+    if not dex_files:
+        logging.warning("Found no dex files to verify")
+        return
+
+    logging.info("Verifying %d dex files...", len(dex_files))
+
+    with multiprocessing.Pool() as pool:
+        result = pool.starmap(
+            _verify_dex,
+            ((dex_file, cmd) for dex_file in dex_files),
+        )
+
+    assert all(result), "Some dex files failed to verify!"
+
+    end = timer()
+    logging.debug("Dex verification finished in {:.2f} seconds".format(end - start))

redex.py

@@ -47,6 +47,7 @@
     relocate_dexen_to_directories,
     remove_comments,
     sign_apk,
+    verify_dexes,
     with_temp_cleanup,
 )
 
@@ -785,6 +786,10 @@ def arg_parser(
         help="Path to JNI summary directory of json files.",
     )
 
+    parser.add_argument(
+        "--verify-dexes", type=str, help="Verify dex files with the supplied command"
+    )
+
     # Manual tool paths.
 
     # Must be subclassed.
@@ -1296,6 +1301,9 @@ def get_compression_list() -> typing.List[CompressionEntry]:
 
 
 def finalize_redex(state: State) -> None:
+    if state.args.verify_dexes:
+        verify_dexes(state.dex_dir, state.args.verify_dexes)
+
     if state.dexen_initial_state is not None:
         dexen_final_state = DexenSnapshot(state.dex_dir)
         assert _assert_val(state.dexen_initial_state).equals(