diff --git a/tools/generate-driver-release-manifest.py b/tools/generate-driver-release-manifest.py index dad99e06..12edcf26 100644 --- a/tools/generate-driver-release-manifest.py +++ b/tools/generate-driver-release-manifest.py @@ -9,9 +9,13 @@ import shutil import subprocess import sys import tempfile +import time from pathlib import Path +METADATA_PROBE_ATTEMPTS = 3 + + def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--assets-dir", required=True, help="driver release staging dir that contains standalone driver assets") @@ -156,17 +160,22 @@ def resolve_host_platform(): return f"{goos}/{goarch}" -def probe_agent_metadata(asset_path: Path): +def probe_agent_metadata_once(asset_path: Path): if os.name != "nt": os.chmod(asset_path, asset_path.stat().st_mode | 0o111) - proc = subprocess.run( - [str(asset_path)], - input='{"id":1,"method":"metadata"}\n', - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - timeout=30, - ) + try: + proc = subprocess.run( + [str(asset_path)], + input='{"id":1,"method":"metadata"}\n', + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=30, + ) + except subprocess.TimeoutExpired as exc: + raise RuntimeError(f"{asset_path.name}: metadata probe timed out after 30 seconds") from exc + except OSError as exc: + raise RuntimeError(f"{asset_path.name}: metadata probe failed to start: {exc}") from exc if proc.returncode != 0: detail = proc.stderr.strip() or f"exit code {proc.returncode}" raise RuntimeError(f"{asset_path.name}: metadata probe failed: {detail}") @@ -188,6 +197,64 @@ def probe_agent_metadata(asset_path: Path): return driver_type, revision +def probe_upx_unpacked_metadata(asset_path: Path): + upx = shutil.which("upx") + if not upx: + return None + + tested = subprocess.run( + [upx, "-t", str(asset_path)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + if tested.returncode != 0: + return None + + with tempfile.TemporaryDirectory(prefix="gonavi-driver-agent-unpacked-") as tmp: + unpacked_path = Path(tmp) / asset_path.name + shutil.copy2(asset_path, unpacked_path) + unpacked = subprocess.run( + [upx, "-d", str(unpacked_path)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + if unpacked.returncode != 0: + detail = unpacked.stderr.strip() or f"exit code {unpacked.returncode}" + print(f"{asset_path.name}: UPX unpack fallback failed: {detail}", file=sys.stderr) + return None + try: + metadata = probe_agent_metadata_once(unpacked_path) + except RuntimeError as exc: + print(f"{asset_path.name}: unpacked metadata probe failed: {exc}", file=sys.stderr) + return None + + print(f"{asset_path.name}: metadata probe recovered from an unpacked UPX copy", file=sys.stderr) + return metadata + + +def probe_agent_metadata(asset_path: Path): + last_error = None + for attempt in range(1, METADATA_PROBE_ATTEMPTS + 1): + try: + return probe_agent_metadata_once(asset_path) + except RuntimeError as exc: + last_error = exc + print( + f"{asset_path.name}: metadata probe attempt " + f"{attempt}/{METADATA_PROBE_ATTEMPTS} failed: {exc}", + file=sys.stderr, + ) + if attempt < METADATA_PROBE_ATTEMPTS: + time.sleep(attempt) + + unpacked_metadata = probe_upx_unpacked_metadata(asset_path) + if unpacked_metadata is not None: + return unpacked_metadata + raise last_error + + def load_asset_provenance(paths): entries = {} for raw_path in paths: diff --git a/tools/generate-driver-release-manifest.test.py b/tools/generate-driver-release-manifest.test.py index c074ce7d..f336ff5a 100644 --- a/tools/generate-driver-release-manifest.test.py +++ b/tools/generate-driver-release-manifest.test.py @@ -105,7 +105,7 @@ class GenerateDriverReleaseManifestTest(unittest.TestCase): ).stdout.strip() return goos, goarch - def _build_metadata_agent(self, output: Path, revision: str): + def _build_metadata_agent(self, output: Path, revision: str, driver: str = "clickhouse"): source = output.parent / "metadata-agent.go" source.write_text( "package main\n" @@ -113,6 +113,35 @@ class GenerateDriverReleaseManifestTest(unittest.TestCase): "func main() {\n" " scanner := bufio.NewScanner(os.Stdin)\n" " for scanner.Scan() {\n" + f" fmt.Println(`{{\"id\":1,\"success\":true,\"data\":{{\"driverType\":\"{driver}\",\"agentRevision\":\"{revision}\",\"protocolSchema\":\"json-lines-v1\"}}}}`)\n" + " }\n" + "}\n", + encoding="utf-8", + ) + output.unlink(missing_ok=True) + proc = subprocess.run( + ["go", "build", "-o", str(output), str(source)], + cwd=ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + self.assertEqual(proc.returncode, 0, proc.stderr) + + def _build_transiently_failing_metadata_agent(self, output: Path, revision: str): + source = output.parent / "transient-metadata-agent.go" + marker = output.parent / "metadata-agent-first-probe" + source.write_text( + "package main\n" + "import (\"bufio\"; \"fmt\"; \"os\")\n" + "func main() {\n" + f" marker := {json.dumps(str(marker))}\n" + " if _, err := os.Stat(marker); os.IsNotExist(err) {\n" + " if err := os.WriteFile(marker, []byte(\"attempted\"), 0600); err != nil { panic(err) }\n" + " os.Exit(127)\n" + " }\n" + " scanner := bufio.NewScanner(os.Stdin)\n" + " for scanner.Scan() {\n" f" fmt.Println(`{{\"id\":1,\"success\":true,\"data\":{{\"driverType\":\"clickhouse\",\"agentRevision\":\"{revision}\",\"protocolSchema\":\"json-lines-v1\"}}}}`)\n" " }\n" "}\n", @@ -157,6 +186,48 @@ class GenerateDriverReleaseManifestTest(unittest.TestCase): self.assertIn("src-stale-agent", proc.stderr) self.assertFalse(output.exists()) + def test_retries_native_metadata_probe_after_transient_exit_127(self): + goos, goarch = self._host_platform() + platform = f"{goos}/{goarch}" + extension = ".exe" if goos == "windows" else "" + revision = "src-transient-probe" + with tempfile.TemporaryDirectory(prefix="gonavi-release-manifest-transient-agent-") as tmp: + tmpdir = Path(tmp) + assets_dir = tmpdir / "drivers" + assets_dir.mkdir(parents=True) + asset = assets_dir / f"clickhouse-driver-agent-{goos}-{goarch}{extension}" + self._build_transiently_failing_metadata_agent(asset, revision) + revision_file = tmpdir / "driver_agent_revisions_gen.go" + revision_file.write_text( + "package db\n" + "var revisions = map[string]string{\n" + f' "clickhouse": "{revision}",\n' + "}\n", + encoding="utf-8", + ) + provenance = tmpdir / "build-provenance.json" + + proc = subprocess.run( + [ + sys.executable, + str(SCRIPT), + "--assets-dir", + str(assets_dir), + "--provenance-output", + str(provenance), + "--revision-file", + str(revision_file), + ], + cwd=ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + self.assertEqual(proc.returncode, 0, proc.stderr) + payload = json.loads(provenance.read_text(encoding="utf-8")) + self.assertEqual(payload["assets"][asset.name]["revision"], revision) + def test_rejects_native_asset_when_sha_bound_provenance_disagrees_with_binary(self): goos, goarch = self._host_platform() platform = f"{goos}/{goarch}" @@ -448,8 +519,8 @@ class GenerateDriverReleaseManifestTest(unittest.TestCase): for path, (driver, platform, content) in fixtures.items(): path.write_bytes(content) os.chmod(path, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR) - if platform == host_platform and driver == "clickhouse": - self._build_metadata_agent(path, revision_by_platform_driver[(platform, driver)]) + if platform == host_platform: + self._build_metadata_agent(path, revision_by_platform_driver[(platform, driver)], driver) provenance_assets = {} for path, (driver, platform, _) in fixtures.items():