diff --git a/uraniborg/docs/automate_observation.md b/uraniborg/docs/automate_observation.md index 2c62914..228ee5f 100644 --- a/uraniborg/docs/automate_observation.md +++ b/uraniborg/docs/automate_observation.md @@ -42,6 +42,242 @@ A requested serial that is not connected is not silently skipped: it is logged as an error, reported as `FAILED` in the final summary, and makes the script exit with code `1`. The remaining requested devices are still observed. +### Ignored Flag Combinations + +Some flags only matter together with others. When a flag would be ignored, the +script logs a warning at startup and then runs as it otherwise would; the exit +code is not affected: + +- `--check_preinstalled_only`, `--verifier_path`, `--no_prefetch`, + `--cache_dir`, `--cache_prefetch_concurrency` and `--cache_prefetch_timeout` + have no effect without `--perform_inclusion_proof_check`. The two integer + flags only count as given when their value differs from the default. +- `--cache_prefetch_concurrency` and `--cache_prefetch_timeout` also have no + effect with `--no_prefetch`, since they only tune pre-fetching. + (`--cache_dir` is still used: verification reads the cache.) +- `--pull-all-apks` together with `--pull-preinstalled-apks-only` pulls only + pre-installed APKs. + +`--perform_inclusion_proof_check` without `--verifier_path` remains an error. + +### Machine-Readable Progress (`--events`) + +Log messages are meant for people and may change between versions. Wrappers +such as CI jobs or device-lab scripts should use `--events` instead: it writes +a stream of progress events as [JSON Lines](https://jsonlines.org/), one JSON +object per line, flushed as each event happens. + +```bash +# Write events to a file (overwritten if it exists): +python3 automate_observation.py --events run.jsonl + +# Or stream them on stdout: +python3 automate_observation.py --events - | my-wrapper +``` + +With `--events -`, stdout carries **only** events. Everything else that would +otherwise have gone to stdout is sent to stderr, including output from child +processes such as the verifier and the Xiaomi manual-install prompt. Logging +always goes to stderr. Without `--events`, nothing changes: log output, exit +codes and the results layout are the same. If the `--events` destination +cannot be opened, the script exits with code `1` before doing anything else. +If writing events fails partway through (for example, the reading process +exits), a warning is logged and the run continues without events. + +#### Event schema (version 1) + +Every event has these fields: + +| Field | Meaning | +| :--- | :--- | +| `v` | Schema version, currently `1` | +| `ts` | UTC timestamp, ISO 8601 with milliseconds, e.g. `2026-01-02T03:04:05.678Z` | +| `type` | One of the event types below | + +Optional fields are left out rather than set to `null`. + +| `type` | Fields | When | +| :--- | :--- | :--- | +| `run_started` | `argv` (arguments, without the script name), `pid` | First event of every run | +| `step` | `step`, `state` (`started` / `finished` / `failed`), `device`\*, `duration_ms`\*\*, `message`\* | Around each phase (see below) | +| `step_progress` | `step`, `device`\*, `done`, `total` | Between a step's `started` and its `finished` / `failed`, for long steps (see [Step progress](#step-progress)) | +| `devices` | `devices` (list of `{serial, unauthorized, model?, product?, device?}`), `selected` (serials to observe), `missing` (requested via `--serial` but not connected) | Once, after listing devices | +| `device_started` | `device` | Before processing each selected device | +| `prompt` | `device`, `kind`, `message`, `expects_input` | The script is waiting for a person (see [Manual intervention](#manual-intervention)) | +| `prompt_resolved` | `device`, `kind`, `outcome` | The wait is over, however it ended | +| `device_finished` | `device`, `status`, `results_dir`\*, `error`\* | After each selected device, and once for each missing serial | +| `run_finished` | `exit_code`, `ok`, `summary`, `error`\* | Last event of a run that ends normally, exits early, raises, or is stopped by Ctrl-C or `SIGTERM` (see [End of stream](#end-of-stream)) | + +\* only when relevant. \*\* only on `finished` / `failed`. + +Both `error` fields are `{reason, message}`: `reason` is a stable code (see +[Error reasons](#error-reasons)) to match on, and `message` is human-readable +text that may change. Step `message` is free text. + +**Steps.** Run-level steps have no `device` field and run in this order: +`build_hubble` (only without `-H`), `verify_hubble`, `check_adb`, +`start_adb_server`, `list_devices`. Per-device steps carry `device`: +`uninstall_previous` (only if Hubble was already installed), `install_hubble`, +`launch_hubble`, `wait_for_results`, `extract_results`, `extract_selinux`, and, +with `--perform_inclusion_proof_check`, `inclusion_proof_prefetch` (only when +pre-fetching is attempted) and `inclusion_proof_check`. A failed +`inclusion_proof_prefetch` is not fatal: verification falls back to fetching +entries on demand. + +#### Step progress + +Some steps take minutes. While they run, `step_progress` events report how far +they have got, e.g. to show "120 of 314 checked". Currently only +`inclusion_proof_check` reports progress: `done` is the number of APK splits +verified so far (found in the log or not) and `total` the number of splits to +verify (splits skipped for missing fields are not counted). + +- Each event carries the same `step` and `device` as the step it belongs to, + and appears only between that step's `started` and `finished` / `failed`. +- The first event has `done` = `0`, once verification starts. For + `inclusion_proof_check`, that is after the package list has been read; if it + cannot be read, the step fails without any `step_progress`. +- `done` only increases, and `total` stays the same within a step. +- Events are throttled: at most one every 2 seconds, plus the first one and the + one where `done` reaches `total`, which is always sent. A step with nothing + to verify sends a single event with `done` = `total` = `0`. +- A step that is stopped early (Ctrl-C, `SIGTERM`, an error) ends with + `failed` without reaching `done` = `total`. + +```json +{"v": 1, "ts": "2026-01-02T03:05:00.000Z", "type": "step", "step": "inclusion_proof_check", "state": "started", "device": "ABCDEF012345"} +{"v": 1, "ts": "2026-01-02T03:05:00.010Z", "type": "step_progress", "step": "inclusion_proof_check", "device": "ABCDEF012345", "done": 0, "total": 314} +{"v": 1, "ts": "2026-01-02T03:05:02.020Z", "type": "step_progress", "step": "inclusion_proof_check", "device": "ABCDEF012345", "done": 11, "total": 314} +... +{"v": 1, "ts": "2026-01-02T03:05:55.400Z", "type": "step_progress", "step": "inclusion_proof_check", "device": "ABCDEF012345", "done": 314, "total": 314} +{"v": 1, "ts": "2026-01-02T03:05:55.410Z", "type": "step", "step": "inclusion_proof_check", "state": "finished", "device": "ABCDEF012345", "duration_ms": 55410} +``` + +#### Manual intervention + +Two situations need a person. Each is reported as a `prompt` event, and is +always followed by a matching `prompt_resolved` (same `device` and `kind`) +once the wait ends, whether it succeeded, failed or was interrupted. + +| `kind` | When | `expects_input` | What to do | +| :--- | :--- | :--- | :--- | +| `xiaomi_manual_install` | During `install_hubble` on Xiaomi phones, where Hubble must be installed by hand | `true` | Install Hubble on the device, then write a newline to the script's stdin. If Hubble is still not installed, the script waits for another newline, without a new `prompt` | +| `adb_backup_confirm` | During `extract_results`, when results must be pulled with `adb backup` | `false` | Tap `Back up my data` on the device. Nothing is read from stdin | + +`prompt_resolved.outcome` says how the wait ended: + +| `outcome` | Meaning | +| :--- | :--- | +| `done` | Hubble is installed, or `adb backup` completed | +| `failed` | `adb backup` failed (the device then fails with `extract_failed`) | +| `stdin_closed` | stdin was closed while waiting for a newline | +| `interrupted` / `terminated` / `unexpected_error` | The wait was cut short by Ctrl-C, `SIGTERM` or an unhandled exception | + +If stdin is closed while the script waits for a newline (for example, the +wrapper exits, or stdin is `/dev/null`), that device fails with reason +`stdin_closed` and the run moves on to the next device. + +**Device status.** `device_finished.status` and `summary[serial].status` use the +same four outcomes as the final log summary: + +| `status` | Meaning | +| :--- | :--- | +| `success` | Results collected, and the inclusion proof check (if requested) completed | +| `partial_check_incomplete` | Results collected, but the inclusion proof check could not complete (for example, the package list was missing or unreadable, or the output could not be written) | +| `partial_error` | Results collected, but an unexpected error occurred afterwards, or the run was stopped (`error.reason` is `interrupted` or `terminated`) | +| `failed` | No results collected (including unauthorized and not-connected devices, and devices stopped before results were collected) | + +> [!IMPORTANT] +> `success` does **not** mean every APK was found in the transparency log. A +> completed check records a per-split `inclusion_proof_verified` result, which +> may be `false`, in `packages_with_inclusion_proof_signal.txt` (or +> `preinstalled_packages_with_inclusion_proof_signal.txt`) inside +> `results_dir`. Read that file to decide whether the APKs are in the log. + +`results_dir` is present whenever results were collected, including both +`partial_*` outcomes. + +**Run outcome.** `run_finished.summary` maps each serial, in the order it was +requested (or discovered, without `--serial`), to `{status, results_dir?}`. +If the run is stopped early (`interrupted`, `terminated` or +`unexpected_error`), `summary` lists only the devices that reached +`device_finished`; devices that had not been processed yet are absent. +`ok` is `true` only when `exit_code` is `0` **and** there is no `error`. Treat +`ok`, not `exit_code`, as the success signal: when the script stops before +processing any device (for example, no device is connected), it still exits +with code `0`, but `run_finished.error` explains why. + +#### Error reasons + +`run_finished.error.reason`: + +| `reason` | Meaning | +| :--- | :--- | +| `unsupported_platform` | The host OS is not supported | +| `android_sdk_not_found` | Rebuilding Hubble (no `-H`) failed: Android SDK not found | +| `hubble_build_failed` | Rebuilding Hubble failed (Gradle or symlink error) | +| `invalid_hubble_apk` | The `-H` path is not an `.apk` | +| `adb_not_found` | `adb` is not installed | +| `adb_server_failed` | The adb server could not be started | +| `adb_devices_failed` | `adb devices` failed (without `--serial`) | +| `no_devices` | No device is connected (without `--serial`) | +| `unexpected_error` | An unhandled exception; `exit_code` is `1` | +| `interrupted` | Interrupted with Ctrl-C; `exit_code` is `130` | +| `terminated` | Stopped by `SIGTERM`; `exit_code` is `143` | + +`device_finished.error.reason`: + +| `reason` | Meaning | `status` | +| :--- | :--- | :--- | +| `not_connected` | Requested via `--serial` but not connected | `failed` | +| `unauthorized` | ADB is not authorized on the device | `failed` | +| `uninstall_failed` | A previous Hubble installation could not be removed | `failed` | +| `install_failed` | Hubble could not be installed | `failed` | +| `launch_failed` | Hubble could not be launched | `failed` | +| `no_results` | Hubble produced no results in time | `failed` | +| `extract_failed` | Results could not be pulled from the device | `failed` | +| `stdin_closed` | stdin was closed while waiting for a manual Hubble install | `failed` | +| `inclusion_proof_check_incomplete` | The inclusion proof check could not complete | `partial_check_incomplete` | +| `unexpected_error` | An unhandled exception while processing the device | `failed` or `partial_error` | +| `interrupted` | Ctrl-C while processing the device | `failed` or `partial_error` | +| `terminated` | `SIGTERM` while processing the device | `failed` or `partial_error` | + +For the last three, `status` is `partial_error` if results had already been +collected, otherwise `failed`. + +Example, for one device with inclusion proofs: + +```json +{"v": 1, "ts": "2026-01-02T03:04:05.678Z", "type": "run_started", "argv": ["--events", "-", "--perform_inclusion_proof_check", "--verifier_path=/path/to/verifier"], "pid": 4242} +{"v": 1, "ts": "2026-01-02T03:04:05.690Z", "type": "step", "step": "verify_hubble", "state": "started"} +{"v": 1, "ts": "2026-01-02T03:04:05.691Z", "type": "step", "step": "verify_hubble", "state": "finished", "duration_ms": 1} +... +{"v": 1, "ts": "2026-01-02T03:04:06.100Z", "type": "devices", "devices": [{"serial": "ABCDEF012345", "unauthorized": false, "model": "Pixel_9"}], "selected": ["ABCDEF012345"], "missing": []} +{"v": 1, "ts": "2026-01-02T03:04:06.101Z", "type": "device_started", "device": "ABCDEF012345"} +{"v": 1, "ts": "2026-01-02T03:04:06.102Z", "type": "step", "step": "install_hubble", "state": "started", "device": "ABCDEF012345"} +... +{"v": 1, "ts": "2026-01-02T03:09:41.310Z", "type": "device_finished", "device": "ABCDEF012345", "status": "success", "results_dir": "/path/to/results/google/.../000"} +{"v": 1, "ts": "2026-01-02T03:09:41.312Z", "type": "run_finished", "exit_code": 0, "ok": true, "summary": {"ABCDEF012345": {"status": "success", "results_dir": "/path/to/results/google/.../000"}}} +``` + +#### End of stream + +`run_finished` cannot be written if the process is killed with `SIGKILL`, the +Python interpreter crashes, or the host goes down. **A stream that ends without +`run_finished` means the run failed.** Do not infer success from the +`device_finished` events seen so far. + +With `--events`, `SIGTERM` (how CI systems usually enforce timeouts and +cancellations) is caught so that `run_finished` can be written. The script then +still terminates by `SIGTERM`, so the parent process sees the same exit status +as before. Without `--events`, `SIGTERM` handling is unchanged. + +#### Compatibility + +The schema is a stable contract. New event types and new fields may be added +without changing `v`, so consumers should ignore types and fields they do not +recognize. Removing a field or changing what it means requires a new `v`. + ## Performing Inclusion Proof Checks You can automatically verify extracted package APK splits against Android Binary @@ -88,6 +324,30 @@ flags: `preinstalled_packages.txt` instead of `packages.txt`, writing results to `preinstalled_packages_with_inclusion_proof_signal.txt`. +### Batch Verification +All APK splits are verified by a single `verifier` run in batch mode +(`--payloads_path`, see the [verifier README](../../verifier_tools/verify/README.md)). +It fetches each log's checkpoint, and searches the log, once for all splits. +This keeps a cold cache (e.g. with `--no_prefetch`, or after pre-fetching +failed) from being filled by many verifiers downloading the same files at +once. Results are reported as the verifier finds them. + +With a `verifier` built before batch mode existed, or for any splits a batch +run gives no result for (e.g. if it crashes), each split is verified by a +`verifier` run of its own instead, one at a time. That is much slower, so a +warning suggests rebuilding the `verifier`. The results file is the same +either way, including the order of packages and splits. + +Log lines about individual splits (shown with `-D`) start with the package and +split they are about, e.g. `com.android.chrome [config.en]: ...`. + +Ctrl-C stops the running verifier and starts no new ones. `SIGTERM` does the +same when `--events` is given (and always for `inclusion_proof_check.py`). +Without `--events`, `automate_observation.py` keeps the default `SIGTERM` +behaviour: the script exits at once, and a verifier that was already running +stops on its own: a per-split verifier after its split (about a second or +two), a batch verifier when it next tries to write a result. + ### Pulling Pre-installed APKs Only By default, `--pull-all-apks` downloads all packages listed in `packages.txt`. You can pass `--pull-preinstalled-apks-only` to download only pre-installed diff --git a/uraniborg/scripts/python/automate_observation.py b/uraniborg/scripts/python/automate_observation.py index 0de14f9..7dc9fda 100644 --- a/uraniborg/scripts/python/automate_observation.py +++ b/uraniborg/scripts/python/automate_observation.py @@ -21,19 +21,25 @@ Hubble data back to host. """ import argparse +import contextlib +import dataclasses +import datetime import io # to convert regular buffer to in-memory bytes buffer for tarfileobj import json import logging import os import shutil # to help move files +import signal import sys import tarfile # to untar decompressed adb backup file import tempfile -from typing import Optional # runtime support for type hints. +import time +from typing import Optional, TextIO # runtime support for type hints. import zlib # to decompress adb backup import inclusion_proof_check import syscall_wrapper +import termination AdbWrapper = syscall_wrapper.AdbWrapper SyscallWrapper = syscall_wrapper.SyscallWrapper @@ -76,6 +82,15 @@ def parse_arguments() -> argparse.Namespace: "the order given. If omitted, every connected " "device is observed. A requested serial that is " "not connected is reported as FAILED.") + parser.add_argument("--events", required=False, default=None, + metavar="PATH", + help="If specified, writes machine-readable progress " + "events as JSON Lines (one JSON object per line) " + "to PATH, overwriting it. Use \"-\" to write events " + "to stdout; everything else the script would print " + "to stdout is then sent to stderr, so stdout " + "carries only events. See " + "docs/automate_observation.md for the schema.") parser.add_argument("--pull-all-apks", required=False, action="count", help="If specified, the script will attempt to download " @@ -132,6 +147,54 @@ def parse_arguments() -> argparse.Namespace: return args +def validate_argument_combinations(args: argparse.Namespace) -> list[str]: + """Finds flags that are ignored because of how they are combined. + + These combinations are not errors, so that existing invocations keep + working; run() logs each returned message as a warning. Combinations that + cannot work at all are rejected by parse_arguments() instead. + + Args: + args: Parsed arguments from parse_arguments(). + + Returns: + One human-readable warning per ignored flag or conflict, in a stable + order. Empty if nothing is ignored. + """ + warnings = [] + # Integer flags always have a value, so only a non-default one counts as + # given. Only prefetch_log_entries() reads them. + prefetch_tuning = [ + ("--cache_prefetch_concurrency", + args.cache_prefetch_concurrency != + inclusion_proof_check.DEFAULT_PREFETCH_CONCURRENCY), + ("--cache_prefetch_timeout", + args.cache_prefetch_timeout != + inclusion_proof_check.DEFAULT_PREFETCH_TIMEOUT), + ] + if not args.perform_inclusion_proof_check: + ignored = [ + ("--check_preinstalled_only", args.check_preinstalled_only), + ("--verifier_path", args.verifier_path is not None), + ("--no_prefetch", args.no_prefetch), + ("--cache_dir", args.cache_dir is not None), + ] + prefetch_tuning + for flag, given in ignored: + if given: + warnings.append("{} has no effect without " + "--perform_inclusion_proof_check.".format(flag)) + elif args.no_prefetch: + # Pre-fetching is skipped, so its tuning flags are unused. --cache_dir is + # not: verification itself still uses the cache. + for flag, given in prefetch_tuning: + if given: + warnings.append("{} has no effect with --no_prefetch.".format(flag)) + if args.pull_all_apks is not None and args.pull_preinstalled_apks_only: + warnings.append("--pull-all-apks and --pull-preinstalled-apks-only were " + "both given; only pre-installed APKs will be pulled.") + return warnings + + def set_up_logging(args: argparse.Namespace) -> logging.Logger: """Sets up various logging parameters. @@ -157,6 +220,335 @@ def set_up_logging(args: argparse.Namespace) -> logging.Logger: return logger +# --- Machine-readable progress events (--events) ------------------------------ +# +# The schema is documented in docs/automate_observation.md and is a stable +# contract: fields and event types may be added without bumping +# EVENTS_SCHEMA_VERSION, but never removed or changed in meaning. + +EVENTS_SCHEMA_VERSION = 1 + +# step_progress events are written at most this often (plus the first and +# the final one of each step). +STEP_PROGRESS_INTERVAL_SECONDS = 2.0 + +# Per-device outcomes. These are the same four outcomes the final log summary +# distinguishes. +STATUS_SUCCESS = "success" +STATUS_PARTIAL_CHECK_INCOMPLETE = "partial_check_incomplete" +STATUS_PARTIAL_ERROR = "partial_error" +STATUS_FAILED = "failed" + +# Error reasons. `error` fields in run_finished and device_finished are +# {reason, message}: reason is one of these stable codes, message is free text. +# +# run_finished.error, when the run stops before processing any device: +REASON_UNSUPPORTED_PLATFORM = "unsupported_platform" +REASON_ANDROID_SDK_NOT_FOUND = "android_sdk_not_found" +REASON_HUBBLE_BUILD_FAILED = "hubble_build_failed" +REASON_INVALID_HUBBLE_APK = "invalid_hubble_apk" +REASON_ADB_NOT_FOUND = "adb_not_found" +REASON_ADB_SERVER_FAILED = "adb_server_failed" +REASON_ADB_DEVICES_FAILED = "adb_devices_failed" +REASON_NO_DEVICES = "no_devices" +# device_finished.error only: +REASON_NOT_CONNECTED = "not_connected" +REASON_UNAUTHORIZED = "unauthorized" +REASON_UNINSTALL_FAILED = "uninstall_failed" +REASON_INSTALL_FAILED = "install_failed" +REASON_LAUNCH_FAILED = "launch_failed" +REASON_NO_RESULTS = "no_results" +REASON_EXTRACT_FAILED = "extract_failed" +REASON_INCLUSION_PROOF_CHECK_INCOMPLETE = "inclusion_proof_check_incomplete" +REASON_STDIN_CLOSED = "stdin_closed" +# Both run_finished.error and device_finished.error: +REASON_UNEXPECTED_ERROR = "unexpected_error" +REASON_INTERRUPTED = "interrupted" +REASON_TERMINATED = "terminated" + +# Kinds of manual intervention reported by prompt / prompt_resolved events. +PROMPT_XIAOMI_MANUAL_INSTALL = "xiaomi_manual_install" +PROMPT_ADB_BACKUP_CONFIRM = "adb_backup_confirm" + +# How a prompt ended, reported by prompt_resolved.outcome. An exception during +# the wait maps to REASON_INTERRUPTED, REASON_TERMINATED or +# REASON_UNEXPECTED_ERROR. +PROMPT_OUTCOME_DONE = "done" +PROMPT_OUTCOME_FAILED = "failed" +PROMPT_OUTCOME_STDIN_CLOSED = REASON_STDIN_CLOSED + + +def _interruption_error(e: BaseException) -> dict: + """Describes an exception that is not an Exception, for device errors.""" + if isinstance(e, KeyboardInterrupt): + return _error(REASON_INTERRUPTED, "Interrupted.") + if isinstance(e, termination.Terminated): + return _error(REASON_TERMINATED, "Terminated.") + return _error(REASON_UNEXPECTED_ERROR, + "{}: {}".format(type(e).__name__, e)) + + +def open_event_stream(path: str) -> TextIO: + """Opens the destination for --events. + + For "-", events go to the process's original stdout. To guarantee that + stdout then carries nothing but events, file descriptor 1 is redirected to + stderr afterwards, so anything else written to stdout -- by this script, + input() prompts, or child processes that inherit stdout -- lands on stderr. + + Args: + path: A file path, or "-" for stdout. + + Returns: + A writable text stream. + + Raises: + OSError: if the destination cannot be opened. + """ + if path == "-": + sys.stdout.flush() + events_fd = os.dup(sys.stdout.fileno()) + os.dup2(sys.stderr.fileno(), sys.stdout.fileno()) + return os.fdopen(events_fd, "w", encoding="utf-8") + return open(path, "w", encoding="utf-8") + + +class _Step: + """Handle yielded by EventEmitter.step() to mark a step as failed.""" + + def __init__(self): + self.failed = False + self.message = None + + def fail(self, message: str) -> str: + """Marks the step failed and returns message, for convenient reuse.""" + self.failed = True + self.message = message + return message + + +class _Prompt: + """Handle yielded by EventEmitter.prompt() to record how the wait ended.""" + + def __init__(self): + self.outcome = PROMPT_OUTCOME_DONE + + +class EventEmitter: + """Writes progress events as JSON Lines. + + Without a stream every method is a no-op, so callers never need to check + whether --events was given. Writing is best-effort: if the destination + breaks (e.g. the reading process went away), a warning is logged once and + the run carries on without events. + """ + + def __init__(self, stream: Optional[TextIO] = None, + logger: Optional[logging.Logger] = None, + clock=None, monotonic=None): + self._stream = stream + self._logger = logger + self._clock = clock or ( + lambda: datetime.datetime.now(datetime.timezone.utc)) + # Used only to throttle step_progress events. + self._monotonic = monotonic or time.monotonic + self._run_finished = False + # serial -> {status, results_dir?}, in device_finished order. + self._finished_devices = {} + + @property + def enabled(self) -> bool: + return self._stream is not None + + def emit(self, event_type: str, **fields): + """Writes one event. Fields whose value is None are omitted.""" + if self._stream is None: + return + timestamp = self._clock().isoformat(timespec="milliseconds") + event = {"v": EVENTS_SCHEMA_VERSION, + "ts": timestamp.replace("+00:00", "Z"), + "type": event_type} + event.update({k: v for k, v in fields.items() if v is not None}) + try: + self._stream.write(json.dumps(event, default=str) + "\n") + self._stream.flush() + except (OSError, ValueError) as e: + # ValueError: write to a closed file. + if self._logger: + self._logger.warning("Failed to write --events output (%s); " + "continuing without events.", e) + self._stream = None + + def device_finished(self, device: str, status: str, + results_dir: Optional[str] = None, + error: Optional[dict] = None): + """Emits device_finished and remembers the outcome for finish_run().""" + entry = {"status": status} + if results_dir is not None: + entry["results_dir"] = results_dir + self._finished_devices[device] = entry + self.emit("device_finished", device=device, status=status, + results_dir=results_dir, error=error) + + def progress_reporter( + self, step: str, device: Optional[str] = None, + interval: float = STEP_PROGRESS_INTERVAL_SECONDS): + """Returns a progress(done, total) callback emitting step_progress. + + Meant to be called often (e.g. once per item); events are throttled: + the first call is always reported, and so is the one where done reaches + total, but in between at most one event is written per `interval` + seconds. A call repeating the last reported `done` is never reported. + """ + last = {"done": None, "at": None} + + def report(done: int, total: int) -> None: + if self._stream is None or done == last["done"]: + return + now = self._monotonic() + if (last["at"] is not None and done != total and + now - last["at"] < interval): + return + last["done"], last["at"] = done, now + self.emit("step_progress", step=step, device=device, done=done, + total=total) + + return report + + @contextlib.contextmanager + def step(self, name: str, device: Optional[str] = None): + """Brackets a phase with step started/finished/failed events. + + The step is reported as failed if the body raises (the exception is + re-raised) or calls fail() on the yielded handle; otherwise as finished. + """ + handle = _Step() + start = time.monotonic() + self.emit("step", step=name, state="started", device=device) + try: + yield handle + except BaseException as e: + self.emit("step", step=name, state="failed", device=device, + duration_ms=int((time.monotonic() - start) * 1000), + message="{}: {}".format(type(e).__name__, e)) + raise + self.emit("step", step=name, + state="failed" if handle.failed else "finished", + device=device, + duration_ms=int((time.monotonic() - start) * 1000), + message=handle.message) + + @contextlib.contextmanager + def prompt(self, device: Optional[str], kind: str, message: str, + expects_input: bool): + """Brackets a wait for manual intervention with prompt/prompt_resolved. + + prompt_resolved is emitted however the body exits, including by an + exception, so every prompt is closed. Its outcome is "done" unless the + body sets another PROMPT_OUTCOME_* on the yielded handle, or raises (then + it is the matching error reason, e.g. "interrupted"). expects_input tells + the reader whether the script is blocked on stdin (write a newline to + continue) or on an action on the device. + """ + handle = _Prompt() + self.emit("prompt", device=device, kind=kind, message=message, + expects_input=expects_input) + try: + yield handle + except BaseException as e: + handle.outcome = _interruption_error(e)["reason"] + raise + finally: + self.emit("prompt_resolved", device=device, kind=kind, + outcome=handle.outcome) + + def finish_run(self, exit_code: int, summary: Optional[dict] = None, + error: Optional[dict] = None): + """Emits run_finished. Only the first call has any effect. + + If summary is None, it defaults to the devices reported through + device_finished() so far, so a run that is cut short still reports the + devices that did finish. + """ + if self._run_finished: + return + self._run_finished = True + self.emit("run_finished", + exit_code=exit_code, + ok=(exit_code == 0 and error is None), + summary=(summary if summary is not None + else dict(self._finished_devices)), + error=error) + + def close(self): + if self._stream is not None: + try: + self._stream.close() + except OSError: + pass + self._stream = None + + +def _error(reason: str, message: str) -> dict: + """Builds an `error` field: a stable REASON_* code plus free text.""" + return {"reason": reason, "message": message} + + +def _device_to_event(device) -> dict: + """Describes a DeviceInfo for the devices event.""" + info = {"serial": device.serial_number, + "unauthorized": bool(device.unauthorized)} + for key, attr in (("model", "model_name"), ("product", "product_name"), + ("device", "device_name")): + value = getattr(device, attr, None) + if isinstance(value, str) and value: + info[key] = value + return info + + +def classify_device(collected: bool, collection_error: bool, + check_incomplete: bool) -> str: + """Maps what happened on one device to one of the STATUS_* values.""" + if not collected: + return STATUS_FAILED + if collection_error: + return STATUS_PARTIAL_ERROR + if check_incomplete: + return STATUS_PARTIAL_CHECK_INCOMPLETE + return STATUS_SUCCESS + + +@dataclasses.dataclass +class DeviceResult: + """The outcome of observing one device, as reported by device_finished.""" + status: str + results_dir: Optional[str] = None + error: Optional[dict] = None + + +@dataclasses.dataclass +class RunState: + """State carried from one device to the next within a run.""" + # Log entries only need pre-fetching once per run. + prefetched: bool = False + + +@dataclasses.dataclass +class _DeviceProgress: + """What process_device() has learned about one device so far.""" + results_dir: Optional[str] = None # set once results are fully collected + error: Optional[dict] = None + collection_error: bool = False + check_incomplete: bool = False + + def result(self) -> DeviceResult: + return DeviceResult( + status=classify_device(self.results_dir is not None, + self.collection_error, self.check_incomplete), + results_dir=self.results_dir, + error=self.error) + + def supported_platform(logger: logging.Logger) -> bool: """Checks if this script is running on supported platform. @@ -321,6 +713,48 @@ def launch_xiaomi_file_explorer( "com.android.fileexplorer.FileExplorerTabActivity") +def wait_for_xiaomi_manual_install(adb_wrapper: syscall_wrapper.AdbWrapper, + serial: str, + logger: logging.Logger, + events: EventEmitter) -> bool: + """Waits, via stdin, for the user to install Hubble by hand on a Xiaomi phone. + + The user is asked to press Enter after installing; this repeats until Hubble + is installed. With --events, the wait is reported as a prompt event so that + a parent process can ask its user and then write a newline to stdin. + + Args: + adb_wrapper: An AdbWrapper for the target device. + serial: The target device's serial number, for events. + logger: A logger object to log messages. + events: The EventEmitter for the run. + + Returns: + True once Hubble is installed; False if stdin was closed first (e.g. the + parent process cancelled the run, or stdin is /dev/null). + """ + if is_hubble_installed(adb_wrapper, logger): + return True + with events.prompt(serial, PROMPT_XIAOMI_MANUAL_INSTALL, + "Install Hubble manually from the \"Downloads\" folder " + "in the \"Files Manager\" app on the device, then press " + "Enter.", + expects_input=True) as prompt: + while True: + logger.warning("Please manually install Hubble by launching the " + "\"Files Manager\" app (it may have been launched " + "for you) and navigate to the \"Downloads\" folder.") + try: + input("Press [ENTER] when you are done.") + except EOFError: + logger.error("Standard input was closed while waiting for Hubble to " + "be installed manually on device %s.", serial) + prompt.outcome = PROMPT_OUTCOME_STDIN_CLOSED + return False + if is_hubble_installed(adb_wrapper, logger): + return True + + def adb_push_hubble(adb_wrapper: syscall_wrapper.AdbWrapper, hubble_path: str): """Drops the Hubble APK onto device (used when direct installation fails). @@ -594,7 +1028,8 @@ def classify_dir_using_build_fingerprint( extract_apks: bool, logger: logging.Logger, pull_preinstalled_only: bool = False, - tmp_dir: str = "/tmp") -> Optional[str]: + tmp_dir: str = "/tmp", + events: Optional[EventEmitter] = None) -> Optional[str]: """Decides which directory in results/ to dump new result to. This is a renewed method that makes use of build fingerprint to do @@ -612,6 +1047,8 @@ def classify_dir_using_build_fingerprint( listed in preinstalled_packages.txt. tmp_dir: Temporary directory on host used for staging build.txt and adb backup artifacts. Defaults to "/tmp". + events: An optional EventEmitter; used to report the `adb backup` + confirmation that the user must give on the device. Returns: A string representing the final directory (on host) where results are pulled @@ -632,7 +1069,15 @@ def classify_dir_using_build_fingerprint( decompressed_backup_filepath = os.path.join(tmp_dir, "hubble_results.tar") logger.warning("Manual intervention required: Please select " "`Back up my data` to proceed") - if not adb_wrapper.backup(compressed_backup_filepath, HUBBLE_PACKAGE_NAME): + with (events or EventEmitter()).prompt( + adb_wrapper.device_serial_number, PROMPT_ADB_BACKUP_CONFIRM, + "Select `Back up my data` on the device to proceed.", + expects_input=False) as prompt: + backed_up = adb_wrapper.backup(compressed_backup_filepath, + HUBBLE_PACKAGE_NAME) + if not backed_up: + prompt.outcome = PROMPT_OUTCOME_FAILED + if not backed_up: logger.error("Failed to use `adb backup` to pull result files.") return None @@ -738,7 +1183,9 @@ def extract_results_and_apks(adb_wrapper: syscall_wrapper.AdbWrapper, logger: logging.Logger, extract_apks=False, pull_preinstalled_only=False, - tmp_dir: str = "/tmp") -> Optional[str]: + tmp_dir: str = "/tmp", + events: Optional[EventEmitter] = None + ) -> Optional[str]: """Extracts results (and optionally APKs) from Hubble's execution. Args: @@ -752,6 +1199,7 @@ def extract_results_and_apks(adb_wrapper: syscall_wrapper.AdbWrapper, from preinstalled_packages.txt. tmp_dir: Temporary directory on host used for staging build.txt and adb backup artifacts. Defaults to "/tmp". + events: An optional EventEmitter, passed on to report manual intervention. Returns: A string representing the final directory (on host) where results are copied @@ -785,7 +1233,8 @@ def extract_results_and_apks(adb_wrapper: syscall_wrapper.AdbWrapper, extract_apks, logger, pull_preinstalled_only=pull_preinstalled_only, - tmp_dir=tmp_dir) + tmp_dir=tmp_dir, + events=events) def extract_selinux_policies(adb_wrapper: syscall_wrapper.AdbWrapper, @@ -891,215 +1340,289 @@ def ensure_android_sdk(hubble_project_dir: str, logger: logging.Logger) -> bool: return True -def main(): - args = parse_arguments() - logger = set_up_logging(args) +def rebuild_hubble(logger: logging.Logger) -> tuple[Optional[str], Optional[dict]]: + """Rebuilds the Hubble debug APK with Gradle and refreshes prebuilts/APK/latest. - if not supported_platform(logger): - logger.error("Sorry, your OS is currently unsupported for this script.") - return + Args: + logger: A logger object to log debug or error messages. - if args.hubble is None: - logger.info("-H flag not used. Rebuilding Hubble...") - script_dir = os.path.dirname(os.path.abspath(__file__)) - hubble_project_dir = os.path.abspath(os.path.join(script_dir, "../../AndroidStudioProject/Hubble")) - if not ensure_android_sdk(hubble_project_dir, logger): - logger.error("Failed to (re)build Hubble APK: Android SDK not found.") - return - gradlew_path = os.path.join(hubble_project_dir, "gradlew") + Returns: + A tuple (apk_path, error). On success apk_path is the path of the + refreshed "latest" symlink and error is None. On failure apk_path is None + and error is a run_finished error dict (see _error). + """ + script_dir = os.path.dirname(os.path.abspath(__file__)) + hubble_project_dir = os.path.abspath(os.path.join(script_dir, "../../AndroidStudioProject/Hubble")) + if not ensure_android_sdk(hubble_project_dir, logger): + msg = "Failed to (re)build Hubble APK: Android SDK not found." + logger.error(msg) + return None, _error(REASON_ANDROID_SDK_NOT_FOUND, msg) + gradlew_path = os.path.join(hubble_project_dir, "gradlew") - logger.info("Running 'gradlew assembleDebug' in %s", hubble_project_dir) - sw = SyscallWrapper(logger) - sw.call_returnable_command([gradlew_path, "assembleDebug"], cwd=hubble_project_dir) - if sw.error_occured: - logger.error("Failed to (re)build Hubble APK: [%d] %s", sw.return_code, sw.error_message) - return + logger.info("Running 'gradlew assembleDebug' in %s", hubble_project_dir) + sw = SyscallWrapper(logger) + sw.call_returnable_command([gradlew_path, "assembleDebug"], cwd=hubble_project_dir) + if sw.error_occured: + logger.error("Failed to (re)build Hubble APK: [%d] %s", sw.return_code, sw.error_message) + return None, _error( + REASON_HUBBLE_BUILD_FAILED, + "Failed to (re)build Hubble APK: [{}] {}".format( + sw.return_code, sw.error_message)) + + for line in sw.result_final: + logger.debug("Gradle output: %s", line) + + latest_symlink_path = os.path.abspath(os.path.join(script_dir, "../../prebuilts/APK/latest")) + os.makedirs(os.path.dirname(latest_symlink_path), exist_ok=True) + if os.path.exists(latest_symlink_path) or os.path.islink(latest_symlink_path): + logger.debug("Removing old symlink: %s", latest_symlink_path) + try: + os.remove(latest_symlink_path) + except Exception as e: + logger.error("Failed to remove old symlink %s: %s", latest_symlink_path, e) + return None, _error( + REASON_HUBBLE_BUILD_FAILED, + "Failed to remove old symlink {}: {}".format(latest_symlink_path, e)) - for line in sw.result_final: - logger.debug("Gradle output: %s", line) + symlink_target = "../../AndroidStudioProject/Hubble/app/build/outputs/apk/debug/app-debug.apk" - latest_symlink_path = os.path.abspath(os.path.join(script_dir, "../../prebuilts/APK/latest")) - os.makedirs(os.path.dirname(latest_symlink_path), exist_ok=True) - if os.path.exists(latest_symlink_path) or os.path.islink(latest_symlink_path): - logger.debug("Removing old symlink: %s", latest_symlink_path) - try: - os.remove(latest_symlink_path) - except Exception as e: - logger.error("Failed to remove old symlink %s: %s", latest_symlink_path, e) - return + logger.info("Creating new symlink 'latest' -> %s", symlink_target) + try: + os.symlink(symlink_target, latest_symlink_path) + except Exception as e: + logger.error("Failed to create symlink %s: %s", latest_symlink_path, e) + return None, _error( + REASON_HUBBLE_BUILD_FAILED, + "Failed to create symlink {}: {}".format(latest_symlink_path, e)) - symlink_target = "../../AndroidStudioProject/Hubble/app/build/outputs/apk/debug/app-debug.apk" + return latest_symlink_path, None - logger.info("Creating new symlink 'latest' -> %s", symlink_target) - try: - os.symlink(symlink_target, latest_symlink_path) - except Exception as e: - logger.error("Failed to create symlink %s: %s", latest_symlink_path, e) - return - args.hubble = latest_symlink_path +def process_device(target_device, args: argparse.Namespace, + logger: logging.Logger, events: EventEmitter, + state: RunState) -> DeviceResult: + """Observes one selected device, from installing Hubble to the proof check. - if not verify_hubble(args, logger): - return + Emits device_started and, however processing ends, device_finished. + Exceptions are logged and reported as the device's error; KeyboardInterrupt + and other BaseExceptions are reported, then re-raised. - if not adb_installed(logger): - return + Args: + target_device: A DeviceInfo from AdbWrapper.devices(). + args: Parsed arguments from parse_arguments(). + logger: A logger object to log debug or error messages. + events: Where to report progress events. May be a disabled emitter. + state: State shared with the other devices of this run. - if not AdbWrapper.start_server(logger): + Returns: + The device's outcome. + """ + serial = target_device.serial_number + progress = _DeviceProgress() + events.emit("device_started", device=serial) + try: + _observe_device(target_device, args, logger, events, state, progress) + except Exception as e: + logger.exception("Unexpected error while processing device %s: %s", + serial, e) + progress.error = _error(REASON_UNEXPECTED_ERROR, + "{}: {}".format(type(e).__name__, e)) + progress.collection_error = True + except BaseException as e: + # KeyboardInterrupt and other non-Exception exits bypass the handler + # above. Record them before propagating, so the device_finished event + # in the finally block does not report an interrupted device as a + # success when its results were already collected. + progress.error = _interruption_error(e) + progress.collection_error = True + raise + finally: + result = progress.result() + events.device_finished(serial, result.status, + results_dir=result.results_dir, + error=result.error) + return result + + +def _observe_device(target_device, args: argparse.Namespace, + logger: logging.Logger, events: EventEmitter, + state: RunState, progress: _DeviceProgress): + """The steps of process_device(); returns early when a step fails. + + Failures are recorded on progress rather than returned, so that + process_device() can report them even when an exception cuts this short. + """ + serial = target_device.serial_number + if target_device.unauthorized: + logger.error("Please authorize device with serial number %s for ADB via " + "device GUI.", serial) + progress.error = _error(REASON_UNAUTHORIZED, + "ADB is not authorized on this device.") return - connected_devices = AdbWrapper.devices(logger) or [] - # With --serial, fall through: every requested serial is then reported as - # missing, FAILED in the summary, and the run exits 1. - if not connected_devices and args.serial is None: - logger.error("No devices connected!") - return + # set up an adb_wrapper to be used throughout for this target device + adb_wrapper = AdbWrapper(serial, logger) - logger.debug("There are %d connected device(s)", len(connected_devices)) - target_devices, missing_serials = select_target_devices( - connected_devices, args.serial, logger) - # Without --serial every connected device is observed; warn in case that - # was not intended. With --serial the user has already chosen explicitly. - if args.serial is None and len(connected_devices) > 1: - logger.warning("More than 1 device connected!") + if is_hubble_installed(adb_wrapper, logger): + logger.debug("Removing previous Hubble installation...") + with events.step("uninstall_previous", serial) as s: + if not remove_previous_installation(adb_wrapper): + logger.error("Failed to remove previous Hubble installation.") + progress.error = _error(REASON_UNINSTALL_FAILED, s.fail( + "Failed to remove previous Hubble installation.")) + return - results = {} - prefetched = False - has_errors = bool(missing_serials) - verification_failed_devices = set() - collection_error_devices = set() - for target_device in target_devices: - try: - if target_device.unauthorized: - logger.error("Please authorize device with serial number %s for ADB via " - "device GUI.", target_device.serial_number) - has_errors = True - continue + with events.step("install_hubble", serial) as s: + if is_xiaomi_phone(adb_wrapper, logger): + logger.info("This is a Xiaomi phone.") + adb_push_hubble(adb_wrapper, args.hubble) + if not launch_xiaomi_file_explorer(adb_wrapper): + logger.error("Failed to launch Xiaomi file explorer") + if not wait_for_xiaomi_manual_install(adb_wrapper, serial, logger, + events): + progress.error = _error(REASON_STDIN_CLOSED, s.fail( + "Standard input closed while waiting for manual Hubble " + "installation.")) + return + else: + logger.info("This is not a Xiaomi phone. Regular workflow continues...") + if not install_hubble(adb_wrapper, args, logger): + logger.error("Error installing Hubble: %s", adb_wrapper.error_message) + progress.error = _error(REASON_INSTALL_FAILED, s.fail( + "Error installing Hubble: {}".format(adb_wrapper.error_message))) + return - # set up an adb_wrapper to be used throughout for this target device - adb_wrapper = AdbWrapper(target_device.serial_number, logger) + with events.step("launch_hubble", serial) as s: + clear_logcat(adb_wrapper) + if not launch_hubble(adb_wrapper): + logger.error("Failed to launch Hubble: %s", adb_wrapper.error_message) + progress.error = _error(REASON_LAUNCH_FAILED, s.fail( + "Failed to launch Hubble: {}".format(adb_wrapper.error_message))) + return - if is_hubble_installed(adb_wrapper, logger): - logger.debug("Removing previous Hubble installation...") - if not remove_previous_installation(adb_wrapper): - logger.error("Failed to remove previous Hubble installation.") - has_errors = True - continue - - if is_xiaomi_phone(adb_wrapper, logger): - logger.info("This is a Xiaomi phone.") - adb_push_hubble(adb_wrapper, args.hubble) - if not launch_xiaomi_file_explorer(adb_wrapper): - logger.error("Failed to launch Xiaomi file explorer") - while not is_hubble_installed(adb_wrapper, logger): - logger.warning("Please manually install Hubble by launching the " - "\"Files Manager\" app (it may have been launched " - "for you) and navigate to the \"Downloads\" folder.") - input("Press [ENTER] when you are done.") - else: - logger.info("This is not a Xiaomi phone. Regular workflow continues...") - if not install_hubble(adb_wrapper, args, logger): - logger.error("Error installing Hubble: %s", adb_wrapper.error_message) - has_errors = True - continue - clear_logcat(adb_wrapper) - - if not launch_hubble(adb_wrapper): - logger.error("Failed to launch Hubble: %s", adb_wrapper.error_message) - has_errors = True - continue - results_source = wait_for_results(adb_wrapper, logger) - - if not results_source: - logger.error("Failed to obtain results from Hubble execution.") - has_errors = True - continue - extract_apks = (args.pull_all_apks is not None) or args.pull_preinstalled_apks_only - with tempfile.TemporaryDirectory() as device_tmp_dir: - results_dir = extract_results_and_apks( - adb_wrapper, - results_source, - args.output, - logger, - extract_apks, - pull_preinstalled_only=args.pull_preinstalled_apks_only, - tmp_dir=device_tmp_dir) - - if not results_dir: - logger.error("Failed to extract results from target device (%s).", - target_device.serial_number) - has_errors = True - continue - - extract_selinux_policies(adb_wrapper, results_dir, logger) - results[target_device.serial_number] = results_dir - - if args.perform_inclusion_proof_check: - pkg_filename = ( - "preinstalled_packages.txt" - if args.check_preinstalled_only - else "packages.txt" - ) - packages_txt_path = os.path.join(results_dir, "results", pkg_filename) - if args.check_preinstalled_only and not os.path.isfile(packages_txt_path): - logger.error( - "preinstalled_packages.txt not found at %s (extracted results at %s). " - "--check_preinstalled_only requires Hubble >= 2.1.0.", - packages_txt_path, results_dir) - has_errors = True - verification_failed_devices.add(target_device.serial_number) - continue - if (not args.no_prefetch and not prefetched and - os.path.isfile(packages_txt_path)): - prefetched = inclusion_proof_check.prefetch_log_entries( - args.verifier_path, - logger, - cache_dir=args.cache_dir, - concurrency=args.cache_prefetch_concurrency, - timeout=args.cache_prefetch_timeout) - if not inclusion_proof_check.perform_inclusion_proof_check( - args.verifier_path, - packages_txt_path, - logger, - cache_dir=args.cache_dir, - concurrency=args.cache_prefetch_concurrency, - timeout=args.cache_prefetch_timeout, - prefetch=False, - preinstalled_only=args.check_preinstalled_only): - has_errors = True - verification_failed_devices.add(target_device.serial_number) - except Exception as e: - logger.exception("Unexpected error while processing device %s: %s", - target_device.serial_number, e) - has_errors = True - if target_device.serial_number in results: - collection_error_devices.add(target_device.serial_number) + with events.step("wait_for_results", serial) as s: + results_source = wait_for_results(adb_wrapper, logger) + if not results_source: + logger.error("Failed to obtain results from Hubble execution.") + progress.error = _error(REASON_NO_RESULTS, s.fail( + "Failed to obtain results from Hubble execution.")) + return - # Summarise in the order devices were requested (or discovered, without - # --serial), including requested serials that were never connected. + with events.step("extract_results", serial) as s: + extract_apks = (args.pull_all_apks is not None) or args.pull_preinstalled_apks_only + with tempfile.TemporaryDirectory() as device_tmp_dir: + results_dir = extract_results_and_apks( + adb_wrapper, + results_source, + args.output, + logger, + extract_apks, + pull_preinstalled_only=args.pull_preinstalled_apks_only, + tmp_dir=device_tmp_dir, + events=events) + + if not results_dir: + logger.error("Failed to extract results from target device (%s).", + serial) + progress.error = _error(REASON_EXTRACT_FAILED, s.fail( + "Failed to extract results from device.")) + return + + with events.step("extract_selinux", serial): + extract_selinux_policies(adb_wrapper, results_dir, logger) + progress.results_dir = results_dir + + if not args.perform_inclusion_proof_check: + return + pkg_filename = ( + "preinstalled_packages.txt" + if args.check_preinstalled_only + else "packages.txt" + ) + packages_txt_path = os.path.join(results_dir, "results", pkg_filename) + if args.check_preinstalled_only and not os.path.isfile(packages_txt_path): + with events.step("inclusion_proof_check", serial) as s: + logger.error( + "preinstalled_packages.txt not found at %s (extracted results at %s). " + "--check_preinstalled_only requires Hubble >= 2.1.0.", + packages_txt_path, results_dir) + progress.error = _error( + REASON_INCLUSION_PROOF_CHECK_INCOMPLETE, + s.fail("preinstalled_packages.txt not found; " + "--check_preinstalled_only requires Hubble >= 2.1.0.")) + progress.check_incomplete = True + return + if (not args.no_prefetch and not state.prefetched and + os.path.isfile(packages_txt_path)): + with events.step("inclusion_proof_prefetch", serial) as s: + state.prefetched = inclusion_proof_check.prefetch_log_entries( + args.verifier_path, + logger, + cache_dir=args.cache_dir, + concurrency=args.cache_prefetch_concurrency, + timeout=args.cache_prefetch_timeout) + if not state.prefetched: + s.fail("Pre-fetching failed; falling back to on-demand fetching.") + with events.step("inclusion_proof_check", serial) as s: + if not inclusion_proof_check.perform_inclusion_proof_check( + args.verifier_path, + packages_txt_path, + logger, + cache_dir=args.cache_dir, + concurrency=args.cache_prefetch_concurrency, + timeout=args.cache_prefetch_timeout, + prefetch=False, + preinstalled_only=args.check_preinstalled_only, + progress=events.progress_reporter("inclusion_proof_check", serial)): + progress.check_incomplete = True + # False means the check could not complete (bad input or unwritable + # output), not that some splits are absent from the log; per-split + # results are in the *_signal.txt output. + progress.error = _error(REASON_INCLUSION_PROOF_CHECK_INCOMPLETE, + s.fail("Inclusion proof check could not " + "complete.")) + + +def print_summary(serials, outcomes: dict, missing_serials, + logger: logging.Logger) -> dict: + """Logs the final outcome line(s) for each device. + + Args: + serials: Serials in the order to report them. + outcomes: serial -> DeviceResult, for every serial in serials. + missing_serials: Serials requested via --serial that were not connected. + logger: A logger object to log messages. + + Returns: + The run_finished summary: serial -> {status, results_dir?}. + """ missing = set(missing_serials) - summary_serials = list(dict.fromkeys( - args.serial if args.serial is not None - else [d.serial_number for d in connected_devices])) - for device in summary_serials: + summary = {} + for device in serials: + outcome = outcomes[device] + status = outcome.status + summary[device] = {"status": status} + if outcome.results_dir is not None: + summary[device]["results_dir"] = outcome.results_dir + if device in missing: logger.error( "FAILED: Requested device %s is not connected (exiting 1).", device) continue - if device not in results: + if status == STATUS_FAILED: logger.error( "FAILED: Hubble data collection failed on connected device %s " "(exiting 1).", device) continue - if device in collection_error_devices: + if status == STATUS_PARTIAL_ERROR: logger.warning( "PARTIAL SUCCESS: Hubble data collection succeeded on connected " "device %s, but an unexpected error occurred during post-collection " "processing (exiting 1).", device) - elif device in verification_failed_devices: + elif status == STATUS_PARTIAL_CHECK_INCOMPLETE: logger.warning( "PARTIAL SUCCESS: Hubble data collection succeeded on connected " "device %s, but inclusion proof verification failed (exiting 1).", @@ -1107,10 +1630,163 @@ def main(): else: logger.info("SUCCESS! Hubble was successfully deployed and executed on " "connected device %s.", device) - logger.info("Hubble output files can be found at: %s", results[device]) + logger.info("Hubble output files can be found at: %s", outcome.results_dir) + return summary + + +def run(args: argparse.Namespace, logger: logging.Logger, + events: EventEmitter) -> int: + """Runs the observation workflow. - if has_errors: - sys.exit(1) + Emits run_finished on every path that returns. + + Args: + args: Parsed arguments from parse_arguments(). + logger: A logger object to log debug or error messages. + events: Where to report progress events. May be a disabled emitter. + + Returns: + The process exit code. Early exits (before any device is processed) + return 0, as they always have; run_finished.error describes them. + """ + events.emit("run_started", argv=sys.argv[1:], pid=os.getpid()) + for warning in validate_argument_combinations(args): + logger.warning(warning) + + def early_exit(reason: str, message: str) -> int: + events.finish_run(0, error=_error(reason, message)) + return 0 + + if not supported_platform(logger): + logger.error("Sorry, your OS is currently unsupported for this script.") + return early_exit(REASON_UNSUPPORTED_PLATFORM, + "Unsupported platform: {}".format(sys.platform)) + + if args.hubble is None: + logger.info("-H flag not used. Rebuilding Hubble...") + with events.step("build_hubble") as s: + hubble_path, build_error = rebuild_hubble(logger) + if build_error: + s.fail(build_error["message"]) + if build_error: + events.finish_run(0, error=build_error) + return 0 + args.hubble = hubble_path + + with events.step("verify_hubble") as s: + if not verify_hubble(args, logger): + s.fail("Hubble APK to be installed must have the .apk extension.") + if s.failed: + return early_exit(REASON_INVALID_HUBBLE_APK, s.message) + + with events.step("check_adb") as s: + if not adb_installed(logger): + s.fail("adb was not found on this system.") + if s.failed: + return early_exit(REASON_ADB_NOT_FOUND, s.message) + + with events.step("start_adb_server") as s: + if not AdbWrapper.start_server(logger): + s.fail("Failed to start the adb server.") + if s.failed: + return early_exit(REASON_ADB_SERVER_FAILED, s.message) + + with events.step("list_devices") as s: + listed_devices = AdbWrapper.devices(logger) + if listed_devices is None: + s.fail("`adb devices` failed.") + connected_devices = listed_devices or [] + # With --serial, fall through: every requested serial is then reported as + # missing, FAILED in the summary, and the run exits 1. + if not connected_devices and args.serial is None: + logger.error("No devices connected!") + if listed_devices is None: + return early_exit(REASON_ADB_DEVICES_FAILED, "`adb devices` failed.") + return early_exit(REASON_NO_DEVICES, "No devices connected!") + + logger.debug("There are %d connected device(s)", len(connected_devices)) + target_devices, missing_serials = select_target_devices( + connected_devices, args.serial, logger) + # Without --serial every connected device is observed; warn in case that + # was not intended. With --serial the user has already chosen explicitly. + if args.serial is None and len(connected_devices) > 1: + logger.warning("More than 1 device connected!") + + events.emit("devices", + devices=[_device_to_event(d) for d in connected_devices], + selected=[d.serial_number for d in target_devices], + missing=missing_serials) + outcomes = {} + for serial in missing_serials: + outcomes[serial] = DeviceResult( + STATUS_FAILED, error=_error(REASON_NOT_CONNECTED, + "Requested device is not connected.")) + events.device_finished(serial, STATUS_FAILED, + error=outcomes[serial].error) + + state = RunState() + has_errors = bool(missing_serials) + for target_device in target_devices: + result = process_device(target_device, args, logger, events, state) + # Decided per device, not from `outcomes`: two connected devices can + # share a serial, and the second must not hide the first one's failure. + has_errors = has_errors or result.status != STATUS_SUCCESS + outcomes[target_device.serial_number] = result + + # Summarise in the order devices were requested (or discovered, without + # --serial), including requested serials that were never connected. + summary_serials = list(dict.fromkeys( + args.serial if args.serial is not None + else [d.serial_number for d in connected_devices])) + summary = print_summary(summary_serials, outcomes, missing_serials, logger) + + exit_code = 1 if has_errors else 0 + events.finish_run(exit_code, summary=summary) + return exit_code + + +def main(): + args = parse_arguments() + logger = set_up_logging(args) + + events = EventEmitter(logger=logger) + if args.events is not None: + try: + events = EventEmitter(open_event_stream(args.events), logger=logger) + except OSError as e: + logger.error("Cannot open --events destination %s: %s", args.events, e) + sys.exit(1) + + # SIGTERM (how CI timeouts and cancellations usually stop a process) would + # otherwise kill the script without a run_finished event. While events are + # being written, turn it into an exception so the run can report itself, + # then die by SIGTERM anyway so the parent sees the usual exit status. + # Without --events, SIGTERM handling is left untouched. + terminated = False + with (termination.sigterm_raises() if events.enabled + else contextlib.nullcontext()): + try: + exit_code = run(args, logger, events) + except termination.Terminated: + events.finish_run(128 + signal.SIGTERM, error=_error( + REASON_TERMINATED, "Terminated by SIGTERM.")) + terminated = True + except KeyboardInterrupt: + events.finish_run(130, error=_error(REASON_INTERRUPTED, + "Interrupted by user.")) + raise + except BaseException as e: + events.finish_run(1, error=_error( + REASON_UNEXPECTED_ERROR, "{}: {}".format(type(e).__name__, e))) + raise + finally: + events.close() + + if terminated: + termination.die_by_sigterm(logger) + return + if exit_code: + sys.exit(exit_code) if __name__ == "__main__": diff --git a/uraniborg/scripts/python/inclusion_proof_check.py b/uraniborg/scripts/python/inclusion_proof_check.py index 63795d3..990deb7 100644 --- a/uraniborg/scripts/python/inclusion_proof_check.py +++ b/uraniborg/scripts/python/inclusion_proof_check.py @@ -17,20 +17,46 @@ """Performs inclusion proof check against packages in packages.txt or preinstalled_packages.txt.""" import argparse +import contextlib import json import logging import os import subprocess import sys import tempfile -from typing import Optional +from typing import Callable, NamedTuple, Optional + +import termination OUTPUT_FILENAME = 'packages_with_inclusion_proof_signal.txt' PREINSTALLED_OUTPUT_FILENAME = ( 'preinstalled_packages_with_inclusion_proof_signal.txt' ) +# progress(done, total); see perform_inclusion_proof_check(). +ProgressCallback = Callable[[int, int], None] DEFAULT_PREFETCH_CONCURRENCY = 16 DEFAULT_PREFETCH_TIMEOUT = 600 +# What a verifier without batch mode prints (Go's flag package, exit code 2) +# when given --payloads_path. Exit code 2 alone is not enough: a Go panic +# exits with 2 too. +_BATCH_UNSUPPORTED_EXIT_CODE = 2 +_BATCH_UNSUPPORTED_MESSAGE = "flag provided but not defined: -payloads_path" + + +class _SplitJob(NamedTuple): + """One APK split to verify.""" + split: dict # The split's entry in the package list; receives the result. + payload: str + label: str # Names the package and split in log lines. + + +def _split_label(package_name: str, split: dict, index: int, + split_count: int) -> str: + """Describes a split for log lines, e.g. 'com.android.chrome [chrome]'.""" + split_name = split.get("name") + if not split_name: + split_name = "split {}".format(index) if split_count > 1 else "base" + return "{} [{}]".format(package_name, split_name) def prefetch_log_entries(verifier_executable: str, @@ -89,8 +115,26 @@ def prefetch_log_entries(verifier_executable: str, def run_verifier(verifier_executable: str, payload_path: str, logger: logging.Logger, - cache_dir: Optional[str] = None) -> bool: - """Runs verifier tool and returns True if inclusion proof is successful.""" + cache_dir: Optional[str] = None, + label: Optional[str] = None) -> bool: + """Runs verifier tool and returns True if inclusion proof is successful. + + If interrupted (KeyboardInterrupt, or an exception raised by a signal + handler), the verifier is killed and the exception propagates. + + Args: + verifier_executable: path to verifier tool. + payload_path: path to the payload file to verify. + logger: logger instance. + cache_dir: optional custom root directory for local cache. + label: names the package and split in log lines. Defaults to + payload_path. + + Returns: + True if the inclusion proof succeeded; False if it failed or the verifier + could not be run. + """ + what = label or payload_path try: cmd = [verifier_executable, f"--payload_path={payload_path}", "--log_type=google_1p_apk"] @@ -98,26 +142,215 @@ def run_verifier(verifier_executable: str, payload_path: str, cmd.append(f"--cache_dir={cache_dir}") with open(payload_path, "r") as f_in: payload = f_in.read() - logger.debug("payload content: %s", payload) - logger.debug("Running verifier: %s", " ".join(cmd)) - result = subprocess.run(cmd, capture_output=True, text=True, check=False) - logger.debug("Verifier stdout: %s", result.stdout) - logger.debug("Verifier stderr: %s", result.stderr) - if ("OK. inclusion check success!" in result.stdout or - "OK. inclusion check success!" in result.stderr): - logger.debug("Verifier check passed.") + logger.debug("%s: payload content: %s", what, payload) + logger.debug("%s: Running verifier: %s", what, " ".join(cmd)) + with subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + text=True) as process: + try: + stdout, stderr = process.communicate() + except BaseException: + _kill(process) + raise + logger.debug("%s: Verifier stdout: %s", what, stdout) + logger.debug("%s: Verifier stderr: %s", what, stderr) + if ("OK. inclusion check success!" in stdout or + "OK. inclusion check success!" in stderr): + logger.debug("%s: Verifier check passed.", what) return True else: - logger.debug("Verifier check failed.") + logger.debug("%s: Verifier check failed.", what) return False except FileNotFoundError: - logger.error("`%s` command not found.", verifier_executable) + logger.error("%s: `%s` command not found.", what, verifier_executable) return False except Exception as e: - logger.error("Error running verifier: %s", e) + logger.error("%s: Error running verifier: %s", what, e) return False +def _kill(process: subprocess.Popen) -> None: + """Kills and reaps a verifier. + + Popen's context manager does not wait for the process after a + KeyboardInterrupt, so wait here to not leave a zombie behind. + """ + process.kill() + process.wait() + + +def _verify_split(job: _SplitJob, verifier_executable: str, + logger: logging.Logger, cache_dir: Optional[str]) -> bool: + """Verifies one split with a verifier run of its own.""" + temp_payload_path = "" + try: + with tempfile.NamedTemporaryFile(mode="w", delete=False, + suffix=".txt") as fp: + fp.write(job.payload) + temp_payload_path = fp.name + + return run_verifier( + verifier_executable, + temp_payload_path, + logger, + cache_dir=cache_dir, + label=job.label) + finally: + if temp_payload_path and os.path.exists(temp_payload_path): + os.remove(temp_payload_path) + + +def _write_payloads_file(jobs: list) -> str: + """Writes the jobs' payloads to a new JSON Lines file; returns its path.""" + fd, path = tempfile.mkstemp(suffix=".jsonl") + try: + with os.fdopen(fd, "w") as fp: + for job in jobs: + fp.write(json.dumps({"payload": job.payload}) + "\n") + except BaseException: + os.remove(path) + raise + return path + + +def _verify_splits_batch(jobs: list, verifier_executable: str, + logger: logging.Logger, cache_dir: Optional[str], + report: Callable[[_SplitJob, bool], None]) -> list: + """Verifies jobs with a single verifier run in batch mode (--payloads_path). + + One run fetches each log's checkpoint and searches its entries once for + all splits, instead of once per split. Calls report(job, verified) as each + result arrives. If interrupted, the verifier is killed and the exception + propagates. + + Returns: + The jobs that got no result, in their original order: all of them if the + verifier does not support batch mode or could not be run, and the rest + if it stopped early or printed malformed results. Empty if the verifier + executable is missing or not executable: then every job is reported as + not verified, since running it per split would fail the same way. + """ + reported = set() + returncode = None + with contextlib.ExitStack() as cleanup: + try: + payloads_path = _write_payloads_file(jobs) + cleanup.callback(os.remove, payloads_path) + # stderr (the verifier's log) goes to a file, so that it cannot fill a + # pipe and block the verifier while stdout is read here. + stderr_file = cleanup.enter_context(tempfile.TemporaryFile(mode="w+")) + except OSError as e: + logger.warning("Could not prepare batch verification: %s", e) + return list(jobs) + + cmd = [verifier_executable, f"--payloads_path={payloads_path}", + "--log_type=google_1p_apk"] + if cache_dir: + cmd.append(f"--cache_dir={cache_dir}") + logger.debug("Running verifier in batch mode: %s", " ".join(cmd)) + try: + process = subprocess.Popen(cmd, stdout=subprocess.PIPE, + stderr=stderr_file, text=True) + except (FileNotFoundError, PermissionError) as e: + # E.g. a wrong --verifier_path. Running it once per split would fail + # the same way for each, so give every split its result now. + logger.error("Cannot run verifier `%s`: %s. Marking %d split(s) as not " + "verified.", verifier_executable, e, len(jobs)) + for job in jobs: + report(job, False) + return [] + except OSError as e: + logger.warning("Could not run verifier in batch mode: %s", e) + return list(jobs) + + with process: + try: + # Results are read as they arrive. In the main thread, signals + # interrupt this blocking read. + for line in process.stdout: + if not line.strip(): + continue + try: + result = json.loads(line) + index, verified = result["index"], result["verified"] + except (ValueError, KeyError, TypeError): + logger.warning("Ignoring malformed verifier result: %r", line) + continue + if (type(index) is not int or not 0 <= index < len(jobs) or + index in reported or type(verified) is not bool): + logger.warning("Ignoring unexpected verifier result: %r", line) + continue + reported.add(index) + job = jobs[index] + if result.get("error"): + logger.warning("%s: Inclusion proof failed: %s", job.label, + result["error"]) + logger.debug("%s: Verifier check %s.", job.label, + "passed" if verified else "failed") + report(job, verified) + returncode = process.wait() + except BaseException: + # E.g. Ctrl-C: do not leave the verifier behind. + _kill(process) + raise + stderr_file.seek(0) + stderr = stderr_file.read() + + logger.debug("Batch verifier stderr: %s", stderr) + remaining = [job for i, job in enumerate(jobs) if i not in reported] + if (returncode == _BATCH_UNSUPPORTED_EXIT_CODE and not reported and + _BATCH_UNSUPPORTED_MESSAGE in stderr): + # An older verifier, which rejects the unknown --payloads_path flag. + logger.warning("Verifier does not support batch mode; verifying %d " + "split(s) one at a time, which is slow. Rebuild the " + "verifier for faster checks.", len(jobs)) + elif remaining and returncode is not None: + logger.warning("Batch verifier exited with code %d after %d of %d " + "results; verifying the remaining %d split(s) one at a " + "time.", returncode, len(reported), len(jobs), + len(remaining)) + return remaining + + +def _verify_splits(jobs: list, verifier_executable: str, + logger: logging.Logger, cache_dir: Optional[str], + progress: Optional[ProgressCallback] = None) -> None: + """Verifies jobs, preferably with one batch-mode verifier run. + + Splits that the batch run gives no result for (e.g. with a verifier that + does not support batch mode) are then verified one at a time, with a + verifier run each. + + Stores each result as job.split["inclusion_proof_verified"]. + + If interrupted (KeyboardInterrupt, or an exception raised by a SIGTERM + handler), the running verifier is killed, no more are started, and the + exception propagates. + """ + total = len(jobs) + done_count = 0 + + def report(job: _SplitJob, verified: bool) -> None: + nonlocal done_count + job.split["inclusion_proof_verified"] = verified + done_count += 1 + if progress is not None: + progress(done_count, total) + + try: + if progress is not None: + progress(0, total) + remaining = jobs + if jobs: + remaining = _verify_splits_batch(jobs, verifier_executable, logger, + cache_dir, report) + for job in remaining: + report(job, _verify_split(job, verifier_executable, logger, cache_dir)) + except (KeyboardInterrupt, termination.Terminated): + logger.warning("Inclusion proof check interrupted after %d of %d " + "split(s).", done_count, total) + raise + + def perform_inclusion_proof_check( verifier_executable: str, packages_file_path: str, @@ -126,7 +359,8 @@ def perform_inclusion_proof_check( concurrency: int = DEFAULT_PREFETCH_CONCURRENCY, timeout: int = DEFAULT_PREFETCH_TIMEOUT, prefetch: bool = True, - preinstalled_only: Optional[bool] = None) -> bool: + preinstalled_only: Optional[bool] = None, + progress: Optional[ProgressCallback] = None) -> bool: """Reads packages.txt or preinstalled_packages.txt and performs inclusion proof check. By default, pre-fetches and locally caches transparency log entries before @@ -144,6 +378,13 @@ def perform_inclusion_proof_check( preinstalled_only: whether the input file is preinstalled_packages.txt (expects 'preinstalledPackages' key). If None, inferred from packages_file_path basename. + progress: optional callback, called in the calling thread as + progress(done, total): once with done=0 when verification + starts (after the input was read and pre-fetching, if any, + is over), then after each split is verified. total is the + number of splits to verify; skipped splits are not counted. + Not called if the input file cannot be read. Calls are not + throttled. Returns: True if the input file was valid and inclusion proof results were @@ -183,6 +424,10 @@ def perform_inclusion_proof_check( concurrency=concurrency, timeout=timeout) + # Collect the splits to verify first, then verify them (in one batch run + # if possible). Each job refers to its split's own dict, so the output + # keeps the input's package and split order. + jobs = [] for package in packages_list: if "name" not in package or "versionCode" not in package: logger.warning("Skipping package due to missing fields: %s", @@ -201,7 +446,8 @@ def perform_inclusion_proof_check( package_name) continue - for split in package["splits"]: + splits = package["splits"] + for index, split in enumerate(splits): split_hash = "" if "hash" in split: split_hash = split["hash"] @@ -216,22 +462,16 @@ def perform_inclusion_proof_check( split["hash"] = split_hash payload = f"{split_hash}\nSHA256(APK)\n{package_name}\n{version_code}\n" - temp_payload_path = "" - try: - with tempfile.NamedTemporaryFile(mode="w", delete=False, - suffix=".txt") as fp: - fp.write(payload) - temp_payload_path = fp.name - - verified = run_verifier( - verifier_executable, - temp_payload_path, - logger, - cache_dir=cache_dir) - split["inclusion_proof_verified"] = verified - finally: - if temp_payload_path and os.path.exists(temp_payload_path): - os.remove(temp_payload_path) + jobs.append(_SplitJob( + split=split, + payload=payload, + label=_split_label(package_name, split, index, len(splits)))) + + if jobs: + logger.info("Verifying %d APK split(s) in one batch verifier run...", + len(jobs)) + _verify_splits(jobs, verifier_executable, logger, cache_dir, + progress=progress) filtered_packages = [] for p in packages_list: @@ -316,14 +556,24 @@ def main(): s_handler.setFormatter(s_format) logger.addHandler(s_handler) - if not perform_inclusion_proof_check( - args.verifier_path, - args.packages_file, - logger, - cache_dir=args.cache_dir, - concurrency=args.cache_prefetch_concurrency, - timeout=args.cache_prefetch_timeout, - prefetch=not args.no_prefetch): + # Turn SIGTERM into an exception so that running verifiers are killed + # instead of being left behind, then die by SIGTERM as before. + terminated = False + with termination.sigterm_raises(): + try: + ok = perform_inclusion_proof_check( + args.verifier_path, + args.packages_file, + logger, + cache_dir=args.cache_dir, + concurrency=args.cache_prefetch_concurrency, + timeout=args.cache_prefetch_timeout, + prefetch=not args.no_prefetch) + except termination.Terminated: + terminated = True + if terminated: + termination.die_by_sigterm(logger) + if not ok: sys.exit(1) diff --git a/uraniborg/scripts/python/termination.py b/uraniborg/scripts/python/termination.py new file mode 100644 index 0000000000..ec97800 --- /dev/null +++ b/uraniborg/scripts/python/termination.py @@ -0,0 +1,83 @@ +#!/usr/bin/python3 +# Copyright 2026 Uraniborg authors. +# +# 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. +# + +"""Turns SIGTERM into an exception, so that scripts can clean up first. + +SIGTERM is how CI timeouts and cancellations usually stop a process. By +default it kills the script at once, leaving e.g. child verifiers behind or +an event stream without its last event. Usage: + + terminated = False + with termination.sigterm_raises(): + try: + ... + except termination.Terminated: + terminated = True # Clean up here. + if terminated: + termination.die_by_sigterm(logger) +""" + +import contextlib +import logging +import os +import signal +import sys + + +class Terminated(BaseException): + """Raised by the SIGTERM handler that sigterm_raises() installs. + + Derives from BaseException, like KeyboardInterrupt, so that the generic + `except Exception` handlers do not swallow it. SyscallWrapper only catches + KeyboardInterrupt, so this also propagates out of adb calls. + """ + + def __init__(self): + super().__init__("received SIGTERM") + + +def raise_terminated(signum, frame): # pylint: disable=unused-argument + """SIGTERM handler that raises Terminated.""" + # Ignore repeated SIGTERMs while cleaning up; die_by_sigterm() restores the + # default action afterwards. + signal.signal(signal.SIGTERM, signal.SIG_IGN) + raise Terminated() + + +@contextlib.contextmanager +def sigterm_raises(): + """Makes SIGTERM raise Terminated within the block. + + The previous handler is restored on leaving the block, however it is left. + """ + previous = signal.signal(signal.SIGTERM, raise_terminated) + try: + yield + finally: + signal.signal(signal.SIGTERM, previous) + + +def die_by_sigterm(logger: logging.Logger) -> None: + """Terminates this process with SIGTERM's default action. + + So the parent process sees the same exit status as without + sigterm_raises(). + """ + logger.error("Terminated by SIGTERM.") + signal.signal(signal.SIGTERM, signal.SIG_DFL) + os.kill(os.getpid(), signal.SIGTERM) + # Not reached unless SIGTERM is blocked; fall back to the conventional code. + sys.exit(128 + signal.SIGTERM) diff --git a/uraniborg/scripts/python/tests/test_automate_observation.py b/uraniborg/scripts/python/tests/test_automate_observation.py index 3e3541a..825d778 100644 --- a/uraniborg/scripts/python/tests/test_automate_observation.py +++ b/uraniborg/scripts/python/tests/test_automate_observation.py @@ -1321,6 +1321,7 @@ def test_extract_results_and_apks_destination_normalization_and_validation( logger, pull_preinstalled_only=True, tmp_dir=staging_dir, + events=None, ) # 2. Destination already ending with "results" + trailing slashes -> does not duplicate "/results" @@ -1342,6 +1343,7 @@ def test_extract_results_and_apks_destination_normalization_and_validation( logger, pull_preinstalled_only=False, tmp_dir="/tmp", + events=None, ) # 3. Empty destination "" -> defaults to /results @@ -1365,6 +1367,7 @@ def test_extract_results_and_apks_destination_normalization_and_validation( logger, pull_preinstalled_only=False, tmp_dir="/tmp", + events=None, ) # 4. Target results_dir exists as a regular file -> logs error and returns None @@ -1873,5 +1876,267 @@ def test_main_without_serial_no_devices_keeps_early_return( m["logger"].error.assert_any_call("No devices connected!") +_NO_PROOF = "has no effect without --perform_inclusion_proof_check." +_BOTH_PULL = ("--pull-all-apks and --pull-preinstalled-apks-only were both " + "given; only pre-installed APKs will be pulled.") +_PROOF = ("--perform_inclusion_proof_check", "--verifier_path=/v") + + +def _warnings_for(monkeypatch: pytest.MonkeyPatch, *extra: str) -> list[str]: + _set_argv(monkeypatch, *extra) + return automate_observation.validate_argument_combinations( + automate_observation.parse_arguments()) + + +@pytest.mark.parametrize( + "extra", + [ + (), + _PROOF, + (*_PROOF, "--check_preinstalled_only", "--cache_dir=/c", + "--cache_prefetch_concurrency=4", "--cache_prefetch_timeout=30"), + # Verification still uses the cache when pre-fetching is off. + (*_PROOF, "--no_prefetch", "--cache_dir=/c"), + ("--pull-all-apks",), + ("--pull-preinstalled-apks-only",), + (*_PROOF, "--pull-all-apks"), + # Integer flags set to their defaults count as not given. + ("--cache_prefetch_concurrency={}".format( + inclusion_proof_check.DEFAULT_PREFETCH_CONCURRENCY), + "--cache_prefetch_timeout={}".format( + inclusion_proof_check.DEFAULT_PREFETCH_TIMEOUT)), + ], + ids=["defaults", "proof_check", "proof_check_all_options", + "no_prefetch_with_cache_dir", "pull_all", + "pull_preinstalled", "proof_check_and_pull_all", + "explicit_int_defaults"], +) +def test_validate_argument_combinations_no_warnings( + monkeypatch: pytest.MonkeyPatch, extra, +): + assert _warnings_for(monkeypatch, *extra) == [] + + +@pytest.mark.parametrize( + "extra, expected", + [ + (("--check_preinstalled_only",), + "--check_preinstalled_only " + _NO_PROOF), + (("--verifier_path=/v",), "--verifier_path " + _NO_PROOF), + (("--no_prefetch",), "--no_prefetch " + _NO_PROOF), + (("--cache_dir=/c",), "--cache_dir " + _NO_PROOF), + (("--cache_prefetch_concurrency=4",), + "--cache_prefetch_concurrency " + _NO_PROOF), + (("--cache_prefetch_timeout=30",), + "--cache_prefetch_timeout " + _NO_PROOF), + ((*_PROOF, "--no_prefetch", "--cache_prefetch_concurrency=4"), + "--cache_prefetch_concurrency has no effect with --no_prefetch."), + ((*_PROOF, "--no_prefetch", "--cache_prefetch_timeout=30"), + "--cache_prefetch_timeout has no effect with --no_prefetch."), + (("--pull-all-apks", "--pull-preinstalled-apks-only"), _BOTH_PULL), + ], + ids=["check_preinstalled_only", "verifier_path", "no_prefetch", + "cache_dir", "cache_prefetch_concurrency", "cache_prefetch_timeout", + "no_prefetch_concurrency", "no_prefetch_timeout", "both_pull_flags"], +) +def test_validate_argument_combinations_one_warning_each( + monkeypatch: pytest.MonkeyPatch, extra, expected, +): + assert _warnings_for(monkeypatch, *extra) == [expected] + + +def test_validate_argument_combinations_reports_all_in_stable_order( + monkeypatch: pytest.MonkeyPatch, +): + assert _warnings_for( + monkeypatch, "--pull-preinstalled-apks-only", "--pull-all-apks", + "--cache_dir=/c", "--check_preinstalled_only") == [ + "--check_preinstalled_only " + _NO_PROOF, + "--cache_dir " + _NO_PROOF, + _BOTH_PULL, + ] + + +def test_validate_argument_combinations_no_prefetch_with_both_tuning_flags( + monkeypatch: pytest.MonkeyPatch, +): + assert _warnings_for( + monkeypatch, *_PROOF, "--no_prefetch", "--cache_prefetch_concurrency=4", + "--cache_prefetch_timeout=30") == [ + "--cache_prefetch_concurrency has no effect with --no_prefetch.", + "--cache_prefetch_timeout has no effect with --no_prefetch.", + ] + # Without the proof check, each flag gets only the "without" warning. + assert _warnings_for( + monkeypatch, "--no_prefetch", "--cache_prefetch_timeout=30") == [ + "--no_prefetch " + _NO_PROOF, + "--cache_prefetch_timeout " + _NO_PROOF, + ] + + +def test_main_logs_argument_warnings_before_doing_any_work( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, +): + m = serial_main_mocks + m["supported_platform"].return_value = False # stop right after the checks + _set_argv(monkeypatch, "--no_prefetch", "--pull-all-apks", + "--pull-preinstalled-apks-only") + + automate_observation.main() + + assert m["logger"].warning.call_args_list == [ + mock.call("--no_prefetch " + _NO_PROOF), mock.call(_BOTH_PULL)] + + +def test_main_logs_no_argument_warnings_by_default( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, +): + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [_make_mock_device("DEV1")] + _set_argv(monkeypatch) + + automate_observation.main() + + m["logger"].warning.assert_not_called() + + +@pytest.mark.parametrize( + "collected, collection_error, check_incomplete, expected", + [ + (False, False, False, "failed"), + (False, True, True, "failed"), + (True, True, False, "partial_error"), + (True, True, True, "partial_error"), + (True, False, True, "partial_check_incomplete"), + (True, False, False, "success"), + ], +) +def test_classify_device(collected, collection_error, check_incomplete, + expected): + assert automate_observation.classify_device( + collected, collection_error, check_incomplete) == expected + + +def _process(monkeypatch, device, *extra, state=None): + _set_argv(monkeypatch, *extra) + args = automate_observation.parse_arguments() + return automate_observation.process_device( + device, args, mock.Mock(), automate_observation.EventEmitter(), + state if state is not None else automate_observation.RunState()) + + +def test_process_device_success(serial_main_mocks, monkeypatch): + result = _process(monkeypatch, _make_mock_device("DEV1")) + assert result == automate_observation.DeviceResult( + status="success", results_dir="/tmp/out/DEV1") + assert serial_main_mocks["processed"] == ["DEV1"] + + +def test_process_device_unauthorized_touches_nothing( + serial_main_mocks, monkeypatch, +): + device = _make_mock_device("DEV1") + device.unauthorized = True + result = _process(monkeypatch, device) + assert result.status == "failed" + assert result.results_dir is None + assert result.error["reason"] == "unauthorized" + serial_main_mocks["AdbWrapper"].assert_not_called() + + +def test_process_device_error_after_collection_is_partial( + serial_main_mocks, monkeypatch, +): + with mock.patch("automate_observation.os.path.isfile", return_value=True), \ + mock.patch("inclusion_proof_check.prefetch_log_entries", + return_value=True), \ + mock.patch("inclusion_proof_check.perform_inclusion_proof_check", + side_effect=RuntimeError("boom")): + result = _process(monkeypatch, _make_mock_device("DEV1"), + "--perform_inclusion_proof_check", "--verifier_path=/v") + assert result == automate_observation.DeviceResult( + status="partial_error", results_dir="/tmp/out/DEV1", + error={"reason": "unexpected_error", "message": "RuntimeError: boom"}) + + +def test_process_device_run_state_prefetches_once( + serial_main_mocks, monkeypatch, +): + state = automate_observation.RunState() + with mock.patch("automate_observation.os.path.isfile", return_value=True), \ + mock.patch("inclusion_proof_check.prefetch_log_entries", + return_value=True) as prefetch, \ + mock.patch("inclusion_proof_check.perform_inclusion_proof_check", + return_value=True): + for serial in ("DEV1", "DEV2"): + result = _process(monkeypatch, _make_mock_device(serial), + "--perform_inclusion_proof_check", "--verifier_path=/v", + state=state) + assert result.status == "success" + prefetch.assert_called_once() + assert state.prefetched is True + + +def test_print_summary_logs_each_outcome_in_order(): + DeviceResult = automate_observation.DeviceResult + outcomes = { + "OK": DeviceResult("success", "/r/OK"), + "GONE": DeviceResult("failed", error={"reason": "not_connected", + "message": "x"}), + "BAD": DeviceResult("failed"), + "ERR": DeviceResult("partial_error", "/r/ERR"), + "CHK": DeviceResult("partial_check_incomplete", "/r/CHK"), + } + logger = mock.Mock() + manager = mock.Mock() + manager.attach_mock(logger.info, "info") + manager.attach_mock(logger.warning, "warning") + manager.attach_mock(logger.error, "error") + + summary = automate_observation.print_summary( + ["OK", "GONE", "BAD", "ERR", "CHK"], outcomes, ["GONE"], logger) + + assert summary == { + "OK": {"status": "success", "results_dir": "/r/OK"}, + "GONE": {"status": "failed"}, + "BAD": {"status": "failed"}, + "ERR": {"status": "partial_error", "results_dir": "/r/ERR"}, + "CHK": {"status": "partial_check_incomplete", "results_dir": "/r/CHK"}, + } + assert list(summary) == ["OK", "GONE", "BAD", "ERR", "CHK"] + calls = [(c[0], c[1][0].split(":")[0], c[1][1:]) for c in manager.mock_calls] + assert calls == [ + ("info", "SUCCESS! Hubble was successfully deployed and executed on " + "connected device %s.", ("OK",)), + ("info", "Hubble output files can be found at", ("/r/OK",)), + ("error", "FAILED", ("GONE",)), + ("error", "FAILED", ("BAD",)), + ("warning", "PARTIAL SUCCESS", ("ERR",)), + ("info", "Hubble output files can be found at", ("/r/ERR",)), + ("warning", "PARTIAL SUCCESS", ("CHK",)), + ("info", "Hubble output files can be found at", ("/r/CHK",)), + ] + assert "not connected" in manager.mock_calls[2][1][0] + assert "unexpected error" in manager.mock_calls[4][1][0] + assert "inclusion proof" in manager.mock_calls[6][1][0] + + +def test_main_duplicate_serial_failure_is_not_hidden_by_later_success( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, +): + """Some devices share a serial; one failing must still make the run fail.""" + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [_make_mock_device("DUP"), + _make_mock_device("DUP")] + m["install_hubble"].side_effect = [False, True] + _set_argv(monkeypatch) + + with pytest.raises(SystemExit) as exc_info: + automate_observation.main() + + assert exc_info.value.code == 1 + assert m["processed"] == ["DUP", "DUP"] + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/uraniborg/scripts/python/tests/test_automate_observation_events.py b/uraniborg/scripts/python/tests/test_automate_observation_events.py new file mode 100644 index 0000000000..202b659 --- /dev/null +++ b/uraniborg/scripts/python/tests/test_automate_observation_events.py @@ -0,0 +1,1170 @@ +#!/usr/bin/python3 +# Copyright 2026 Google LLC +# +# 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. +# + +"""Unit tests for automate_observation.py's --events progress stream.""" + +import contextlib +import datetime +import io +import json +import os +from pathlib import Path +import signal +import subprocess +import sys +import textwrap +import time +from unittest import mock + +import pytest + +# Ensure uraniborg/scripts/python is on sys.path +SCRIPT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +sys.path.insert(0, SCRIPT_DIR) + +import automate_observation +import termination +# Shared fixture and helpers for driving main() with every collaborator mocked. +from test_automate_observation import _make_mock_device # pylint: disable=g-importing-member +from test_automate_observation import _set_argv # pylint: disable=g-importing-member +from test_automate_observation import serial_main_mocks # pylint: disable=unused-import,g-importing-member + +EventEmitter = automate_observation.EventEmitter + +_FIXED_TIME = datetime.datetime(2026, 1, 2, 3, 4, 5, 678000, + tzinfo=datetime.timezone.utc) + + +def _emitter(stream=None, logger=None): + return EventEmitter(stream if stream is not None else io.StringIO(), + logger=logger, clock=lambda: _FIXED_TIME) + + +def _parse(stream: io.StringIO) -> list[dict]: + return [json.loads(line) for line in stream.getvalue().splitlines()] + + +def _read_events(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text().splitlines()] + + +def _steps(events: list[dict], device=None) -> list[tuple[str, str]]: + return [(e["step"], e["state"]) for e in events + if e["type"] == "step" and e.get("device") == device] + + +def _of_type(events: list[dict], event_type: str) -> list[dict]: + return [e for e in events if e["type"] == event_type] + + +# --- EventEmitter --------------------------------------------------------------- + + +def test_emitter_without_stream_is_a_noop(): + emitter = EventEmitter() + assert not emitter.enabled + emitter.emit("anything", x=1) + with emitter.step("s") as s: + s.fail("nope") + emitter.finish_run(1) + emitter.close() # must not raise + + +def test_emit_envelope_drops_none_and_flushes_each_event(): + stream = mock.Mock(wraps=io.StringIO()) + emitter = EventEmitter(stream, clock=lambda: _FIXED_TIME) + + emitter.emit("thing", device="D1", results_dir=None, count=2) + assert stream.flush.call_count == 1 + emitter.emit("other") + assert stream.flush.call_count == 2 + + lines = stream.getvalue().splitlines() + assert json.loads(lines[0]) == { + "v": 1, "ts": "2026-01-02T03:04:05.678Z", "type": "thing", + "device": "D1", "count": 2, + } + assert json.loads(lines[1]) == { + "v": 1, "ts": "2026-01-02T03:04:05.678Z", "type": "other"} + + +def test_emit_write_failure_warns_once_and_disables(): + stream = mock.Mock() + stream.write.side_effect = BrokenPipeError("reader went away") + logger = mock.Mock() + emitter = EventEmitter(stream, logger=logger) + + emitter.emit("a") + emitter.emit("b") + + assert not emitter.enabled + assert stream.write.call_count == 1 + logger.warning.assert_called_once() + + +def test_step_reports_finished_failed_and_exceptions(): + stream = io.StringIO() + emitter = _emitter(stream) + + with emitter.step("ok", device="D1"): + pass + with emitter.step("soft_fail") as s: + assert s.fail("went wrong") == "went wrong" + with pytest.raises(RuntimeError): + with emitter.step("boom", device="D1"): + raise RuntimeError("kaput") + + events = _parse(stream) + assert [(e["step"], e["state"], e.get("device")) for e in events] == [ + ("ok", "started", "D1"), ("ok", "finished", "D1"), + ("soft_fail", "started", None), ("soft_fail", "failed", None), + ("boom", "started", "D1"), ("boom", "failed", "D1"), + ] + assert "message" not in events[1] + assert events[3]["message"] == "went wrong" + assert events[5]["message"] == "RuntimeError: kaput" + for e in events[1::2]: + assert isinstance(e["duration_ms"], int) and e["duration_ms"] >= 0 + for e in events[0::2]: + assert "duration_ms" not in e + + +def test_step_reports_finished_when_body_continues_a_loop(): + """`continue` inside the with-block must still close the step.""" + stream = io.StringIO() + emitter = _emitter(stream) + for _ in range(1): + with emitter.step("loop") as s: + s.fail("skipped") + continue + assert [(e["step"], e["state"]) for e in _parse(stream)] == [ + ("loop", "started"), ("loop", "failed")] + + +class _FakeMonotonic: + + def __init__(self): + self.now = 100.0 + + def __call__(self): + return self.now + + +def _progress(events: list[dict]) -> list[tuple[int, int]]: + return [(e["done"], e["total"]) for e in _of_type(events, "step_progress")] + + +def test_progress_reporter_event_shape(): + stream = io.StringIO() + emitter = EventEmitter(stream, clock=lambda: _FIXED_TIME) + emitter.progress_reporter("inclusion_proof_check", "D1")(0, 5) + emitter.progress_reporter("other_step")(0, 2) + assert _parse(stream) == [ + {"v": 1, "ts": "2026-01-02T03:04:05.678Z", "type": "step_progress", + "step": "inclusion_proof_check", "device": "D1", "done": 0, + "total": 5}, + {"v": 1, "ts": "2026-01-02T03:04:05.678Z", "type": "step_progress", + "step": "other_step", "done": 0, "total": 2}, + ] + + +def test_progress_reporter_throttles_but_reports_first_and_last(): + stream = io.StringIO() + clock = _FakeMonotonic() + emitter = EventEmitter(stream, clock=lambda: _FIXED_TIME, monotonic=clock) + report = emitter.progress_reporter("s", "D1", interval=2.0) + + report(0, 10) # first: reported + report(1, 10) # too soon + clock.now += 1.9 + report(2, 10) # still too soon + clock.now += 0.1 + report(3, 10) # 2 s after the last event: reported + report(4, 10) # too soon + clock.now += 5 + report(5, 10) # reported + report(5, 10) # same done again: never reported + report(10, 10) # last: always reported, even right away + clock.now += 5 + report(10, 10) # repeat of the last: not reported + + assert _progress(_parse(stream)) == [(0, 10), (3, 10), (5, 10), (10, 10)] + + +def test_progress_reporter_empty_step_reports_once(): + stream = io.StringIO() + report = _emitter(stream).progress_reporter("s") + report(0, 0) + assert _progress(_parse(stream)) == [(0, 0)] + + +def test_progress_reporter_without_stream_is_a_noop(): + report = EventEmitter().progress_reporter("s", "D1") + report(0, 3) + report(3, 3) # must not raise + + +def test_finish_run_is_emitted_once_with_ok_semantics(): + stream = io.StringIO() + emitter = _emitter(stream) + emitter.finish_run(0, summary={"D1": {"status": "success"}}) + emitter.finish_run(1, error={"reason": "x", "message": "y"}) # ignored + (event,) = _parse(stream) + assert event["type"] == "run_finished" + assert event["exit_code"] == 0 + assert event["ok"] is True + assert event["summary"] == {"D1": {"status": "success"}} + assert "error" not in event + + for exit_code, error, ok in [(1, None, False), + (0, {"reason": "r", "message": "m"}, False)]: + stream = io.StringIO() + _emitter(stream).finish_run(exit_code, error=error) + (event,) = _parse(stream) + assert event["ok"] is ok + assert event["summary"] == {} + + +def test_device_to_event_skips_empty_and_non_string_attributes(): + dev = automate_observation.syscall_wrapper.DeviceInfo() + dev.serial_number = "S1" + dev.model_name = "Pixel_9" + dev.product_name = "" + assert automate_observation._device_to_event(dev) == { + "serial": "S1", "unauthorized": False, "model": "Pixel_9"} + # Mock devices (as used throughout the tests) must not leak Mock reprs. + assert automate_observation._device_to_event(_make_mock_device("S2")) == { + "serial": "S2", "unauthorized": False} + + +# --- main() with --events -------------------------------------------------------- + + +def test_main_events_full_run_with_missing_serial( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + m = serial_main_mocks + unauthorized = _make_mock_device("DEV_UNAUTH") + unauthorized.unauthorized = True + m["AdbWrapper"].devices.return_value = [ + _make_mock_device("DEV1"), unauthorized, _make_mock_device("DEV_OTHER")] + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path), + "-s", "DEV1", "-s", "GONE", "-s", "DEV_UNAUTH") + + with pytest.raises(SystemExit) as exc_info: + automate_observation.main() + assert exc_info.value.code == 1 + + events = _read_events(events_path) + assert all(e["v"] == 1 and e["ts"].endswith("Z") for e in events) + assert events[0]["type"] == "run_started" + assert events[0]["argv"][-6:] == [ + "-s", "DEV1", "-s", "GONE", "-s", "DEV_UNAUTH"] + assert events[-1]["type"] == "run_finished" + assert len(_of_type(events, "run_finished")) == 1 + + assert _steps(events) == [ + ("verify_hubble", "started"), ("verify_hubble", "finished"), + ("check_adb", "started"), ("check_adb", "finished"), + ("start_adb_server", "started"), ("start_adb_server", "finished"), + ("list_devices", "started"), ("list_devices", "finished"), + ] + + (devices,) = _of_type(events, "devices") + assert devices["devices"] == [ + {"serial": "DEV1", "unauthorized": False}, + {"serial": "DEV_UNAUTH", "unauthorized": True}, + {"serial": "DEV_OTHER", "unauthorized": False}, + ] + assert devices["selected"] == ["DEV1", "DEV_UNAUTH"] + assert devices["missing"] == ["GONE"] + + # Every selected or missing serial gets exactly one device_finished; + # missing ones get no device_started. + assert [e["device"] for e in _of_type(events, "device_started")] == [ + "DEV1", "DEV_UNAUTH"] + finished = {e["device"]: e for e in _of_type(events, "device_finished")} + assert set(finished) == {"GONE", "DEV1", "DEV_UNAUTH"} + assert finished["GONE"]["status"] == "failed" + assert finished["GONE"]["error"] == { + "reason": "not_connected", + "message": "Requested device is not connected."} + assert finished["DEV1"]["status"] == "success" + assert finished["DEV1"]["results_dir"] == "/tmp/out/DEV1" + assert "error" not in finished["DEV1"] + assert finished["DEV_UNAUTH"]["status"] == "failed" + assert "results_dir" not in finished["DEV_UNAUTH"] + assert finished["DEV_UNAUTH"]["error"]["reason"] == "unauthorized" + assert "not authorized" in finished["DEV_UNAUTH"]["error"]["message"] + + assert _steps(events, "DEV1") == [ + ("install_hubble", "started"), ("install_hubble", "finished"), + ("launch_hubble", "started"), ("launch_hubble", "finished"), + ("wait_for_results", "started"), ("wait_for_results", "finished"), + ("extract_results", "started"), ("extract_results", "finished"), + ("extract_selinux", "started"), ("extract_selinux", "finished"), + ] + assert _steps(events, "DEV_UNAUTH") == [] + + run_finished = events[-1] + assert run_finished["exit_code"] == 1 + assert run_finished["ok"] is False + assert "error" not in run_finished + assert run_finished["summary"] == { + "DEV1": {"status": "success", "results_dir": "/tmp/out/DEV1"}, + "GONE": {"status": "failed"}, + "DEV_UNAUTH": {"status": "failed"}, + } + assert list(run_finished["summary"]) == ["DEV1", "GONE", "DEV_UNAUTH"] + + +def test_main_events_success_exits_0_with_ok( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [_make_mock_device("DEV1")] + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path)) + + automate_observation.main() # no SystemExit + + run_finished = _read_events(events_path)[-1] + assert run_finished["type"] == "run_finished" + assert run_finished["exit_code"] == 0 + assert run_finished["ok"] is True + + +def test_main_events_previous_install_and_install_failure( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [_make_mock_device("DEV1")] + m["is_hubble_installed"].return_value = True + with mock.patch("automate_observation.remove_previous_installation", + return_value=True): + m["install_hubble"].return_value = False + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path)) + with pytest.raises(SystemExit): + automate_observation.main() + + events = _read_events(events_path) + install_failed = [e for e in events if e["type"] == "step" + and e["step"] == "install_hubble" + and e["state"] == "failed"] + assert _steps(events, "DEV1") == [ + ("uninstall_previous", "started"), ("uninstall_previous", "finished"), + ("install_hubble", "started"), ("install_hubble", "failed"), + ] + assert install_failed[0]["message"].startswith("Error installing Hubble:") + (finished,) = _of_type(events, "device_finished") + assert finished["status"] == "failed" + assert finished["error"] == {"reason": "install_failed", + "message": install_failed[0]["message"]} + + +@pytest.mark.parametrize( + "prefetch_ok, check_result, expected_status, expected_check_state", + [ + (True, True, "success", "finished"), + (False, True, "success", "finished"), + (True, False, "partial_check_incomplete", "failed"), + (True, RuntimeError("verifier crashed"), "partial_error", "failed"), + ], + ids=["verified", "prefetch_failed", "verification_failed", "crash"], +) +def test_main_events_inclusion_proof_outcomes( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, + prefetch_ok, check_result, expected_status, expected_check_state, +): + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [_make_mock_device("DEV1")] + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path), + "--perform_inclusion_proof_check", "--verifier_path=/v") + check_kwargs = ({"side_effect": check_result} + if isinstance(check_result, Exception) + else {"return_value": check_result}) + with mock.patch("automate_observation.os.path.isfile", return_value=True), \ + mock.patch("inclusion_proof_check.prefetch_log_entries", + return_value=prefetch_ok), \ + mock.patch("inclusion_proof_check.perform_inclusion_proof_check", + **check_kwargs): + if expected_status == "success": + automate_observation.main() + else: + with pytest.raises(SystemExit): + automate_observation.main() + + events = _read_events(events_path) + steps = _steps(events, "DEV1") + assert steps[-4:] == [ + ("inclusion_proof_prefetch", "started"), + ("inclusion_proof_prefetch", "finished" if prefetch_ok else "failed"), + ("inclusion_proof_check", "started"), + ("inclusion_proof_check", expected_check_state), + ] + (finished,) = _of_type(events, "device_finished") + assert finished["status"] == expected_status + # Results were collected in every case, so the directory is always reported. + assert finished["results_dir"] == "/tmp/out/DEV1" + assert events[-1]["summary"]["DEV1"]["status"] == expected_status + if expected_status == "partial_error": + assert finished["error"] == {"reason": "unexpected_error", + "message": "RuntimeError: verifier crashed"} + if expected_status == "partial_check_incomplete": + # A False return means the check could not complete, not that APKs are + # missing from the log; the message must not claim otherwise. + assert finished["error"] == { + "reason": "inclusion_proof_check_incomplete", + "message": "Inclusion proof check could not complete."} + + +def test_main_events_inclusion_proof_progress_inside_step( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [_make_mock_device("DEV1")] + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path), + "--perform_inclusion_proof_check", "--verifier_path=/v", + "--no_prefetch") + clock = _FakeMonotonic() + real_emitter = automate_observation.EventEmitter + + def emitter_with_fake_clock(*args, **kwargs): + return real_emitter(*args, monotonic=clock, **kwargs) + + def fake_check(*args, progress, **kwargs): + for done in range(5): + progress(done, 4) + clock.now += 1.5 # 1.5 s per split: every other one is reported + return True + + with mock.patch("automate_observation.os.path.isfile", return_value=True), \ + mock.patch("automate_observation.EventEmitter", + side_effect=emitter_with_fake_clock), \ + mock.patch("inclusion_proof_check.perform_inclusion_proof_check", + side_effect=fake_check): + automate_observation.main() + + events = _read_events(events_path) + check = [i for i, e in enumerate(events) + if e.get("step") == "inclusion_proof_check"] + started, finished = check[0], check[-1] + assert events[started]["type"] == "step" + assert events[started]["state"] == "started" + progress = events[started + 1:finished] + assert [(e["type"], e["step"], e["device"], e["done"], e["total"]) + for e in progress] == [ + ("step_progress", "inclusion_proof_check", "DEV1", 0, 4), + ("step_progress", "inclusion_proof_check", "DEV1", 2, 4), + ("step_progress", "inclusion_proof_check", "DEV1", 4, 4), + ] + assert events[finished]["state"] == "finished" + # No progress events outside the step. + assert len(_of_type(events, "step_progress")) == 3 + + +def test_main_events_check_preinstalled_only_missing_file( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [_make_mock_device("DEV1")] + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path), + "--perform_inclusion_proof_check", "--verifier_path=/v", + "--check_preinstalled_only") + with mock.patch("automate_observation.os.path.isfile", return_value=False): + with pytest.raises(SystemExit): + automate_observation.main() + + events = _read_events(events_path) + assert _steps(events, "DEV1")[-2:] == [ + ("inclusion_proof_check", "started"), + ("inclusion_proof_check", "failed"), + ] + (finished,) = _of_type(events, "device_finished") + assert finished["status"] == "partial_check_incomplete" + assert finished["error"]["reason"] == "inclusion_proof_check_incomplete" + assert "preinstalled_packages.txt not found" in finished["error"]["message"] + + +def _fail_uninstall(m, stack): + m["is_hubble_installed"].return_value = True + stack.enter_context(mock.patch( + "automate_observation.remove_previous_installation", return_value=False)) + + +def _fail_install(m, stack): # pylint: disable=unused-argument + m["install_hubble"].return_value = False + + +def _fail_launch(m, stack): # pylint: disable=unused-argument + m["launch_hubble"].return_value = False + + +def _fail_wait(m, stack): # pylint: disable=unused-argument + m["wait_for_results"].return_value = None + + +def _fail_extract(m, stack): # pylint: disable=unused-argument + m["extract_results_and_apks"].side_effect = lambda *a, **kw: None + + +@pytest.mark.parametrize( + "setup, failed_step, reason", + [ + (_fail_uninstall, "uninstall_previous", "uninstall_failed"), + (_fail_install, "install_hubble", "install_failed"), + (_fail_launch, "launch_hubble", "launch_failed"), + (_fail_wait, "wait_for_results", "no_results"), + (_fail_extract, "extract_results", "extract_failed"), + ], + ids=["uninstall", "install", "launch", "wait", "extract"], +) +def test_main_events_device_error_reason_per_failed_step( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, + setup, failed_step, reason, +): + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [_make_mock_device("DEV1")] + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path)) + with contextlib.ExitStack() as stack: + setup(m, stack) + with pytest.raises(SystemExit): + automate_observation.main() + + events = _read_events(events_path) + (step_failed,) = [e for e in events if e["type"] == "step" + and e["state"] == "failed"] + assert step_failed["step"] == failed_step + (finished,) = _of_type(events, "device_finished") + assert finished["status"] == "failed" + # The reason is a stable code; the message matches the failed step's. + assert finished["error"] == {"reason": reason, + "message": step_failed["message"]} + + +@pytest.mark.parametrize( + "mock_name, value, reason, failed_step", + [ + ("supported_platform", False, "unsupported_platform", None), + ("verify_hubble", False, "invalid_hubble_apk", "verify_hubble"), + ("adb_installed", False, "adb_not_found", "check_adb"), + ("AdbWrapper.start_server", False, "adb_server_failed", + "start_adb_server"), + ("AdbWrapper.devices", [], "no_devices", None), + ("AdbWrapper.devices", None, "adb_devices_failed", "list_devices"), + ], +) +def test_main_events_early_exits_keep_exit_0_and_report_reason( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, + mock_name, value, reason, failed_step, +): + m = serial_main_mocks + if "." in mock_name: + owner, attr = mock_name.split(".") + getattr(m[owner], attr).return_value = value + else: + m[mock_name].return_value = value + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path)) + + automate_observation.main() # early exits still return exit code 0 + + events = _read_events(events_path) + assert events[0]["type"] == "run_started" + run_finished = events[-1] + assert run_finished["type"] == "run_finished" + assert len(_of_type(events, "run_finished")) == 1 + assert run_finished["exit_code"] == 0 + assert run_finished["ok"] is False + assert run_finished["summary"] == {} + assert run_finished["error"]["reason"] == reason + assert run_finished["error"]["message"] + failed = [e["step"] for e in events + if e["type"] == "step" and e["state"] == "failed"] + assert failed == ([failed_step] if failed_step else []) + assert _of_type(events, "devices") == [] + + +def test_main_events_build_failure( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + events_path = tmp_path / "events.jsonl" + monkeypatch.setattr(sys, "argv", ["automate_observation.py", + "--events", str(events_path)]) + with mock.patch("automate_observation.ensure_android_sdk", + return_value=False): + automate_observation.main() + + events = _read_events(events_path) + assert _steps(events) == [("build_hubble", "started"), + ("build_hubble", "failed")] + assert events[-1]["error"] == { + "reason": "android_sdk_not_found", + "message": "Failed to (re)build Hubble APK: Android SDK not found.", + } + serial_main_mocks["verify_hubble"].assert_not_called() + + +def test_main_events_unexpected_exception_reports_and_reraises( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + m = serial_main_mocks + m["AdbWrapper"].devices.side_effect = ValueError("bad adb output") + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path)) + + with pytest.raises(ValueError): + automate_observation.main() + + events = _read_events(events_path) + assert _steps(events)[-1] == ("list_devices", "failed") + assert events[-1]["type"] == "run_finished" + assert events[-1]["exit_code"] == 1 + assert events[-1]["error"] == {"reason": "unexpected_error", + "message": "ValueError: bad adb output"} + + +@pytest.mark.parametrize( + "interrupt_at, expected_status", + [ + ("wait_for_results", "failed"), + # After results are collected: must not be reported as success. + ("prefetch", "partial_error"), + ("check", "partial_error"), + ], +) +def test_main_events_keyboard_interrupt_during_device( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, + interrupt_at, expected_status, +): + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [_make_mock_device("DEV1")] + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path), + "--perform_inclusion_proof_check", "--verifier_path=/v") + if interrupt_at == "wait_for_results": + m["wait_for_results"].side_effect = KeyboardInterrupt + prefetch = mock.patch( + "inclusion_proof_check.prefetch_log_entries", + side_effect=KeyboardInterrupt if interrupt_at == "prefetch" else None, + return_value=True) + check = mock.patch( + "inclusion_proof_check.perform_inclusion_proof_check", + side_effect=KeyboardInterrupt if interrupt_at == "check" else None, + return_value=True) + + with mock.patch("automate_observation.os.path.isfile", return_value=True), \ + prefetch, check: + with pytest.raises(KeyboardInterrupt): + automate_observation.main() + + events = _read_events(events_path) + (finished,) = _of_type(events, "device_finished") + assert finished["status"] == expected_status + assert finished["error"] == {"reason": "interrupted", + "message": "Interrupted."} + if expected_status == "partial_error": + assert finished["results_dir"] == "/tmp/out/DEV1" + assert events[-1]["type"] == "run_finished" + assert events[-1]["exit_code"] == 130 + assert events[-1]["error"]["reason"] == "interrupted" + + +def _two_devices_interrupted_on_second(m, exc): + """DEV1 completes; DEV2 raises exc while waiting for Hubble's results.""" + m["AdbWrapper"].devices.return_value = [ + _make_mock_device("DEV1"), _make_mock_device("DEV2")] + m["wait_for_results"].side_effect = [ + "/sdcard/hubble/results", exc] + + +def test_main_events_interrupt_reports_partial_summary( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + """Devices that finished before Ctrl-C stay in run_finished.summary.""" + m = serial_main_mocks + _two_devices_interrupted_on_second(m, KeyboardInterrupt) + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path), + "-s", "GONE", "-s", "DEV1", "-s", "DEV2") + + with pytest.raises(KeyboardInterrupt): + automate_observation.main() + + run_finished = _read_events(events_path)[-1] + assert run_finished["error"]["reason"] == "interrupted" + assert run_finished["summary"] == { + "GONE": {"status": "failed"}, + "DEV1": {"status": "success", "results_dir": "/tmp/out/DEV1"}, + "DEV2": {"status": "failed"}, + } + + +def test_main_events_unexpected_error_after_devices_reports_partial_summary( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [_make_mock_device("DEV1")] + + def _info(msg, *args): + if msg.startswith("SUCCESS!"): # fail while printing the log summary + raise RuntimeError("logging broke") + m["logger"].info.side_effect = _info + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path)) + + with pytest.raises(RuntimeError): + automate_observation.main() + + run_finished = _read_events(events_path)[-1] + assert run_finished["error"]["reason"] == "unexpected_error" + assert run_finished["summary"] == { + "DEV1": {"status": "success", "results_dir": "/tmp/out/DEV1"}} + + +def test_finish_run_defaults_to_recorded_devices_and_explicit_summary_wins(): + stream = io.StringIO() + emitter = _emitter(stream) + emitter.device_finished("A", "success", results_dir="/r/A") + emitter.device_finished("B", "failed", + error={"reason": "unexpected_error", + "message": "boom"}) + emitter.finish_run(1) + events = _parse(stream) + assert events[0] == {"v": 1, "ts": "2026-01-02T03:04:05.678Z", + "type": "device_finished", "device": "A", + "status": "success", "results_dir": "/r/A"} + assert events[1]["error"] == {"reason": "unexpected_error", + "message": "boom"} + assert events[-1]["summary"] == {"A": {"status": "success", + "results_dir": "/r/A"}, + "B": {"status": "failed"}} + + stream = io.StringIO() + emitter = _emitter(stream) + emitter.device_finished("A", "success") + emitter.finish_run(0, summary={"X": {"status": "success"}}) + assert _parse(stream)[-1]["summary"] == {"X": {"status": "success"}} + + +def test_main_sigterm_handler_only_installed_with_events( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + """Without --events SIGTERM is untouched; with it, restored afterwards.""" + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [_make_mock_device("DEV1")] + seen = [] + + def _wait(*unused_args): + seen.append(signal.getsignal(signal.SIGTERM)) + return "/sdcard/hubble/results" + m["wait_for_results"].side_effect = _wait + before = signal.getsignal(signal.SIGTERM) + + _set_argv(monkeypatch) + automate_observation.main() + _set_argv(monkeypatch, "--events", str(tmp_path / "events.jsonl")) + automate_observation.main() + + assert seen == [before, termination.raise_terminated] + assert signal.getsignal(signal.SIGTERM) == before + + +def test_main_events_terminated_in_process( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + """Terminated -> run_finished{terminated}, then die by SIGTERM (patched).""" + m = serial_main_mocks + _two_devices_interrupted_on_second(m, termination.Terminated()) + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path)) + before = signal.getsignal(signal.SIGTERM) + + with mock.patch("termination.die_by_sigterm") as die: + automate_observation.main() + die.assert_called_once_with(mock.ANY) + assert signal.getsignal(signal.SIGTERM) == before + + events = _read_events(events_path) + finished = {e["device"]: e for e in _of_type(events, "device_finished")} + assert finished["DEV2"]["status"] == "failed" + assert finished["DEV2"]["error"] == {"reason": "terminated", + "message": "Terminated."} + run_finished = events[-1] + assert run_finished["type"] == "run_finished" + assert run_finished["exit_code"] == 143 + assert run_finished["ok"] is False + assert run_finished["error"] == {"reason": "terminated", + "message": "Terminated by SIGTERM."} + assert run_finished["summary"] == { + "DEV1": {"status": "success", "results_dir": "/tmp/out/DEV1"}, + "DEV2": {"status": "failed"}, + } + + +_SIGTERM_DRIVER = textwrap.dedent(""" + import sys, time + from unittest import mock + sys.path.insert(0, {script_dir!r}) + import automate_observation as ao + + def dev(serial): + d = mock.Mock(); d.serial_number = serial; d.unauthorized = False + return d + + waits = iter(["/sdcard/hubble/results", None]) + def wait_for_results(*unused): + result = next(waits) + if result is None: + time.sleep(60) # DEV2 blocks here until the test sends SIGTERM + return result + + patches = dict( + supported_platform=True, verify_hubble=True, adb_installed=True, + is_hubble_installed=False, is_xiaomi_phone=False, install_hubble=True, + clear_logcat=None, launch_hubble=True, extract_selinux_policies=None) + for name, value in patches.items(): + mock.patch.object(ao, name, return_value=value).start() + adb = mock.patch.object(ao, "AdbWrapper").start() + adb.start_server.return_value = True + adb.devices.return_value = [dev("DEV1"), dev("DEV2")] + mock.patch.object(ao, "wait_for_results", + side_effect=wait_for_results).start() + mock.patch.object(ao, "extract_results_and_apks", + return_value="/tmp/out/DEV1").start() + sys.argv = ["automate_observation.py", "-H", "x.apk", + "--events", {events_path!r}] + ao.main() +""") + + +def test_events_sigterm_in_real_process(tmp_path: Path): + """A real SIGTERM mid-run: run_finished is written, exit status is -SIGTERM.""" + events_path = tmp_path / "events.jsonl" + driver = tmp_path / "driver.py" + driver.write_text(_SIGTERM_DRIVER.format(script_dir=SCRIPT_DIR, + events_path=str(events_path))) + proc = subprocess.Popen([sys.executable, str(driver)], + stdout=subprocess.PIPE, stderr=subprocess.PIPE, + text=True) + try: + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + text = events_path.read_text() if events_path.exists() else "" + if '"step": "wait_for_results", "state": "started", "device": "DEV2"' in text: + break + if proc.poll() is not None: + pytest.fail("driver exited early: " + proc.stderr.read()) + time.sleep(0.05) + else: + pytest.fail("DEV2 never reached wait_for_results") + proc.send_signal(signal.SIGTERM) + _, stderr = proc.communicate(timeout=30) + finally: + if proc.poll() is None: + proc.kill() + + assert proc.returncode == -signal.SIGTERM, stderr + events = _read_events(events_path) + assert _steps(events, "DEV2")[-1] == ("wait_for_results", "failed") + finished = {e["device"]: e for e in _of_type(events, "device_finished")} + assert finished["DEV2"]["error"] == {"reason": "terminated", + "message": "Terminated."} + run_finished = events[-1] + assert run_finished["type"] == "run_finished" + assert run_finished["exit_code"] == 143 + assert run_finished["error"]["reason"] == "terminated" + assert run_finished["summary"] == { + "DEV1": {"status": "success", "results_dir": "/tmp/out/DEV1"}, + "DEV2": {"status": "failed"}, + } + assert "Terminated by SIGTERM." in stderr + + +def test_main_events_unopenable_path_exits_1( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + bad_path = tmp_path / "no" / "such" / "dir" / "events.jsonl" + _set_argv(monkeypatch, "--events", str(bad_path)) + with pytest.raises(SystemExit) as exc_info: + automate_observation.main() + assert exc_info.value.code == 1 + serial_main_mocks["supported_platform"].assert_not_called() + + +def test_main_without_events_writes_nothing( + serial_main_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + serial_main_mocks["AdbWrapper"].devices.return_value = [ + _make_mock_device("DEV1")] + monkeypatch.chdir(tmp_path) + _set_argv(monkeypatch) + automate_observation.main() + assert list(tmp_path.iterdir()) == [] + + +# --- --events - (stdout) in a real process ----------------------------------------- + + +def test_events_to_stdout_end_to_end_is_pure_json(tmp_path: Path): + """A real run that stops early: stdout must carry only JSON Lines.""" + not_an_apk = tmp_path / "hubble.txt" + not_an_apk.write_text("") + proc = subprocess.run( + [sys.executable, os.path.join(SCRIPT_DIR, "automate_observation.py"), + "-H", str(not_an_apk), "--events", "-"], + capture_output=True, text=True, timeout=60, check=False) + + assert proc.returncode == 0, proc.stderr + events = [json.loads(line) for line in proc.stdout.splitlines()] + assert [e["type"] for e in events] == [ + "run_started", "step", "step", "run_finished"] + assert _steps(events) == [("verify_hubble", "started"), + ("verify_hubble", "failed")] + assert events[-1]["error"]["reason"] == "invalid_hubble_apk" + # Human-readable logging still goes to stderr. + assert ".apk extension" in proc.stderr + + +def test_open_event_stream_dash_moves_other_stdout_writes_to_stderr(): + """print(), input() prompts and child processes must not pollute stdout.""" + code = textwrap.dedent(""" + import subprocess, sys + sys.path.insert(0, {script_dir!r}) + import automate_observation as ao + emitter = ao.EventEmitter(ao.open_event_stream("-")) + print("noise from print") + subprocess.run([sys.executable, "-c", "print('noise from child')"]) + emitter.emit("hello") + emitter.close() + """).format(script_dir=SCRIPT_DIR) + proc = subprocess.run([sys.executable, "-c", code], capture_output=True, + text=True, timeout=60, check=False) + + assert proc.returncode == 0, proc.stderr + (line,) = proc.stdout.splitlines() + assert json.loads(line)["type"] == "hello" + assert "noise from print" in proc.stderr + assert "noise from child" in proc.stderr + + +# --- Manual-intervention prompts ------------------------------------------------ + + +def test_prompt_is_always_followed_by_prompt_resolved(): + stream = io.StringIO() + emitter = _emitter(stream) + with emitter.prompt("DEV1", "some_kind", "Do something.", + expects_input=True): + pass + with pytest.raises(KeyboardInterrupt): + with emitter.prompt(None, "other_kind", "Wait.", expects_input=False): + raise KeyboardInterrupt() + assert _parse(stream) == [ + {"v": 1, "ts": "2026-01-02T03:04:05.678Z", "type": "prompt", + "device": "DEV1", "kind": "some_kind", "message": "Do something.", + "expects_input": True}, + {"v": 1, "ts": "2026-01-02T03:04:05.678Z", "type": "prompt_resolved", + "device": "DEV1", "kind": "some_kind", "outcome": "done"}, + {"v": 1, "ts": "2026-01-02T03:04:05.678Z", "type": "prompt", + "kind": "other_kind", "message": "Wait.", "expects_input": False}, + {"v": 1, "ts": "2026-01-02T03:04:05.678Z", "type": "prompt_resolved", + "kind": "other_kind", "outcome": "interrupted"}, + ] + + +@pytest.mark.parametrize( + "body, expected_outcome", + [ + (lambda p: None, "done"), + (lambda p: setattr(p, "outcome", "failed"), "failed"), + (lambda p: setattr(p, "outcome", "stdin_closed"), "stdin_closed"), + ], + ids=["default", "failed", "stdin_closed"], +) +def test_prompt_outcome_set_by_body(body, expected_outcome): + stream = io.StringIO() + with _emitter(stream).prompt("DEV1", "k", "m", expects_input=True) as p: + body(p) + assert _parse(stream)[-1]["outcome"] == expected_outcome + + +@pytest.mark.parametrize( + "exc, expected_outcome", + [ + (KeyboardInterrupt(), "interrupted"), + (termination.Terminated(), "terminated"), + (RuntimeError("boom"), "unexpected_error"), + ], + ids=["ctrl_c", "sigterm", "exception"], +) +def test_prompt_outcome_from_exception_overrides_body(exc, expected_outcome): + stream = io.StringIO() + with pytest.raises(type(exc)): + with _emitter(stream).prompt("DEV1", "k", "m", expects_input=True) as p: + p.outcome = "failed" + raise exc + assert _parse(stream)[-1] == { + "v": 1, "ts": "2026-01-02T03:04:05.678Z", "type": "prompt_resolved", + "device": "DEV1", "kind": "k", "outcome": expected_outcome} + + +def test_prompt_without_stream_is_a_noop(): + with EventEmitter().prompt("DEV1", "k", "m", expects_input=True): + pass + + +@pytest.fixture +def xiaomi_mocks(serial_main_mocks): + m = serial_main_mocks + m["AdbWrapper"].devices.return_value = [_make_mock_device("DEV1")] + m["is_xiaomi_phone"].return_value = True + with mock.patch("automate_observation.adb_push_hubble"), \ + mock.patch("automate_observation.launch_xiaomi_file_explorer", + return_value=True): + yield m + + +def test_main_events_xiaomi_manual_install_prompt( + xiaomi_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + m = xiaomi_mocks + # Not installed beforehand; then still missing after the first Enter, and + # installed after the second. + m["is_hubble_installed"].side_effect = [False, False, False, True] + stdin = io.StringIO("\n\n") + monkeypatch.setattr(sys, "stdin", stdin) + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path)) + extract = m["extract_results_and_apks"].side_effect + emitter_enabled_at_extract = [] + + def _extract(*args, **kwargs): + emitter_enabled_at_extract.append(kwargs["events"].enabled) + return extract(*args, **kwargs) + + m["extract_results_and_apks"].side_effect = _extract + + automate_observation.main() + + assert stdin.read() == "" # both newlines consumed by the real input() + events = _read_events(events_path) + install = [e for e in events if e["type"] in ("prompt", "prompt_resolved") + or (e["type"] == "step" and e["step"] == "install_hubble")] + assert [(e["type"], e.get("state")) for e in install] == [ + ("step", "started"), ("prompt", None), ("prompt_resolved", None), + ("step", "finished")] + prompt = install[1] + assert prompt["device"] == "DEV1" + assert prompt["kind"] == "xiaomi_manual_install" + assert prompt["expects_input"] is True + assert "press Enter" in prompt["message"] + assert (install[2]["device"], install[2]["kind"], install[2]["outcome"]) == ( + "DEV1", "xiaomi_manual_install", "done") + (finished,) = _of_type(events, "device_finished") + assert finished["status"] == "success" + # The live emitter is also handed down for the adb backup prompt. + assert emitter_enabled_at_extract == [True] + + +def test_main_events_xiaomi_already_installed_after_push_has_no_prompt( + xiaomi_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + m = xiaomi_mocks + m["is_hubble_installed"].side_effect = [False, True] + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, "--events", str(events_path)) + with mock.patch("builtins.input") as fake_input: + automate_observation.main() + fake_input.assert_not_called() + assert not _of_type(_read_events(events_path), "prompt") + + +@pytest.mark.parametrize("with_events", [True, False]) +def test_main_xiaomi_stdin_closed_fails_device_and_stops_waiting( + xiaomi_mocks, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, + with_events, +): + m = xiaomi_mocks + m["is_hubble_installed"].return_value = False # never installed + monkeypatch.setattr(sys, "stdin", io.StringIO("")) # EOF on first read + events_path = tmp_path / "events.jsonl" + _set_argv(monkeypatch, + *(["--events", str(events_path)] if with_events else [])) + + with pytest.raises(SystemExit) as exc_info: + automate_observation.main() + + assert exc_info.value.code == 1 + m["launch_hubble"].assert_not_called() + # Reported as a clear error, not as an unexpected exception. + m["logger"].exception.assert_not_called() + assert any("Standard input was closed" in c.args[0] + for c in m["logger"].error.call_args_list) + if not with_events: + assert not events_path.exists() + return + events = _read_events(events_path) + assert [(e["type"], e.get("outcome")) for e in events + if e["type"] in ("prompt", "prompt_resolved")] == [ + ("prompt", None), ("prompt_resolved", "stdin_closed")] + assert _steps(events, "DEV1") == [ + ("install_hubble", "started"), ("install_hubble", "failed")] + (finished,) = _of_type(events, "device_finished") + assert finished["status"] == "failed" + assert finished["error"]["reason"] == "stdin_closed" + assert events[-1]["summary"] == {"DEV1": {"status": "failed"}} + + +def test_adb_backup_confirmation_is_reported_as_prompt(tmp_path: Path): + stream = io.StringIO() + emitter = _emitter(stream) + mock_adb = mock.Mock() + mock_adb.device_serial_number = "DEV1" + mock_adb.pull.return_value = False # forces the adb backup fallback + events_during_backup = [] + + def fake_backup(ab_path, pkg_name): # pylint: disable=unused-argument + events_during_backup.extend(_parse(stream)) + return False + + mock_adb.backup.side_effect = fake_backup + + assert automate_observation.classify_dir_using_build_fingerprint( + mock_adb, "/sdcard/hubble/results", str(tmp_path / "results"), + extract_apks=False, logger=mock.Mock(), tmp_dir=str(tmp_path), + events=emitter) is None + + # The prompt is open while `adb backup` blocks, and closed afterwards even + # though the backup failed. + assert [e["type"] for e in events_during_backup] == ["prompt"] + assert events_during_backup[0]["device"] == "DEV1" + assert events_during_backup[0]["kind"] == "adb_backup_confirm" + assert events_during_backup[0]["expects_input"] is False + assert [(e["type"], e["kind"], e.get("outcome")) + for e in _parse(stream)] == [ + ("prompt", "adb_backup_confirm", None), + ("prompt_resolved", "adb_backup_confirm", "failed")] + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/uraniborg/scripts/python/tests/test_inclusion_proof_check.py b/uraniborg/scripts/python/tests/test_inclusion_proof_check.py index 75dc5da..dc54619 100644 --- a/uraniborg/scripts/python/tests/test_inclusion_proof_check.py +++ b/uraniborg/scripts/python/tests/test_inclusion_proof_check.py @@ -20,8 +20,12 @@ import logging import os from pathlib import Path +import signal import subprocess import sys +import textwrap +import threading +import time from unittest import mock import pytest @@ -30,6 +34,11 @@ import inclusion_proof_check +_SCRIPT = os.path.abspath(os.path.join( + os.path.dirname(__file__), "..", "inclusion_proof_check.py")) +_OK = (0, "OK. inclusion check success!", "") +_NOT_FOUND = (1, "", "inclusion check failed") + @pytest.fixture def logger() -> logging.Logger: @@ -38,6 +47,123 @@ def logger() -> logging.Logger: return log +_UNSUPPORTED_BATCH = (2, [], "flag provided but not defined: -payloads_path\n") + + +def _is_batch(cmd) -> bool: + return any(arg.startswith("--payloads_path=") for arg in cmd) + + +class _FakeProcess: + """subprocess.Popen stand-in for one verifier run.""" + + def __init__(self, cmd, result_for, batch_for, **kwargs): + self.cmd = cmd + self.returncode = None + self.killed = False + self._result_for = result_for + if _is_batch(cmd): + self._batch_returncode, lines, stderr = batch_for(cmd) + self.stdout = iter(lines) + kwargs["stderr"].write(stderr) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + def communicate(self, input=None, timeout=None): # pylint: disable=redefined-builtin + self.returncode, stdout, stderr = self._result_for(self.cmd) + return stdout, stderr + + def poll(self): + return self.returncode + + def wait(self): + if self.returncode is None: + self.returncode = self._batch_returncode + return self.returncode + + def kill(self): + self.killed = True + self.returncode = -signal.SIGKILL + + +def _fake_popen(result_for, batch_for=lambda cmd: _UNSUPPORTED_BATCH): + """Returns a subprocess.Popen stand-in for verifier runs. + + Args: + result_for: for a single-payload run, called with the command when the + process is waited for (in communicate(), like a real + process's work); returns (returncode, stdout, stderr). + batch_for: for a batch run (--payloads_path), called with the command when + the process starts; returns (returncode, stdout lines, stderr). + By default, the verifier does not support batch mode. + """ + def factory(cmd, **kwargs): + return _FakeProcess(cmd, result_for, batch_for, **kwargs) + return mock.Mock(side_effect=factory) + + +def _split_cmds(popen) -> list: + """The single-payload verifier commands a fake Popen was called with.""" + return [c.args[0] for c in popen.call_args_list if not _is_batch(c.args[0])] + + +def _batch_cmds(popen) -> list: + return [c.args[0] for c in popen.call_args_list if _is_batch(c.args[0])] + + +def _batch_payloads(cmd) -> list: + """Reads the payloads of a batch verifier command.""" + prefix = "--payloads_path=" + path = next(arg[len(prefix):] for arg in cmd if arg.startswith(prefix)) + return [json.loads(line)["payload"] + for line in Path(path).read_text().splitlines()] + + +def _batch_result(index, verified, **extra) -> str: + return json.dumps(dict(index=index, verified=verified, **extra)) + "\n" + + +def _payload_of(cmd) -> str: + """Reads the payload file a verifier command refers to.""" + prefix = "--payload_path=" + path = next(arg[len(prefix):] for arg in cmd if arg.startswith(prefix)) + return Path(path).read_text() + + +def _write_packages(path: Path, packages, key="packages") -> Path: + path.write_text(json.dumps({key: packages})) + return path + + +def _write_fake_verifier(tmp_path: Path, body: str) -> Path: + """Writes an executable Python script that stands in for the verifier.""" + path = tmp_path / "fake_verifier" + path.write_text("#!{}\n".format(sys.executable) + textwrap.dedent(body)) + path.chmod(0o755) + return path + + +def _wait_for(predicate, timeout=10.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.05) + return False + + +def _process_exists(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + return True + + @mock.patch("subprocess.run") def test_prefetch_log_entries_default( mock_run: mock.MagicMock, logger: logging.Logger @@ -124,83 +250,94 @@ def test_prefetch_log_entries_file_not_found( mock_error.assert_called_once() -@mock.patch("subprocess.run") -def test_run_verifier_with_cache_dir( - mock_run: mock.MagicMock, tmp_path: Path, logger: logging.Logger -): - mock_run.return_value = mock.Mock( - returncode=0, - stdout="OK. inclusion check success!", - stderr="", - ) +def test_run_verifier_with_cache_dir(tmp_path: Path, logger: logging.Logger): payload_file = tmp_path / "payload.txt" payload_file.write_text("hash\nSHA256(APK)\ncom.example\n1\n") - verified = inclusion_proof_check.run_verifier( - "/path/to/verifier", - str(payload_file), - logger, - cache_dir="/custom/cache", - ) + with mock.patch("subprocess.Popen", _fake_popen(lambda cmd: _OK)) as popen: + verified = inclusion_proof_check.run_verifier( + "/path/to/verifier", + str(payload_file), + logger, + cache_dir="/custom/cache", + ) assert verified is True - mock_run.assert_called_once_with( + popen.assert_called_once_with( [ "/path/to/verifier", f"--payload_path={payload_file}", "--log_type=google_1p_apk", "--cache_dir=/custom/cache", ], - capture_output=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, text=True, - check=False, ) +def test_run_verifier_names_split_in_log_lines( + tmp_path: Path, logger: logging.Logger, caplog: pytest.LogCaptureFixture +): + payload_file = tmp_path / "payload.txt" + payload_file.write_text("hash\nSHA256(APK)\ncom.example\n1\n") + + with caplog.at_level(logging.DEBUG, logger=logger.name): + with mock.patch("subprocess.Popen", _fake_popen(lambda cmd: _NOT_FOUND)): + verified = inclusion_proof_check.run_verifier( + "/path/to/verifier", str(payload_file), logger, + label="com.example [config.en]") + + assert verified is False + assert caplog.records + for record in caplog.records: + assert record.getMessage().startswith("com.example [config.en]: ") + + +def test_run_verifier_file_not_found_is_failure( + tmp_path: Path, logger: logging.Logger +): + payload_file = tmp_path / "payload.txt" + payload_file.write_text("payload") + + with mock.patch("subprocess.Popen", side_effect=FileNotFoundError()): + assert inclusion_proof_check.run_verifier( + "/bad/verifier", str(payload_file), logger) is False + + @mock.patch("subprocess.run") def test_perform_inclusion_proof_check_with_prefetch( mock_run: mock.MagicMock, tmp_path: Path, logger: logging.Logger ): - def side_effect(cmd, **kwargs): - if "--fetch_entries" in cmd: - return mock.Mock(returncode=0, stdout="prefetched", stderr="") - return mock.Mock( - returncode=0, - stdout="OK. inclusion check success!", - stderr="", + mock_run.return_value = mock.Mock(returncode=0, stdout="prefetched", stderr="") + + packages_file = _write_packages(tmp_path / "packages.txt", [{ + "name": "com.google.android.gm", + "versionCode": 123, + "splits": [{"hash": "abc123hash"}], + }]) + + with mock.patch("subprocess.Popen", _fake_popen(lambda cmd: _OK)) as popen: + success = inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", + str(packages_file), + logger, + cache_dir="/tmp/cache", + concurrency=32, + timeout=45, + prefetch=True, ) - mock_run.side_effect = side_effect - - packages_file = tmp_path / "packages.txt" - packages_file.write_text(json.dumps({ - "packages": [{ - "name": "com.google.android.gm", - "versionCode": 123, - "splits": [{"hash": "abc123hash"}], - }] - })) - - success = inclusion_proof_check.perform_inclusion_proof_check( - "/path/to/verifier", - str(packages_file), - logger, - cache_dir="/tmp/cache", - concurrency=32, - timeout=45, - prefetch=True, - ) - assert success is True - assert mock_run.call_count == 2 - + mock_run.assert_called_once() prefetch_cmd = mock_run.call_args_list[0][0][0] assert "--fetch_entries" in prefetch_cmd assert "--concurrency=32" in prefetch_cmd assert "--cache_dir=/tmp/cache" in prefetch_cmd assert mock_run.call_args_list[0].kwargs["timeout"] == 45 - verify_cmd = mock_run.call_args_list[1][0][0] + assert len(_split_cmds(popen)) == 1 + verify_cmd = _split_cmds(popen)[0] assert "--log_type=google_1p_apk" in verify_cmd assert "--cache_dir=/tmp/cache" in verify_cmd @@ -214,32 +351,25 @@ def side_effect(cmd, **kwargs): def test_perform_inclusion_proof_check_prefetch_disabled( mock_run: mock.MagicMock, tmp_path: Path, logger: logging.Logger ): - mock_run.return_value = mock.Mock( - returncode=0, - stdout="OK. inclusion check success!", - stderr="", - ) - - packages_file = tmp_path / "packages.txt" - packages_file.write_text(json.dumps({ - "packages": [{ - "name": "com.google.android.gm", - "versionCode": 123, - "splits": [{"hash": "abc123hash"}], - }] - })) - - success = inclusion_proof_check.perform_inclusion_proof_check( - "/path/to/verifier", - str(packages_file), - logger, - cache_dir="/tmp/cache", - prefetch=False, - ) + packages_file = _write_packages(tmp_path / "packages.txt", [{ + "name": "com.google.android.gm", + "versionCode": 123, + "splits": [{"hash": "abc123hash"}], + }]) + + with mock.patch("subprocess.Popen", _fake_popen(lambda cmd: _OK)) as popen: + success = inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", + str(packages_file), + logger, + cache_dir="/tmp/cache", + prefetch=False, + ) assert success is True - assert mock_run.call_count == 1 - verify_cmd = mock_run.call_args_list[0][0][0] + mock_run.assert_not_called() + assert len(_split_cmds(popen)) == 1 + verify_cmd = _split_cmds(popen)[0] assert "--fetch_entries" not in verify_cmd assert "--log_type=google_1p_apk" in verify_cmd @@ -253,35 +383,26 @@ def test_perform_inclusion_proof_check_prefetch_disabled( def test_perform_inclusion_proof_check_fail_open_on_prefetch_failure( mock_run: mock.MagicMock, tmp_path: Path, logger: logging.Logger ): - def side_effect(cmd, **kwargs): - if "--fetch_entries" in cmd: - return mock.Mock(returncode=1, stdout="", stderr="prefetch network timeout") - return mock.Mock( - returncode=0, - stdout="OK. inclusion check success!", - stderr="", + mock_run.return_value = mock.Mock( + returncode=1, stdout="", stderr="prefetch network timeout") + + packages_file = _write_packages(tmp_path / "packages.txt", [{ + "name": "com.google.android.gm", + "versionCode": 123, + "splits": [{"hash": "abc123hash"}], + }]) + + with mock.patch("subprocess.Popen", _fake_popen(lambda cmd: _OK)) as popen: + success = inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", + str(packages_file), + logger, + prefetch=True, ) - mock_run.side_effect = side_effect - - packages_file = tmp_path / "packages.txt" - packages_file.write_text(json.dumps({ - "packages": [{ - "name": "com.google.android.gm", - "versionCode": 123, - "splits": [{"hash": "abc123hash"}], - }] - })) - - success = inclusion_proof_check.perform_inclusion_proof_check( - "/path/to/verifier", - str(packages_file), - logger, - prefetch=True, - ) - assert success is True - assert mock_run.call_count == 2 + mock_run.assert_called_once() + assert len(_split_cmds(popen)) == 1 output_file = tmp_path / inclusion_proof_check.OUTPUT_FILENAME assert output_file.exists() @@ -338,16 +459,7 @@ def test_perform_inclusion_proof_check_empty_packages_list_skips_prefetch_and_su def test_perform_inclusion_proof_check_with_preinstalled_packages_and_metadata( mock_run: mock.MagicMock, tmp_path: Path, logger: logging.Logger ): - def side_effect(cmd, **kwargs): - if "--fetch_entries" in cmd: - return mock.Mock(returncode=0, stdout="prefetched", stderr="") - return mock.Mock( - returncode=0, - stdout="OK. inclusion check success!", - stderr="", - ) - - mock_run.side_effect = side_effect + mock_run.return_value = mock.Mock(returncode=0, stdout="prefetched", stderr="") # Also create an existing full-run output file to verify it is NOT overwritten existing_full_output = tmp_path / inclusion_proof_check.OUTPUT_FILENAME @@ -381,13 +493,14 @@ def side_effect(cmd, **kwargs): ], })) - success = inclusion_proof_check.perform_inclusion_proof_check( - "/path/to/verifier", - str(preinstalled_file), - logger, - prefetch=True, - preinstalled_only=True, - ) + with mock.patch("subprocess.Popen", _fake_popen(lambda cmd: _OK)): + success = inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", + str(preinstalled_file), + logger, + prefetch=True, + preinstalled_only=True, + ) assert success is True # Full-run output artifact remains intact @@ -414,6 +527,529 @@ def side_effect(cmd, **kwargs): assert pkg1["splits"][0]["inclusion_proof_verified"] is True +def _many_packages(): + """Packages with several splits, in a deliberately non-sorted order.""" + return [ + {"name": "com.z.last", "versionCode": 3, "splits": [ + {"name": "base", "hash": "z0"}, + {"name": "config.en", "hash": "z1-ok"}, + {"name": "config.xxhdpi", "hash": "z2"}, + ]}, + {"name": "com.a.first", "versionCode": 1, "hash": "a-ok"}, + {"name": "com.m.middle", "versionCode": 2, "splits": [ + {"hash": "m0-ok"}, + {"hash": "m1"}, + ]}, + {"name": "com.no.hash", "versionCode": 4}, + ] + + +def test_perform_inclusion_proof_check_one_at_a_time_results_and_order( + tmp_path: Path, logger: logging.Logger +): + packages_file = _write_packages(tmp_path / "packages.txt", _many_packages()) + + def result_for(cmd): + split_hash = _payload_of(cmd).split("\n")[0] + # Finish in a scrambled order, so results complete out of order. + time.sleep((hash(split_hash) % 5) / 500) + return _OK if split_hash.endswith("-ok") else _NOT_FOUND + + with mock.patch("subprocess.Popen", _fake_popen(result_for)) as popen: + success = inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", str(packages_file), logger, prefetch=False) + + assert success is True + assert len(_split_cmds(popen)) == 6 + result_json = json.loads( + (tmp_path / inclusion_proof_check.OUTPUT_FILENAME).read_text()) + assert [ + (p["name"], [(s["hash"], s["inclusion_proof_verified"]) + for s in p["splits"]]) + for p in result_json["packages"] + ] == [ + ("com.z.last", [("z0", False), ("z1-ok", True), ("z2", False)]), + ("com.a.first", [("a-ok", True)]), + ("com.m.middle", [("m0-ok", True), ("m1", False)]), + ] + + +def test_perform_inclusion_proof_check_reports_progress( + tmp_path: Path, logger: logging.Logger +): + packages_file = _write_packages(tmp_path / "packages.txt", _many_packages()) + calls = [] + threads = set() + + def progress(done, total): + calls.append((done, total)) + threads.add(threading.get_ident()) + + with mock.patch("subprocess.Popen", _fake_popen(lambda cmd: _OK)): + assert inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", str(packages_file), logger, prefetch=False, progress=progress) + + # Skipped packages (com.no.hash) are not counted. + assert calls == [(done, 6) for done in range(7)] + assert threads == {threading.get_ident()} # Always the calling thread. + + +@mock.patch("subprocess.run") +def test_perform_inclusion_proof_check_progress_starts_after_prefetch( + mock_run: mock.MagicMock, tmp_path: Path, logger: logging.Logger +): + order = [] + mock_run.side_effect = lambda *a, **k: order.append("prefetch") or mock.Mock( + returncode=0) + packages_file = _write_packages(tmp_path / "packages.txt", _many_packages()) + + with mock.patch("subprocess.Popen", _fake_popen(lambda cmd: _OK)): + inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", str(packages_file), logger, prefetch=True, + progress=lambda done, total: order.append((done, total))) + + assert order[:2] == ["prefetch", (0, 6)] + + +def test_perform_inclusion_proof_check_progress_empty_and_invalid_input( + tmp_path: Path, logger: logging.Logger +): + calls = [] + empty = _write_packages(tmp_path / "packages.txt", []) + assert inclusion_proof_check.perform_inclusion_proof_check( + "/v", str(empty), logger, prefetch=False, + progress=lambda *a: calls.append(a)) + assert calls == [(0, 0)] + + calls.clear() + invalid = tmp_path / "bad" / "packages.txt" + invalid.parent.mkdir() + invalid.write_text("{not json") + assert not inclusion_proof_check.perform_inclusion_proof_check( + "/v", str(invalid), logger, prefetch=False, + progress=lambda *a: calls.append(a)) + assert not inclusion_proof_check.perform_inclusion_proof_check( + "/v", str(tmp_path / "missing.txt"), logger, prefetch=False, + progress=lambda *a: calls.append(a)) + assert calls == [] + + +def test_perform_inclusion_proof_check_gives_each_split_its_own_payload( + tmp_path: Path, logger: logging.Logger +): + packages_file = _write_packages(tmp_path / "packages.txt", _many_packages()) + seen = {} + lock = threading.Lock() + + def result_for(cmd): + path = next(a for a in cmd if a.startswith("--payload_path=")) + payload = _payload_of(cmd) + time.sleep(0.02) + # The file still holds this split's payload after other workers ran. + assert _payload_of(cmd) == payload + with lock: + seen[path] = payload + return _OK + + with mock.patch("subprocess.Popen", _fake_popen(result_for)): + assert inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", str(packages_file), logger, prefetch=False) + + assert len(seen) == 6 # One temp file per split... + assert sorted(seen.values()) == sorted([ + "z0\nSHA256(APK)\ncom.z.last\n3\n", + "z1-ok\nSHA256(APK)\ncom.z.last\n3\n", + "z2\nSHA256(APK)\ncom.z.last\n3\n", + "a-ok\nSHA256(APK)\ncom.a.first\n1\n", + "m0-ok\nSHA256(APK)\ncom.m.middle\n2\n", + "m1\nSHA256(APK)\ncom.m.middle\n2\n", + ]) + for path in seen: # ...removed afterwards. + assert not os.path.exists(path[len("--payload_path="):]) + + +def test_perform_inclusion_proof_check_labels_name_package_and_split( + tmp_path: Path, logger: logging.Logger, caplog: pytest.LogCaptureFixture +): + packages_file = _write_packages(tmp_path / "packages.txt", _many_packages()) + + with caplog.at_level(logging.DEBUG, logger=logger.name): + with mock.patch("subprocess.Popen", _fake_popen(lambda cmd: _OK)): + inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", str(packages_file), logger, prefetch=False) + + passed = sorted( + r.getMessage()[:-len(": Verifier check passed.")] + for r in caplog.records + if r.getMessage().endswith(": Verifier check passed.")) + assert passed == [ + "com.a.first [base]", + "com.m.middle [split 0]", + "com.m.middle [split 1]", + "com.z.last [base]", + "com.z.last [config.en]", + "com.z.last [config.xxhdpi]", + ] + + +_HANGING_VERIFIER = """ + import os, sys, time + if any(a.startswith("--payloads_path=") for a in sys.argv): + # No batch mode: what Go's flag package does with an unknown flag. + print("flag provided but not defined: -payloads_path", file=sys.stderr) + sys.exit(2) + pid_dir = os.environ["FAKE_VERIFIER_PID_DIR"] + open(os.path.join(pid_dir, str(os.getpid())), "w").close() + time.sleep(60) +""" + +# Supports batch mode: reports the first payload, then hangs. +_HANGING_BATCH_VERIFIER = """ + import os, sys, time + pid_dir = os.environ["FAKE_VERIFIER_PID_DIR"] + print('{"index": 0, "verified": true}', flush=True) + open(os.path.join(pid_dir, str(os.getpid())), "w").close() + time.sleep(60) +""" + + +def test_interrupt_kills_running_verifier_and_starts_no_more( + tmp_path: Path, logger: logging.Logger, monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture +): + pid_dir = tmp_path / "pids" + pid_dir.mkdir() + monkeypatch.setenv("FAKE_VERIFIER_PID_DIR", str(pid_dir)) + verifier = _write_fake_verifier(tmp_path, _HANGING_VERIFIER) + packages_file = _write_packages(tmp_path / "packages.txt", [ + {"name": "com.example.{}".format(i), "versionCode": 1, + "hash": "h{}".format(i)} + for i in range(6) + ]) + + def interrupt_when_started(): + if _wait_for(lambda: os.listdir(pid_dir)): + os.kill(os.getpid(), signal.SIGINT) + + threading.Thread(target=interrupt_when_started, daemon=True).start() + started = time.monotonic() + with pytest.raises(KeyboardInterrupt): + inclusion_proof_check.perform_inclusion_proof_check( + str(verifier), str(packages_file), logger, prefetch=False) + + assert time.monotonic() - started < 20 + pids = [int(name) for name in os.listdir(pid_dir)] + # Pending splits were never started. + assert len(pids) == 1 + for pid in pids: + assert _wait_for(lambda: not _process_exists(pid), timeout=5) + assert not (tmp_path / inclusion_proof_check.OUTPUT_FILENAME).exists() + assert "Inclusion proof check interrupted after 0 of 6 split(s)." in [ + r.getMessage() for r in caplog.records] + + +@pytest.mark.parametrize("verifier_body", [ + _HANGING_VERIFIER, _HANGING_BATCH_VERIFIER], + ids=["one_at_a_time", "batch"]) +def test_main_sigterm_kills_running_verifier( + tmp_path: Path, verifier_body: str +): + pid_dir = tmp_path / "pids" + pid_dir.mkdir() + verifier = _write_fake_verifier(tmp_path, verifier_body) + packages_file = _write_packages(tmp_path / "packages.txt", [ + {"name": "com.example.{}".format(i), "versionCode": 1, + "hash": "h{}".format(i)} + for i in range(4) + ]) + + process = subprocess.Popen( + [sys.executable, _SCRIPT, "--packages_file", str(packages_file), + "--verifier_path", str(verifier), "--no_prefetch"], + env=dict(os.environ, FAKE_VERIFIER_PID_DIR=str(pid_dir)), + stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True) + try: + assert _wait_for(lambda: os.listdir(pid_dir)) + time.sleep(0.2) # Let the batch result be read, if any. + process.send_signal(signal.SIGTERM) + _, stderr = process.communicate(timeout=20) + finally: + if process.poll() is None: + process.kill() + + assert process.returncode == -signal.SIGTERM + assert "Terminated by SIGTERM." in stderr + pids = [int(name) for name in os.listdir(pid_dir)] + assert len(pids) == 1 # No second verifier was started. + for pid in pids: + assert _wait_for(lambda: not _process_exists(pid), timeout=5) + assert not (tmp_path / inclusion_proof_check.OUTPUT_FILENAME).exists() + + +def test_interrupt_kills_batch_verifier( + tmp_path: Path, logger: logging.Logger, monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture +): + pid_dir = tmp_path / "pids" + pid_dir.mkdir() + monkeypatch.setenv("FAKE_VERIFIER_PID_DIR", str(pid_dir)) + verifier = _write_fake_verifier(tmp_path, _HANGING_BATCH_VERIFIER) + packages_file = _write_packages(tmp_path / "packages.txt", [ + {"name": "com.example.{}".format(i), "versionCode": 1, + "hash": "h{}".format(i)} + for i in range(3) + ]) + calls = [] + + def interrupt_when_started(): + if _wait_for(lambda: os.listdir(pid_dir) and len(calls) >= 2): + os.kill(os.getpid(), signal.SIGINT) + + threading.Thread(target=interrupt_when_started, daemon=True).start() + started = time.monotonic() + with pytest.raises(KeyboardInterrupt): + inclusion_proof_check.perform_inclusion_proof_check( + str(verifier), str(packages_file), logger, prefetch=False, + progress=lambda *a: calls.append(a)) + + assert time.monotonic() - started < 20 + # The result printed before the interrupt was reported as it arrived. + assert calls == [(0, 3), (1, 3)] + pids = [int(name) for name in os.listdir(pid_dir)] + assert len(pids) == 1 # No per-split verifiers were started. + assert _wait_for(lambda: not _process_exists(pids[0]), timeout=5) + assert "Inclusion proof check interrupted after 1 of 3 split(s)." in [ + r.getMessage() for r in caplog.records] + + +def _results_by_hash(tmp_path: Path) -> dict: + result_json = json.loads( + (tmp_path / inclusion_proof_check.OUTPUT_FILENAME).read_text()) + return {s["hash"]: s["inclusion_proof_verified"] + for p in result_json["packages"] for s in p["splits"]} + + +_MANY_PACKAGES_RESULTS = { + "z0": False, "z1-ok": True, "z2": False, "a-ok": True, "m0-ok": True, + "m1": False, +} + + +def test_perform_inclusion_proof_check_batch_mode( + tmp_path: Path, logger: logging.Logger +): + packages_file = _write_packages(tmp_path / "packages.txt", _many_packages()) + seen = {} + calls = [] + + def batch_for(cmd): + payloads = _batch_payloads(cmd) + seen["payloads"] = payloads + # Results arrive in any order. + lines = [_batch_result(i, p.split("\n")[0].endswith("-ok")) + for i, p in reversed(list(enumerate(payloads)))] + return 0, lines, "INFO Verified payloads\n" + + with mock.patch("subprocess.Popen", + _fake_popen(lambda cmd: pytest.fail("per-split run"), + batch_for)) as popen: + assert inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", str(packages_file), logger, prefetch=False, + cache_dir="/tmp/cache", progress=lambda *a: calls.append(a)) + + (cmd,) = _batch_cmds(popen) + assert _split_cmds(popen) == [] + assert "--log_type=google_1p_apk" in cmd + assert "--cache_dir=/tmp/cache" in cmd + # Payloads in input order, same format as per-split payload files. + assert seen["payloads"] == [ + "z0\nSHA256(APK)\ncom.z.last\n3\n", + "z1-ok\nSHA256(APK)\ncom.z.last\n3\n", + "z2\nSHA256(APK)\ncom.z.last\n3\n", + "a-ok\nSHA256(APK)\ncom.a.first\n1\n", + "m0-ok\nSHA256(APK)\ncom.m.middle\n2\n", + "m1\nSHA256(APK)\ncom.m.middle\n2\n", + ] + payloads_path = next(a for a in cmd if a.startswith("--payloads_path=")) + assert not os.path.exists(payloads_path[len("--payloads_path="):]) + assert _results_by_hash(tmp_path) == _MANY_PACKAGES_RESULTS + assert calls == [(done, 6) for done in range(7)] + + +def test_perform_inclusion_proof_check_falls_back_without_batch_mode( + tmp_path: Path, logger: logging.Logger, caplog: pytest.LogCaptureFixture +): + packages_file = _write_packages(tmp_path / "packages.txt", _many_packages()) + + def result_for(cmd): + return _OK if _payload_of(cmd).split("\n")[0].endswith("-ok") else _NOT_FOUND + + with caplog.at_level(logging.INFO, logger=logger.name): + with mock.patch("subprocess.Popen", _fake_popen(result_for)) as popen: + assert inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", str(packages_file), logger, prefetch=False) + + assert len(_batch_cmds(popen)) == 1 + assert len(_split_cmds(popen)) == 6 + assert _results_by_hash(tmp_path) == _MANY_PACKAGES_RESULTS + messages = [r.getMessage() for r in caplog.records] + # Slow, and fixed by rebuilding the verifier: worth a warning. + assert ("Verifier does not support batch mode; verifying 6 split(s) one at " + "a time, which is slow. Rebuild the verifier for faster checks." + in [r.getMessage() for r in caplog.records + if r.levelno == logging.WARNING]) + + +@pytest.mark.parametrize("returncode", [0, 1]) +def test_perform_inclusion_proof_check_batch_leftovers_verified_per_split( + tmp_path: Path, logger: logging.Logger, caplog: pytest.LogCaptureFixture, + returncode: int +): + packages_file = _write_packages(tmp_path / "packages.txt", _many_packages()) + calls = [] + + def batch_for(cmd): + return returncode, [ + _batch_result(1, True), # z1-ok + "not json\n", + _batch_result(1, False), # Duplicate: ignored. + _batch_result(99, True), # Out of range: ignored. + _batch_result(True, True), # Not an index: ignored. + json.dumps({"index": 2}) + "\n", # No verdict: ignored. + _batch_result(5, False, error="inclusion check error"), # m1 + "\n", + ], "verifier log\n" + + per_split = [] + lock = threading.Lock() + + def result_for(cmd): + split_hash = _payload_of(cmd).split("\n")[0] + with lock: + per_split.append(split_hash) + return _OK if split_hash.endswith("-ok") else _NOT_FOUND + + with caplog.at_level(logging.DEBUG, logger=logger.name): + with mock.patch("subprocess.Popen", _fake_popen(result_for, batch_for)): + assert inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", str(packages_file), logger, prefetch=False, + progress=lambda *a: calls.append(a)) + + # Only splits without a batch result were verified one by one. + assert sorted(per_split) == ["a-ok", "m0-ok", "z0", "z2"] + assert _results_by_hash(tmp_path) == _MANY_PACKAGES_RESULTS + assert [c[0] for c in calls] == list(range(7)) + assert {c[1] for c in calls} == {6} + + warnings = [r.getMessage() for r in caplog.records + if r.levelno == logging.WARNING] + assert "com.m.middle [split 1]: Inclusion proof failed: inclusion check error" in warnings + assert len([w for w in warnings if w.startswith("Ignoring ")]) == 5 + assert ("Batch verifier exited with code {} after 2 of 6 results; verifying " + "the remaining 4 split(s) one at a time.".format(returncode) + in warnings) + assert "Batch verifier stderr: verifier log\n" in [ + r.getMessage() for r in caplog.records] + + +def test_perform_inclusion_proof_check_batch_crash_is_not_missing_batch_mode( + tmp_path: Path, logger: logging.Logger, caplog: pytest.LogCaptureFixture +): + """A Go panic also exits 2; only the flag error means an old verifier.""" + packages_file = _write_packages(tmp_path / "packages.txt", _many_packages()) + + def batch_for(cmd): + return 2, [], "panic: runtime error: index out of range\n\ngoroutine 1\n" + + def result_for(cmd): + return _OK if _payload_of(cmd).split("\n")[0].endswith("-ok") else _NOT_FOUND + + with mock.patch("subprocess.Popen", + _fake_popen(result_for, batch_for)) as popen: + assert inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", str(packages_file), logger, prefetch=False) + + # Every split is still verified, one per verifier run... + assert len(_split_cmds(popen)) == 6 + assert _results_by_hash(tmp_path) == _MANY_PACKAGES_RESULTS + # ...but the crash is reported as such. + messages = [(r.levelno, r.getMessage()) for r in caplog.records] + assert not any("does not support batch mode" in m for _, m in messages) + assert (logging.WARNING, + "Batch verifier exited with code 2 after 0 of 6 results; verifying " + "the remaining 6 split(s) one at a time.") in messages + + +@pytest.mark.parametrize("executable", [False, True], + ids=["missing", "not_executable"]) +def test_perform_inclusion_proof_check_batch_verifier_cannot_run( + tmp_path: Path, logger: logging.Logger, caplog: pytest.LogCaptureFixture, + executable: bool +): + packages_file = _write_packages(tmp_path / "packages.txt", _many_packages()) + verifier = tmp_path / "verifier" + if executable: + verifier.write_text("not a program") + verifier.chmod(0o644) + + with mock.patch("subprocess.Popen", wraps=subprocess.Popen) as popen: + assert inclusion_proof_check.perform_inclusion_proof_check( + str(verifier), str(packages_file), logger, prefetch=False) + + assert set(_results_by_hash(tmp_path).values()) == {False} + assert len(_results_by_hash(tmp_path)) == 6 + # Tried once, in batch mode, not once more per split. + assert popen.call_count == 1 + errors = [r.getMessage() for r in caplog.records + if r.levelno == logging.ERROR] + assert len(errors) == 1 + assert errors[0].startswith("Cannot run verifier `{}`".format(verifier)) + assert errors[0].endswith("Marking 6 split(s) as not verified.") + + +@pytest.mark.parametrize("error", [FileNotFoundError, PermissionError]) +def test_perform_inclusion_proof_check_batch_temp_file_error_falls_back( + tmp_path: Path, logger: logging.Logger, caplog: pytest.LogCaptureFixture, + error: type +): + """E.g. a bad TMPDIR is not blamed on the verifier.""" + packages_file = _write_packages(tmp_path / "packages.txt", _many_packages()) + + def result_for(cmd): + return _OK if _payload_of(cmd).split("\n")[0].endswith("-ok") else _NOT_FOUND + + with mock.patch.object(inclusion_proof_check, "_write_payloads_file", + side_effect=error("no temp dir")), \ + mock.patch("subprocess.Popen", _fake_popen(result_for)) as popen: + assert inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", str(packages_file), logger, prefetch=False) + + assert not _batch_cmds(popen) + assert len(_split_cmds(popen)) == 6 + assert _results_by_hash(tmp_path) == _MANY_PACKAGES_RESULTS + messages = [r.getMessage() for r in caplog.records] + assert "Could not prepare batch verification: no temp dir" in messages + assert not any(m.startswith("Cannot run verifier") for m in messages) + + +def test_error_during_check_is_not_logged_as_interrupted( + tmp_path: Path, logger: logging.Logger, caplog: pytest.LogCaptureFixture +): + packages_file = _write_packages(tmp_path / "packages.txt", _many_packages()) + + def progress(done, total): + if done: + raise RuntimeError("event stream broke") + + with mock.patch("subprocess.Popen", _fake_popen(lambda cmd: _OK)): + with pytest.raises(RuntimeError): + inclusion_proof_check.perform_inclusion_proof_check( + "/path/to/verifier", str(packages_file), logger, prefetch=False, + progress=progress) + + assert not any("interrupted" in r.getMessage() for r in caplog.records) + + def test_main_exits_nonzero_on_failure( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ): @@ -432,5 +1068,23 @@ def test_main_exits_nonzero_on_failure( assert exc_info.value.code == 1 +@pytest.mark.parametrize("error", [KeyboardInterrupt, RuntimeError]) +def test_main_restores_sigterm_handler_on_exception( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, error: type +): + monkeypatch.setattr(sys, "argv", [ + "inclusion_proof_check.py", + f"--packages_file={tmp_path / 'packages.txt'}", + "--verifier_path=/path/to/verifier", + ]) + before = signal.getsignal(signal.SIGTERM) + with mock.patch.object(inclusion_proof_check, + "perform_inclusion_proof_check", + side_effect=error): + with pytest.raises(error): + inclusion_proof_check.main() + assert signal.getsignal(signal.SIGTERM) == before + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/verifier_tools/verify/README.md b/verifier_tools/verify/README.md index ebc40e4..77b19ff 100644 --- a/verifier_tools/verify/README.md +++ b/verifier_tools/verify/README.md @@ -51,6 +51,26 @@ where `log_type` is one of the following: * `google_1p_apk` (for Google Product Applications) * `mainline_module` (for Android Mainline Modules) +### Batch Verification Mode + +To verify many candidate binaries in one run: +``` +$ ./verifier --payloads_path=${PAYLOADS_PATH} --log_type= [--cache_dir=] +``` +The input is a [JSON Lines](https://jsonlines.org/) file with one payload per line: +``` +{"payload": "\n\n\n\n"} +``` +Each payload gets the same result as a separate `--payload_path` run, but each log's checkpoint is fetched, and its entries searched, only once for all of them. This is much faster than one run per payload, especially on a cold cache. + +One JSON result per payload is written to stdout, as soon as it is known (so not necessarily in input order). `index` is the payload's 0-based record number in the input: the first `{"payload": ...}` record is `0`, the next `1`, and so on (blank lines do not count): +``` +{"index":0,"verified":true,"log":"google_1p_apk (2026/02 Tessera)"} +{"index":2,"verified":false,"error":"inclusion check error in tlog.CheckRecord ..."} +{"index":1,"verified":false} +``` +`error` is set if a payload was found in a log but its inclusion proof failed; a payload that is in no log has neither `log` nor `error`. The exit code is 0 once every result has been written, whether or not the payloads verified, and 1 if the input cannot be read or the run is interrupted (in which case some payloads have no result). + ### Pre-fetching & Offline Cache Mode To pre-fetch and locally cache all entry tiles or legacy info files up to the current checkpoint (without requiring a payload or running an inclusion proof): @@ -69,6 +89,7 @@ This enables: | --- | --- | --- | | `--log_type`, `--log-type` | Target transparency log (`pixel`, `google_1p_code`, `google_1p_apk`, `mainline_module`). Required. | `""` | | `--payload_path`, `--payload-path` | Path to the payload file describing the candidate binary. Required for verification mode. | `""` | +| `--payloads_path`, `--payloads-path` | Path to a JSON Lines file of payloads. Required for batch verification mode; cannot be combined with `--payload_path`. | `""` | | `--fetch_entries`, `--fetch-entries` | Pre-fetch and cache all log entries locally up to the latest checkpoint. | `false` | | `--concurrency` | Number of concurrent workers for fetching Tessera entry tiles. | `16` | | `--cache_dir`, `--cache-dir` | Custom root directory for local cache. If unspecified, defaults to system cache. | OS user cache dir | diff --git a/verifier_tools/verify/cmd/verifier/verifier.go b/verifier_tools/verify/cmd/verifier/verifier.go index cd188b6..bcd57cc 100644 --- a/verifier_tools/verify/cmd/verifier/verifier.go +++ b/verifier_tools/verify/cmd/verifier/verifier.go @@ -22,11 +22,15 @@ package main import ( "bytes" "context" + "encoding/json" "flag" "fmt" + "io" "log/slog" + "maps" "os" "os/signal" + "slices" "syscall" "github.com/android/android-binary-transparency/verifier_tools/verify/internal/checkpoint" @@ -83,6 +87,7 @@ var mainlineModuleLogPubKey []byte var ( payloadPath = flag.String("payload_path", "", "Path to the payload describing the binary of interest.") + payloadsPath = flag.String("payloads_path", "", "Path to a JSON Lines file of payloads to verify in one run, one {\"payload\": \"...\"} object per line. Writes one JSON result per payload to stdout.") logType = flag.String("log_type", "", "Which log: 'pixel' or 'google_1p_code' or 'google_1p_apk' or 'mainline_module'.") fetchEntries = flag.Bool("fetch_entries", false, "Pre-fetch and cache all entries/tiles locally for the specified --log_type, without performing an inclusion proof.") concurrency = flag.Int("concurrency", tiles.DefaultTesseraFetchConcurrency, "Number of concurrent workers for fetching Tessera entry tiles.") @@ -92,6 +97,7 @@ var ( func init() { flag.StringVar(logType, "log-type", "", "Alias for --log_type.") flag.StringVar(payloadPath, "payload-path", "", "Alias for --payload_path.") + flag.StringVar(payloadsPath, "payloads-path", "", "Alias for --payloads_path.") flag.StringVar(cacheDir, "cache-dir", "", "Alias for --cache_dir.") flag.BoolVar(fetchEntries, "fetch-entries", false, "Alias for --fetch_entries.") @@ -102,14 +108,17 @@ Modes: 1. Verify binary inclusion in transparency log: %s --log_type= --payload_path= [--cache_dir=] - 2. Pre-fetch and cache entries locally for offline verification: + 2. Verify many binaries in one run (JSON Lines in, JSON Lines out): + %s --log_type= --payloads_path= [--cache_dir=] + + 3. Pre-fetch and cache entries locally for offline verification: %s --log_type= --fetch_entries [--concurrency=16] [--cache_dir=] Supported log types: pixel, google_1p_code, google_1p_apk, mainline_module Flags: -`, os.Args[0], os.Args[0], os.Args[0]) +`, os.Args[0], os.Args[0], os.Args[0], os.Args[0]) flag.PrintDefaults() } } @@ -291,8 +300,17 @@ func main() { return } + if *payloadsPath != "" { + if *payloadPath != "" { + slog.Error("specify only one of '--payload_path' and '--payloads_path'") + flag.Usage() + os.Exit(1) + } + os.Exit(runBatch(ctx, targets, *payloadsPath, os.Stdout)) + } + if *payloadPath == "" { - slog.Error("must specify either '--payload_path' to verify a binary, or '--fetch_entries' to pre-fetch log entries") + slog.Error("must specify '--payload_path' or '--payloads_path' to verify binaries, or '--fetch_entries' to pre-fetch log entries") flag.Usage() os.Exit(1) } @@ -301,16 +319,137 @@ func main() { slog.Error("Unable to open file", "path", *payloadPath, "error", err) os.Exit(1) } - // Payload should not contain excessive leading or trailing whitespace. - payloadBytes := bytes.TrimSpace(b) - payloadBytes = append(payloadBytes, '\n') + payloadBytes := normalizePayload(b) if string(b) != string(payloadBytes) { slog.Info("Reformatted payload content", "from", b, "to", payloadBytes) } - var verified bool + var result payloadResult + if err := verifyPayloads(ctx, targets, [][]byte{payloadBytes}, func(r payloadResult) { result = r }); err != nil { + slog.Error("FAILURE: verification interrupted", "error", err) + os.Exit(1) + } + switch { + case result.Verified: + slog.Info("OK. inclusion check success!", "log", result.Log) + case result.Error != "": + slog.Error("FAILURE: " + result.Error) + os.Exit(1) + default: + slog.Error("FAILURE: payload not verified in any log") + os.Exit(1) + } +} + +// normalizePayload trims excessive leading or trailing whitespace from a +// payload and terminates it with a single newline, as it appears in the logs. +func normalizePayload(b []byte) []byte { + p := append([]byte(nil), bytes.TrimSpace(b)...) + return append(p, '\n') +} + +// payloadResult is the outcome for one payload; in --payloads_path mode it is +// written to stdout as one JSON line. +type payloadResult struct { + // Index is the payload's 0-based record number in the input. + Index int `json:"index"` + // Verified is true if the payload's inclusion proof succeeded. + Verified bool `json:"verified"` + // Log names the log the payload was verified in, if Verified. + Log string `json:"log,omitempty"` + // Error is set if the payload was found in a log but could not be proven + // included. A payload that is simply not in any log has no error. + Error string `json:"error,omitempty"` +} + +// readPayloads reads a JSON Lines file of {"payload": "..."} objects. Blank +// lines are ignored, as are unknown fields. +func readPayloads(r io.Reader) ([][]byte, error) { + var payloads [][]byte + dec := json.NewDecoder(r) + for { + var rec struct { + Payload *string `json:"payload"` + } + if err := dec.Decode(&rec); err == io.EOF { + return payloads, nil + } else if err != nil { + return nil, fmt.Errorf("record %d: %w", len(payloads), err) + } + if rec.Payload == nil { + return nil, fmt.Errorf("record %d: missing \"payload\"", len(payloads)) + } + payloads = append(payloads, normalizePayload([]byte(*rec.Payload))) + } +} + +// runBatch verifies every payload in the JSON Lines file at path and writes +// one payloadResult per payload to out. Results are written as soon as they +// are known, so they are not necessarily in input order. +// +// Returns the process exit code: 0 once every result has been written, +// whether or not the payloads were verified; 1 if the input cannot be read, +// the results cannot be written, or ctx is cancelled first (in which case +// some payloads have no result). +func runBatch(ctx context.Context, targets []logTarget, path string, out io.Writer) int { + f, err := os.Open(path) + if err != nil { + slog.Error("Unable to open file", "path", path, "error", err) + return 1 + } + payloads, err := readPayloads(f) + f.Close() + if err != nil { + slog.Error("Malformed payloads file", "path", path, "error", err) + return 1 + } + + enc := json.NewEncoder(out) + var writeErr error + verified := 0 + err = verifyPayloads(ctx, targets, payloads, func(r payloadResult) { + if r.Verified { + verified++ + } + if writeErr == nil { + writeErr = enc.Encode(r) + } + }) + if err != nil { + slog.Error("FAILURE: verification interrupted", "error", err) + return 1 + } + if writeErr != nil { + slog.Error("FAILURE: unable to write results", "error", writeErr) + return 1 + } + slog.Info("Verified payloads", "total", len(payloads), "verified", verified) + return 0 +} + +// verifyPayloads checks each (normalized) payload against targets, in order, +// and calls emit exactly once per payload. +// +// A payload's outcome is the same as verifying it alone: the first log that +// contains it decides, and a failed proof there is final. But each checkpoint +// is fetched, and each log searched, once for all payloads. +// +// If ctx is cancelled, verifyPayloads returns ctx.Err() without calling emit +// for payloads it has not finished, rather than reporting them as not found. +func verifyPayloads(ctx context.Context, targets []logTarget, payloads [][]byte, emit func(payloadResult)) error { + pending := make(map[int]bool, len(payloads)) + for i := range payloads { + pending[i] = true + } + for _, target := range targets { - slog.Info("Checking log", "log", target.name, "url", target.baseURL) + if len(pending) == 0 { + break + } + if err := ctx.Err(); err != nil { + return err + } + slog.Info("Checking log", "log", target.name, "url", target.baseURL, "payloads", len(pending)) root, err := checkpoint.FromURLWithPathContext(ctx, target.baseURL, target.checkpointPath, target.verifier) if err != nil { slog.Warn("Failed to read checkpoint", "log", target.name, "error", err) @@ -318,71 +457,109 @@ func main() { } logSize := int64(root.Size) - var binaryInfoIndex int64 - var found bool - - if target.isTessera { - idx, ok, err := tiles.TesseraFindPayloadIndex(target.baseURL, logSize, payloadBytes) - if err != nil { - slog.Warn("Failed to search Tessera entry tiles", "log", target.name, "error", err) - continue - } - binaryInfoIndex = idx - found = ok - } else { - for _, filename := range target.binaryInfoFilenames { - m, err := tiles.BinaryInfosIndex(target.baseURL, filename, logSize) - if err != nil { - slog.Warn("Failed to load binary info map", "log", target.name, "file", filename, "error", err) - continue - } - if idx, ok := m[string(payloadBytes)]; ok { - binaryInfoIndex = idx - found = true - break - } - } - } - - if !found { - slog.Info("Payload not found in log", "log", target.name) - continue + found := findPayloads(target, logSize, payloads, slices.Sorted(maps.Keys(pending))) + if missing := len(pending) - len(found); missing > 0 { + slog.Info("Payload not found in log", "log", target.name, "count", missing) } var th tlog.Hash copy(th[:], root.Hash) - r := tiles.HashReader{ URL: target.baseURL, TileHeight: target.tileHeight, TreeSize: logSize, IsTessera: target.isTessera, + TileCache: make(map[string][]byte), } - slog.Debug("tlog.ProveRecord", "log", target.name, "logSize", logSize, "binaryInfoIndex", binaryInfoIndex) - rp, err := tlog.ProveRecord(logSize, binaryInfoIndex, r) - if err != nil { - slog.Error("error in tlog.ProveRecord", "log", target.name, "error", err) - os.Exit(1) + for _, i := range slices.Sorted(maps.Keys(found)) { + if err := ctx.Err(); err != nil { + return err + } + delete(pending, i) + emit(proveInclusion(target, r, logSize, th, i, payloads[i], found[i])) } + } + + // A cancelled checkpoint fetch above is only logged, so check again before + // declaring the remaining payloads not found. + if err := ctx.Err(); err != nil { + return err + } + for _, i := range slices.Sorted(maps.Keys(pending)) { + emit(payloadResult{Index: i}) + } + return nil +} - leafHash, err := tiles.PayloadHash(payloadBytes) +// findPayloads returns the leaf index in target of each payloads[i], for i in +// indices, that the log contains. +func findPayloads(target logTarget, logSize int64, payloads [][]byte, indices []int) map[int]int64 { + found := make(map[int]int64) + if target.isTessera { + wanted := make([][]byte, 0, len(indices)) + for _, i := range indices { + wanted = append(wanted, payloads[i]) + } + m, err := tiles.TesseraFindPayloadIndices(target.baseURL, logSize, wanted) if err != nil { - slog.Error("error hashing payload", "error", err) - os.Exit(1) + // m still holds the exact matches found before the error; the + // other payloads move on to the next log, as in a single-payload + // run that hits the same error. + slog.Warn("Failed to search Tessera entry tiles", "log", target.name, "found", len(m), "error", err) } - - if err := tlog.CheckRecord(rp, logSize, th, binaryInfoIndex, leafHash); err != nil { - slog.Error("FAILURE: inclusion check error in tlog.CheckRecord", "log", target.name, "error", err) - os.Exit(1) + for _, i := range indices { + if idx, ok := m[string(bytes.TrimSpace(payloads[i]))]; ok { + found[i] = idx + } } + return found + } - slog.Info("OK. inclusion check success!", "log", target.name) - verified = true - break + // Load each info file only if some payloads are still not found, like a + // single-payload run does. + remaining := indices + for _, filename := range target.binaryInfoFilenames { + if len(remaining) == 0 { + break + } + m, err := tiles.BinaryInfosIndex(target.baseURL, filename, logSize) + if err != nil { + slog.Warn("Failed to load binary info map", "log", target.name, "file", filename, "error", err) + continue + } + var next []int + for _, i := range remaining { + if idx, ok := m[string(payloads[i])]; ok { + found[i] = idx + } else { + next = append(next, i) + } + } + remaining = next } + return found +} - if !verified { - slog.Error("FAILURE: payload not verified in any log") - os.Exit(1) +// proveInclusion proves that payload is the leaf at leafIndex in target's tree +// of size logSize and root hash rootHash. +func proveInclusion(target logTarget, r tiles.HashReader, logSize int64, rootHash tlog.Hash, index int, payload []byte, leafIndex int64) payloadResult { + result := payloadResult{Index: index} + slog.Debug("tlog.ProveRecord", "log", target.name, "logSize", logSize, "binaryInfoIndex", leafIndex) + rp, err := tlog.ProveRecord(logSize, leafIndex, r) + if err != nil { + result.Error = fmt.Sprintf("error in tlog.ProveRecord for log %s: %v", target.name, err) + return result + } + leafHash, err := tiles.PayloadHash(payload) + if err != nil { + result.Error = fmt.Sprintf("error hashing payload: %v", err) + return result + } + if err := tlog.CheckRecord(rp, logSize, rootHash, leafIndex, leafHash); err != nil { + result.Error = fmt.Sprintf("inclusion check error in tlog.CheckRecord for log %s: %v", target.name, err) + return result } + result.Verified = true + result.Log = target.name + return result } diff --git a/verifier_tools/verify/cmd/verifier/verifier_batch_test.go b/verifier_tools/verify/cmd/verifier/verifier_batch_test.go new file mode 100644 index 0000000000..e7641e7 --- /dev/null +++ b/verifier_tools/verify/cmd/verifier/verifier_batch_test.go @@ -0,0 +1,431 @@ +package main + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/base64" + "encoding/binary" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + + "github.com/android/android-binary-transparency/verifier_tools/verify/internal/tiles" + "github.com/google/go-cmp/cmp" + "golang.org/x/mod/sumdb/note" + "golang.org/x/mod/sumdb/tlog" +) + +// fakeLog serves a small signed transparency log over HTTP: a checkpoint, +// hash tiles and either Tessera entry tiles or a legacy package_info.txt. +type fakeLog struct { + t *testing.T + origin string + tessera bool + tileHeight int + entries [][]byte + // corruptHashTiles serves zeroed hash tiles, so inclusion proofs fail. + corruptHashTiles bool + // failFirstEntryTile makes the oldest Tessera entry tile return HTTP 500. + failFirstEntryTile bool + + server *httptest.Server + verifier note.Verifier + checkpointRequests atomic.Int64 + infoFileRequests atomic.Int64 + entryTileRequests atomic.Int64 + hashes []tlog.Hash + signedCheckpointTxt []byte +} + +func (l *fakeLog) readHashes(indices []int64) ([]tlog.Hash, error) { + out := make([]tlog.Hash, len(indices)) + for i, idx := range indices { + if idx >= int64(len(l.hashes)) { + return nil, fmt.Errorf("no stored hash %d", idx) + } + out[i] = l.hashes[idx] + } + return out, nil +} + +func (l *fakeLog) start() *fakeLog { + l.t.Helper() + for i, e := range l.entries { + hs, err := tlog.StoredHashes(int64(i), e, tlog.HashReaderFunc(l.readHashes)) + if err != nil { + l.t.Fatalf("StoredHashes: %v", err) + } + l.hashes = append(l.hashes, hs...) + } + root, err := tlog.TreeHash(int64(len(l.entries)), tlog.HashReaderFunc(l.readHashes)) + if err != nil { + l.t.Fatalf("TreeHash: %v", err) + } + + name := strings.TrimSuffix(l.origin, "\n") + skey, vkey, err := note.GenerateKey(rand.Reader, name) + if err != nil { + l.t.Fatalf("GenerateKey: %v", err) + } + signer, err := note.NewSigner(skey) + if err != nil { + l.t.Fatalf("NewSigner: %v", err) + } + if l.verifier, err = note.NewVerifier(vkey); err != nil { + l.t.Fatalf("NewVerifier: %v", err) + } + text := fmt.Sprintf("%s%d\n%s\n", l.origin, len(l.entries), base64.StdEncoding.EncodeToString(root[:])) + if l.signedCheckpointTxt, err = note.Sign(¬e.Note{Text: text}, signer); err != nil { + l.t.Fatalf("Sign: %v", err) + } + + l.server = httptest.NewServer(http.HandlerFunc(l.serve)) + l.t.Cleanup(l.server.Close) + return l +} + +func (l *fakeLog) serve(w http.ResponseWriter, r *http.Request) { + p := strings.TrimPrefix(r.URL.Path, "/") + switch { + case p == "checkpoint": + l.checkpointRequests.Add(1) + w.Write(l.signedCheckpointTxt) + case p == PackageInfoFilename && !l.tessera: + l.infoFileRequests.Add(1) + var records []string + for i, e := range l.entries { + records = append(records, fmt.Sprintf("%d\n%s", i, bytes.TrimSpace(e))) + } + w.Write([]byte(strings.Join(records, "\n\n"))) + case strings.HasPrefix(p, "tile/entries/") && l.tessera: + l.entryTileRequests.Add(1) + var n, width int + rest := strings.TrimPrefix(p, "tile/entries/") + if _, err := fmt.Sscanf(rest, "%03d.p/%d", &n, &width); err != nil { + if _, err := fmt.Sscanf(rest, "%03d", &n); err != nil { + http.NotFound(w, r) + return + } + width = 256 + } + if n == 0 && l.failFirstEntryTile { + http.Error(w, "boom", http.StatusInternalServerError) + return + } + var buf bytes.Buffer + for _, e := range l.entries[n*256 : n*256+width] { + binary.Write(&buf, binary.BigEndian, uint16(len(e))) + buf.Write(e) + } + w.Write(buf.Bytes()) + case strings.HasPrefix(p, "tile/"): + if l.tessera { + p = fmt.Sprintf("tile/%d/%s", l.tileHeight, strings.TrimPrefix(p, "tile/")) + } + tile, err := tlog.ParseTilePath(p) + if err != nil || tile.H != l.tileHeight { + http.NotFound(w, r) + return + } + data, err := tlog.ReadTileData(tile, tlog.HashReaderFunc(l.readHashes)) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + if l.corruptHashTiles { + data = make([]byte, len(data)) + } + w.Write(data) + default: + http.NotFound(w, r) + } +} + +func (l *fakeLog) target(name string) logTarget { + target := logTarget{ + name: name, + baseURL: l.server.URL, + checkpointPath: "checkpoint", + verifier: l.verifier, + tileHeight: l.tileHeight, + isTessera: l.tessera, + } + if !l.tessera { + target.binaryInfoFilenames = []string{PackageInfoFilename} + } + return target +} + +func testPayload(name string) []byte { + return []byte(fmt.Sprintf("hash_%s\nSHA256(Signed APK)\n%s\n1\n", name, name)) +} + +// testLogs returns a Tessera log followed by a legacy log. "both" is in each; +// "new" only in the Tessera log (in its last, partial entry tile); "old" only +// in the legacy log. +func testLogs(t *testing.T, corruptTessera bool) (*fakeLog, *fakeLog) { + t.Helper() + tiles.SetCacheDir(t.TempDir()) + t.Cleanup(func() { tiles.SetCacheDir("") }) + + var tesseraEntries [][]byte + for i := 0; i < 300; i++ { + switch i { + case 10: + tesseraEntries = append(tesseraEntries, testPayload("both")) + case 290: + tesseraEntries = append(tesseraEntries, testPayload("new")) + default: + tesseraEntries = append(tesseraEntries, testPayload(fmt.Sprintf("t%d", i))) + } + } + tessera := (&fakeLog{ + t: t, + origin: "android.transparency.goog/google1p/apk/2026/1\n", + tessera: true, + tileHeight: 8, + entries: tesseraEntries, + corruptHashTiles: corruptTessera, + }).start() + legacy := (&fakeLog{ + t: t, + origin: "gstatic.com/android/binary_transparency/google1p/apk/2026/0\n", + tileHeight: 2, + entries: [][]byte{testPayload("l0"), testPayload("both"), testPayload("l2"), testPayload("old"), testPayload("l4")}, + }).start() + return tessera, legacy +} + +func collect(t *testing.T, ctx context.Context, targets []logTarget, payloads [][]byte) (map[int]payloadResult, error) { + t.Helper() + got := make(map[int]payloadResult) + err := verifyPayloads(ctx, targets, payloads, func(r payloadResult) { + if _, dup := got[r.Index]; dup { + t.Errorf("result for payload %d emitted twice", r.Index) + } + got[r.Index] = r + }) + return got, err +} + +func TestVerifyPayloads(t *testing.T) { + tessera, legacy := testLogs(t, false) + targets := []logTarget{tessera.target("tessera"), legacy.target("legacy")} + + payloads := [][]byte{ + normalizePayload(testPayload("both")), + normalizePayload(testPayload("missing")), + normalizePayload(testPayload("old")), + normalizePayload(append([]byte("\n "), testPayload("new")...)), + } + got, err := collect(t, context.Background(), targets, payloads) + if err != nil { + t.Fatalf("verifyPayloads: %v", err) + } + want := map[int]payloadResult{ + // The first log that contains a payload decides. + 0: {Index: 0, Verified: true, Log: "tessera"}, + 1: {Index: 1}, + 2: {Index: 2, Verified: true, Log: "legacy"}, + 3: {Index: 3, Verified: true, Log: "tessera"}, + } + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("verifyPayloads mismatch (-want +got):\n%s", diff) + } + + // Each log is consulted once for all payloads. + if n := tessera.checkpointRequests.Load(); n != 1 { + t.Errorf("Tessera checkpoint fetched %d times, want 1", n) + } + if n := legacy.checkpointRequests.Load(); n != 1 { + t.Errorf("legacy checkpoint fetched %d times, want 1", n) + } + if n := legacy.infoFileRequests.Load(); n != 1 { + t.Errorf("legacy info file fetched %d times, want 1", n) + } + if n := tessera.entryTileRequests.Load(); n != 2 { + t.Errorf("Tessera entry tiles fetched %d times, want 2", n) + } +} + +func TestVerifyPayloadsSkipsLaterLogsWhenAllFound(t *testing.T) { + tessera, legacy := testLogs(t, false) + targets := []logTarget{tessera.target("tessera"), legacy.target("legacy")} + + got, err := collect(t, context.Background(), targets, [][]byte{normalizePayload(testPayload("new"))}) + if err != nil { + t.Fatalf("verifyPayloads: %v", err) + } + if want := (payloadResult{Index: 0, Verified: true, Log: "tessera"}); got[0] != want { + t.Errorf("got %+v, want %+v", got[0], want) + } + if n := legacy.checkpointRequests.Load(); n != 0 { + t.Errorf("legacy checkpoint fetched %d times, want 0", n) + } + // "new" is in the last entry tile, so the first tile is never read. + if n := tessera.entryTileRequests.Load(); n != 1 { + t.Errorf("Tessera entry tiles fetched %d times, want 1", n) + } +} + +func TestVerifyPayloadsEntryTileErrorMatchesSinglePayloadRuns(t *testing.T) { + tessera, legacy := testLogs(t, false) + tessera.failFirstEntryTile = true // Tile 0 holds "both"; tile 1 holds "new". + targets := []logTarget{tessera.target("tessera"), legacy.target("legacy")} + + names := []string{"new", "both", "old", "missing"} + var payloads [][]byte + for _, name := range names { + payloads = append(payloads, normalizePayload(testPayload(name))) + } + got, err := collect(t, context.Background(), targets, payloads) + if err != nil { + t.Fatalf("verifyPayloads: %v", err) + } + want := map[int]payloadResult{ + // Found in tile 1, before the failing tile: kept. + 0: {Index: 0, Verified: true, Log: "tessera"}, + // Its Tessera search fails, so it is found in the next log. + 1: {Index: 1, Verified: true, Log: "legacy"}, + 2: {Index: 2, Verified: true, Log: "legacy"}, + 3: {Index: 3}, + } + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("verifyPayloads mismatch (-want +got):\n%s", diff) + } + + // Each result is what verifying that payload alone gives. + for i, p := range payloads { + single, err := collect(t, context.Background(), targets, [][]byte{p}) + if err != nil { + t.Fatalf("verifyPayloads(%s): %v", names[i], err) + } + r := single[0] + r.Index = i + if r != got[i] { + t.Errorf("%s: batch result %+v, single-payload result %+v", names[i], got[i], r) + } + } +} + +func TestVerifyPayloadsFailedProofIsFinal(t *testing.T) { + tessera, legacy := testLogs(t, true) + targets := []logTarget{tessera.target("tessera"), legacy.target("legacy")} + + got, err := collect(t, context.Background(), targets, [][]byte{ + normalizePayload(testPayload("both")), + normalizePayload(testPayload("old")), + }) + if err != nil { + t.Fatalf("verifyPayloads: %v", err) + } + // "both" is found in the Tessera log, whose proof fails; like a + // single-payload run, the legacy log is not tried for it. + if r := got[0]; r.Verified || r.Error == "" || r.Log != "" { + t.Errorf("payload 0 = %+v, want unverified with an error", r) + } + if want := (payloadResult{Index: 1, Verified: true, Log: "legacy"}); got[1] != want { + t.Errorf("payload 1 = %+v, want %+v", got[1], want) + } +} + +func TestVerifyPayloadsCancelled(t *testing.T) { + tessera, legacy := testLogs(t, false) + targets := []logTarget{tessera.target("tessera"), legacy.target("legacy")} + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + got, err := collect(t, ctx, targets, [][]byte{normalizePayload(testPayload("missing"))}) + if err == nil { + t.Errorf("verifyPayloads with cancelled context returned nil error") + } + if len(got) != 0 { + t.Errorf("verifyPayloads with cancelled context emitted %v, want nothing", got) + } +} + +func TestReadPayloads(t *testing.T) { + for _, tc := range []struct { + desc string + input string + want [][]byte + wantErr bool + }{ + {desc: "empty", input: ""}, + { + desc: "normalized, blank lines and unknown fields ignored", + input: "{\"payload\": \"a\\nb\"}\n\n{\"payload\": \" c\\n\\n\", \"label\": \"x\"}\n", + want: [][]byte{[]byte("a\nb\n"), []byte("c\n")}, + }, + {desc: "missing payload", input: "{\"payload\": \"a\"}\n{\"label\": \"x\"}\n", wantErr: true}, + {desc: "malformed", input: "{\"payload\": \"a\"}\nnot json\n", wantErr: true}, + {desc: "wrong type", input: "{\"payload\": 1}\n", wantErr: true}, + } { + t.Run(tc.desc, func(t *testing.T) { + got, err := readPayloads(strings.NewReader(tc.input)) + if (err != nil) != tc.wantErr { + t.Fatalf("readPayloads error = %v, wantErr %v", err, tc.wantErr) + } + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Errorf("readPayloads mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestRunBatch(t *testing.T) { + tessera, legacy := testLogs(t, false) + targets := []logTarget{tessera.target("tessera"), legacy.target("legacy")} + + dir := t.TempDir() + path := filepath.Join(dir, "payloads.jsonl") + var in bytes.Buffer + enc := json.NewEncoder(&in) + for _, name := range []string{"missing", "old", "new"} { + enc.Encode(map[string]string{"payload": string(testPayload(name))}) + } + if err := os.WriteFile(path, in.Bytes(), 0o600); err != nil { + t.Fatal(err) + } + + var out bytes.Buffer + if code := runBatch(context.Background(), targets, path, &out); code != 0 { + t.Fatalf("runBatch exit code = %d, want 0", code) + } + got := make(map[int]payloadResult) + for _, line := range strings.Split(strings.TrimSpace(out.String()), "\n") { + var r payloadResult + if err := json.Unmarshal([]byte(line), &r); err != nil { + t.Fatalf("output line %q: %v", line, err) + } + got[r.Index] = r + } + want := map[int]payloadResult{ + 0: {Index: 0}, + 1: {Index: 1, Verified: true, Log: "legacy"}, + 2: {Index: 2, Verified: true, Log: "tessera"}, + } + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("runBatch results mismatch (-want +got):\n%s", diff) + } + // Optional fields are omitted when empty. + if !strings.Contains(out.String(), `{"index":0,"verified":false}`) { + t.Errorf("runBatch output %q lacks compact not-found result", out.String()) + } + + malformed := filepath.Join(dir, "malformed.jsonl") + os.WriteFile(malformed, []byte("nope\n"), 0o600) + for _, p := range []string{malformed, filepath.Join(dir, "absent.jsonl")} { + out.Reset() + if code := runBatch(context.Background(), targets, p, &out); code != 1 || out.Len() != 0 { + t.Errorf("runBatch(%s) = %d with output %q, want 1 and no output", filepath.Base(p), code, out.String()) + } + } +} diff --git a/verifier_tools/verify/internal/tiles/reader.go b/verifier_tools/verify/internal/tiles/reader.go index adeb597..1818f0e 100644 --- a/verifier_tools/verify/internal/tiles/reader.go +++ b/verifier_tools/verify/internal/tiles/reader.go @@ -31,6 +31,10 @@ type HashReader struct { TileHeight int TreeSize int64 IsTessera bool + // TileCache, if non-nil, keeps downloaded hash tiles (path -> content) + // across ReadHashes calls, e.g. for several proofs against the same tree. + // Not safe for concurrent use. If nil, tiles are kept for one call only. + TileCache map[string][]byte } // Domain separation prefix for Merkle tree hashing with second preimage @@ -42,7 +46,10 @@ const ( // ReadHashes implements tlog.HashReader's ReadHashes. // See: https://pkg.go.dev/golang.org/x/mod/sumdb/tlog#HashReader. func (h HashReader) ReadHashes(indices []int64) ([]tlog.Hash, error) { - tiles := make(map[string][]byte) // cache tile path -> content + tiles := h.TileCache // cache tile path -> content + if tiles == nil { + tiles = make(map[string][]byte) + } hashes := make([]tlog.Hash, 0, len(indices)) for _, index := range indices { // A tlog index is a pointer to a hash at a given level in the tree. @@ -595,16 +602,48 @@ func FetchAllTesseraEntries(ctx context.Context, logBaseURL string, treeSize int // in the log. It returns the highest index if the payload appears more than once. // Returns (index, true, nil) if found, (-1, false, nil) if not found. func TesseraFindPayloadIndex(logBaseURL string, treeSize int64, targetPayload []byte) (int64, bool, error) { - if treeSize <= 0 { + found, err := TesseraFindPayloadIndices(logBaseURL, treeSize, [][]byte{targetPayload}) + if err != nil { + return -1, false, err + } + idx, ok := found[string(bytes.TrimSpace(targetPayload))] + if !ok { return -1, false, nil } + return idx, true, nil +} + +// TesseraFindPayloadIndices searches the entry tiles once for several +// payloads. It returns a map from each found payload, with surrounding +// whitespace trimmed, to its highest index in the log; payloads that are not +// found are absent. +// +// Tiles are read from latest to oldest, as in TesseraFindPayloadIndex, and +// the search stops as soon as every payload has been found. A payload that is +// not in the log therefore makes the search read every tile, but only once for +// all payloads. +// +// If a tile cannot be read, the search stops and returns the payloads found +// so far together with the error. Those matches are still exact: they were +// found in newer tiles, so a search for one payload alone would have stopped +// there too, without reaching the failing tile. Only the payloads not yet +// found are affected, as they would be when searched for alone. +func TesseraFindPayloadIndices(logBaseURL string, treeSize int64, targetPayloads [][]byte) (map[string]int64, error) { + found := make(map[string]int64) + wanted := make(map[string]bool, len(targetPayloads)) + for _, p := range targetPayloads { + wanted[string(bytes.TrimSpace(p))] = true + } + if treeSize <= 0 || len(wanted) == 0 { + return found, nil + } numTiles := (treeSize + 255) / 256 - target := bytes.TrimSpace(targetPayload) // Search in reverse (latest to oldest) to find recent packages with typically - // fewer HTTP requests on cold caches. - for tileN := numTiles - 1; tileN >= 0; tileN-- { + // fewer HTTP requests on cold caches. The first match seen for a payload is + // therefore its highest index. + for tileN := numTiles - 1; tileN >= 0 && len(found) < len(wanted); tileN-- { w := 256 if (tileN+1)*256 > treeSize { w = int(treeSize - tileN*256) @@ -612,20 +651,24 @@ func TesseraFindPayloadIndex(logBaseURL string, treeSize int64, targetPayload [] b, err := readCachedEntryTile(logBaseURL, tileN, w) if err != nil { - return -1, false, fmt.Errorf("failed to fetch entry tile %d (width %d): %w", tileN, w, err) + return found, fmt.Errorf("failed to fetch entry tile %d (width %d): %w", tileN, w, err) } entries, err := ParseEntryBundle(b) if err != nil { - return -1, false, fmt.Errorf("failed to parse entry tile %d: %w", tileN, err) + return found, fmt.Errorf("failed to parse entry tile %d: %w", tileN, err) } for idx := len(entries) - 1; idx >= 0; idx-- { - if bytes.Equal(bytes.TrimSpace(entries[idx]), target) { - return tileN*256 + int64(idx), true, nil + key := string(bytes.TrimSpace(entries[idx])) + if !wanted[key] { + continue + } + if _, seen := found[key]; !seen { + found[key] = tileN*256 + int64(idx) } } } - return -1, false, nil + return found, nil } diff --git a/verifier_tools/verify/internal/tiles/reader_test.go b/verifier_tools/verify/internal/tiles/reader_test.go index 386aa40..8458c9e 100644 --- a/verifier_tools/verify/internal/tiles/reader_test.go +++ b/verifier_tools/verify/internal/tiles/reader_test.go @@ -416,6 +416,121 @@ func TestTesseraFindPayloadIndex(t *testing.T) { } } +func TestTesseraFindPayloadIndices(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("HOME", tempDir) + t.Setenv("XDG_CACHE_HOME", tempDir) + + entry := func(i int) []byte { return []byte(fmt.Sprintf("hash_%d\nhash_desc\npackage_%d\n%d\n", i, i, i)) } + dup := []byte("hash_dup\nhash_desc\npackage_dup\n1\n") + + // Tile 0 holds indices 0..255 (dup at 5 and 250); tile 1 holds 256..258 + // (dup at 257). + var tile0Entries [][]byte + for i := 0; i < 256; i++ { + if i == 5 || i == 250 { + tile0Entries = append(tile0Entries, dup) + } else { + tile0Entries = append(tile0Entries, entry(i)) + } + } + tile0Data := createTestEntryBundle(tile0Entries) + tile1Data := createTestEntryBundle([][]byte{entry(256), dup, entry(258)}) + + var tile0Requests, tile1Requests atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/tile/entries/000": + tile0Requests.Add(1) + w.Write(tile0Data) + case "/tile/entries/001.p/3": + tile1Requests.Add(1) + w.Write(tile1Data) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + // All payloads in the latest tile: tile 0 is never read. Surrounding + // whitespace is ignored, as in TesseraFindPayloadIndex. + got, err := TesseraFindPayloadIndices(server.URL, 259, [][]byte{entry(258), append([]byte(" "), dup...)}) + if err != nil { + t.Fatalf("TesseraFindPayloadIndices error: %v", err) + } + want := map[string]int64{string(bytes.TrimSpace(entry(258))): 258, string(bytes.TrimSpace(dup)): 257} + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("TesseraFindPayloadIndices mismatch (-want +got):\n%s", diff) + } + if tile0Requests.Load() != 0 || tile1Requests.Load() != 1 { + t.Errorf("requests (tile0, tile1) = (%d, %d), want (0, 1)", tile0Requests.Load(), tile1Requests.Load()) + } + + // Payloads across both tiles plus one that is absent: every tile is read, + // each only once, and the absent payload is omitted. + got, err = TesseraFindPayloadIndices(server.URL, 259, [][]byte{entry(3), entry(256), []byte("absent"), entry(3)}) + if err != nil { + t.Fatalf("TesseraFindPayloadIndices error: %v", err) + } + want = map[string]int64{string(bytes.TrimSpace(entry(3))): 3, string(bytes.TrimSpace(entry(256))): 256} + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("TesseraFindPayloadIndices mismatch (-want +got):\n%s", diff) + } + if tile0Requests.Load() != 1 || tile1Requests.Load() != 1 { + t.Errorf("requests (tile0, tile1) = (%d, %d), want (1, 1)", tile0Requests.Load(), tile1Requests.Load()) + } + + // No payloads, or an empty tree: nothing to search. + for _, tc := range []struct { + size int64 + payloads [][]byte + }{{259, nil}, {0, [][]byte{entry(3)}}} { + got, err := TesseraFindPayloadIndices(server.URL, tc.size, tc.payloads) + if err != nil || len(got) != 0 { + t.Errorf("TesseraFindPayloadIndices(size=%d, %d payloads) = (%v, %v), want empty map", tc.size, len(tc.payloads), got, err) + } + } +} + +func TestTesseraFindPayloadIndicesKeepsMatchesOnTileError(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("HOME", tempDir) + t.Setenv("XDG_CACHE_HOME", tempDir) + + entry := func(i int) []byte { return []byte(fmt.Sprintf("hash_%d\nhash_desc\npackage_%d\n%d\n", i, i, i)) } + tile1Data := createTestEntryBundle([][]byte{entry(256), entry(257)}) + + // Tile 0 (the oldest) fails; tile 1 is served. + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/tile/entries/001.p/2": + w.Write(tile1Data) + case "/tile/entries/000": + http.Error(w, "boom", http.StatusInternalServerError) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + // entry(3) would be in tile 0, so the search reaches it and fails. The + // match from tile 1 is exact and must be kept alongside the error. + got, err := TesseraFindPayloadIndices(server.URL, 258, [][]byte{entry(257), entry(3)}) + if err == nil { + t.Fatalf("TesseraFindPayloadIndices error = nil, want tile 0 failure") + } + want := map[string]int64{string(bytes.TrimSpace(entry(257))): 257} + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("TesseraFindPayloadIndices partial result mismatch (-want +got):\n%s", diff) + } + + // A single-payload search that is satisfied by tile 1 never reaches the + // failing tile. + if idx, found, err := TesseraFindPayloadIndex(server.URL, 258, entry(257)); err != nil || !found || idx != 257 { + t.Errorf("TesseraFindPayloadIndex = (%d, %v, %v), want (257, true, nil)", idx, found, err) + } +} + func createTestEntryBundle(entries [][]byte) []byte { var buf bytes.Buffer for _, entry := range entries {