mirror of
https://github.com/NixOS/nixpkgs.git
synced 2024-11-21 22:43:01 +00:00
293 lines
10 KiB
Python
293 lines
10 KiB
Python
import functools
|
|
import hashlib
|
|
import json
|
|
import multiprocessing as mp
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tomllib
|
|
from pathlib import Path
|
|
from typing import Any, TypedDict, cast
|
|
|
|
import requests
|
|
from requests.adapters import HTTPAdapter, Retry
|
|
|
|
eprint = functools.partial(print, file=sys.stderr)
|
|
|
|
|
|
def load_toml(path: Path) -> dict[str, Any]:
|
|
with open(path, "rb") as f:
|
|
return tomllib.load(f)
|
|
|
|
|
|
def download_file_with_checksum(url: str, destination_path: Path) -> str:
|
|
retries = Retry(
|
|
total=5,
|
|
backoff_factor=0.5,
|
|
status_forcelist=[500, 502, 503, 504]
|
|
)
|
|
session = requests.Session()
|
|
session.mount('http://', HTTPAdapter(max_retries=retries))
|
|
session.mount('https://', HTTPAdapter(max_retries=retries))
|
|
|
|
sha256_hash = hashlib.sha256()
|
|
with session.get(url, stream=True) as response:
|
|
if not response.ok:
|
|
raise Exception(f"Failed to fetch file from {url}. Status code: {response.status_code}")
|
|
with open(destination_path, "wb") as file:
|
|
for chunk in response.iter_content(1024): # Download in chunks
|
|
if chunk: # Filter out keep-alive chunks
|
|
file.write(chunk)
|
|
sha256_hash.update(chunk)
|
|
|
|
# Compute the final checksum
|
|
checksum = sha256_hash.hexdigest()
|
|
return checksum
|
|
|
|
|
|
def get_download_url_for_tarball(pkg: dict[str, Any]) -> str:
|
|
# TODO: support other registries
|
|
# maybe fetch config.json from the registry root and get the dl key
|
|
# See: https://doc.rust-lang.org/cargo/reference/registry-index.html#index-configuration
|
|
if pkg["source"] != "registry+https://github.com/rust-lang/crates.io-index":
|
|
raise Exception("Only the default crates.io registry is supported.")
|
|
|
|
return f"https://crates.io/api/v1/crates/{pkg["name"]}/{pkg["version"]}/download"
|
|
|
|
|
|
def download_tarball(pkg: dict[str, Any], out_dir: Path) -> None:
|
|
|
|
url = get_download_url_for_tarball(pkg)
|
|
filename = f"{pkg["name"]}-{pkg["version"]}.tar.gz"
|
|
|
|
# TODO: allow legacy checksum specification, see importCargoLock for example
|
|
# also, don't forget about the other usage of the checksum
|
|
expected_checksum = pkg["checksum"]
|
|
|
|
tarball_out_dir = out_dir / "tarballs" / filename
|
|
eprint(f"Fetching {url} -> tarballs/{filename}")
|
|
|
|
calculated_checksum = download_file_with_checksum(url, tarball_out_dir)
|
|
|
|
if calculated_checksum != expected_checksum:
|
|
raise Exception(f"Hash mismatch! File fetched from {url} had checksum {calculated_checksum}, expected {expected_checksum}.")
|
|
|
|
|
|
def download_git_tree(url: str, git_sha_rev: str, out_dir: Path) -> None:
|
|
|
|
tree_out_dir = out_dir / "git" / git_sha_rev
|
|
eprint(f"Fetching {url}#{git_sha_rev} -> git/{git_sha_rev}")
|
|
|
|
cmd = ["nix-prefetch-git", "--builder", "--quiet", "--fetch-submodules", "--url", url, "--rev", git_sha_rev, "--out", str(tree_out_dir)]
|
|
subprocess.check_output(cmd)
|
|
|
|
|
|
GIT_SOURCE_REGEX = re.compile("git\\+(?P<url>[^?]+)(\\?(?P<type>rev|tag|branch)=(?P<value>.*))?#(?P<git_sha_rev>.*)")
|
|
|
|
|
|
class GitSourceInfo(TypedDict):
|
|
url: str
|
|
type: str | None
|
|
value: str | None
|
|
git_sha_rev: str
|
|
|
|
|
|
def parse_git_source(source: str) -> GitSourceInfo:
|
|
match = GIT_SOURCE_REGEX.match(source)
|
|
if match is None:
|
|
raise Exception(f"Unable to process git source: {source}.")
|
|
return cast(GitSourceInfo, match.groupdict(default=None))
|
|
|
|
|
|
def create_vendor_staging(lockfile_path: Path, out_dir: Path) -> None:
|
|
cargo_toml = load_toml(lockfile_path)
|
|
|
|
git_packages: list[dict[str, Any]] = []
|
|
registry_packages: list[dict[str, Any]] = []
|
|
|
|
for pkg in cargo_toml["package"]:
|
|
# ignore local dependenices
|
|
if "source" not in pkg.keys():
|
|
eprint(f"Skipping local dependency: {pkg["name"]}")
|
|
continue
|
|
source = pkg["source"]
|
|
|
|
if source.startswith("git+"):
|
|
git_packages.append(pkg)
|
|
elif source.startswith("registry+"):
|
|
registry_packages.append(pkg)
|
|
else:
|
|
raise Exception(f"Can't process source: {source}.")
|
|
|
|
git_sha_rev_to_url: dict[str, str] = {}
|
|
for pkg in git_packages:
|
|
source_info = parse_git_source(pkg["source"])
|
|
git_sha_rev_to_url[source_info["git_sha_rev"]] = source_info["url"]
|
|
|
|
out_dir.mkdir(exist_ok=True)
|
|
shutil.copy(lockfile_path, out_dir / "Cargo.lock")
|
|
|
|
# fetch git trees sequentially, since fetching concurrently leads to flaky behaviour
|
|
if len(git_packages) != 0:
|
|
(out_dir / "git").mkdir()
|
|
for git_sha_rev, url in git_sha_rev_to_url.items():
|
|
download_git_tree(url, git_sha_rev, out_dir)
|
|
|
|
# run tarball download jobs in parallel, with at most 5 concurrent download jobs
|
|
with mp.Pool(min(5, mp.cpu_count())) as pool:
|
|
if len(registry_packages) != 0:
|
|
(out_dir / "tarballs").mkdir()
|
|
tarball_args_gen = ((pkg, out_dir) for pkg in registry_packages)
|
|
pool.starmap(download_tarball, tarball_args_gen)
|
|
|
|
|
|
def get_manifest_metadata(manifest_path: Path) -> dict[str, Any]:
|
|
cmd = ["cargo", "metadata", "--format-version", "1", "--no-deps", "--manifest-path", str(manifest_path)]
|
|
output = subprocess.check_output(cmd)
|
|
return json.loads(output)
|
|
|
|
|
|
def try_get_crate_manifest_path_from_mainfest_path(manifest_path: Path, crate_name: str) -> Path | None:
|
|
metadata = get_manifest_metadata(manifest_path)
|
|
|
|
for pkg in metadata["packages"]:
|
|
if pkg["name"] == crate_name:
|
|
return Path(pkg["manifest_path"])
|
|
|
|
return None
|
|
|
|
|
|
def find_crate_manifest_in_tree(tree: Path, crate_name: str) -> Path:
|
|
# in some cases Cargo.toml is not located at the top level, so we also look at subdirectories
|
|
manifest_paths = tree.glob("**/Cargo.toml")
|
|
|
|
for manifest_path in manifest_paths:
|
|
res = try_get_crate_manifest_path_from_mainfest_path(manifest_path, crate_name)
|
|
if res is not None:
|
|
return res
|
|
|
|
raise Exception(f"Couldn't find manifest for crate {crate_name} inside {tree}.")
|
|
|
|
|
|
def copy_and_patch_git_crate_subtree(git_tree: Path, crate_name: str, crate_out_dir: Path) -> None:
|
|
crate_manifest_path = find_crate_manifest_in_tree(git_tree, crate_name)
|
|
crate_tree = crate_manifest_path.parent
|
|
|
|
eprint(f"Copying to {crate_out_dir}")
|
|
shutil.copytree(crate_tree, crate_out_dir)
|
|
crate_out_dir.chmod(0o755)
|
|
|
|
with open(crate_manifest_path, "r") as f:
|
|
manifest_data = f.read()
|
|
|
|
if "workspace" in manifest_data:
|
|
crate_manifest_metadata = get_manifest_metadata(crate_manifest_path)
|
|
workspace_root = Path(crate_manifest_metadata["workspace_root"])
|
|
|
|
root_manifest_path = workspace_root / "Cargo.toml"
|
|
manifest_path = crate_out_dir / "Cargo.toml"
|
|
|
|
manifest_path.chmod(0o644)
|
|
eprint(f"Patching {manifest_path}")
|
|
|
|
cmd = ["replace-workspace-values", str(manifest_path), str(root_manifest_path)]
|
|
subprocess.check_output(cmd)
|
|
|
|
|
|
def extract_crate_tarball_contents(tarball_path: Path, crate_out_dir: Path) -> None:
|
|
eprint(f"Unpacking to {crate_out_dir}")
|
|
crate_out_dir.mkdir()
|
|
cmd = ["tar", "xf", str(tarball_path), "-C", str(crate_out_dir), "--strip-components=1"]
|
|
subprocess.check_output(cmd)
|
|
|
|
|
|
def create_vendor(vendor_staging_dir: Path, out_dir: Path) -> None:
|
|
lockfile_path = vendor_staging_dir / "Cargo.lock"
|
|
out_dir.mkdir(exist_ok=True)
|
|
shutil.copy(lockfile_path, out_dir / "Cargo.lock")
|
|
|
|
cargo_toml = load_toml(lockfile_path)
|
|
|
|
config_lines = [
|
|
'[source.vendored-sources]',
|
|
'directory = "@vendor@"',
|
|
'[source.crates-io]',
|
|
'replace-with = "vendored-sources"',
|
|
]
|
|
|
|
seen_source_keys = set()
|
|
for pkg in cargo_toml["package"]:
|
|
|
|
# ignore local dependenices
|
|
if "source" not in pkg.keys():
|
|
continue
|
|
|
|
source: str = pkg["source"]
|
|
|
|
dir_name = f"{pkg["name"]}-{pkg["version"]}"
|
|
crate_out_dir = out_dir / dir_name
|
|
|
|
if source.startswith("git+"):
|
|
|
|
source_info = parse_git_source(pkg["source"])
|
|
git_sha_rev = source_info["git_sha_rev"]
|
|
git_tree = vendor_staging_dir / "git" / git_sha_rev
|
|
|
|
copy_and_patch_git_crate_subtree(git_tree, pkg["name"], crate_out_dir)
|
|
|
|
# git based crates allow having no checksum information
|
|
with open(crate_out_dir / ".cargo-checksum.json", "w") as f:
|
|
json.dump({"files": {}}, f)
|
|
|
|
source_key = source[0:source.find("#")]
|
|
|
|
if source_key in seen_source_keys:
|
|
continue
|
|
|
|
seen_source_keys.add(source_key)
|
|
|
|
config_lines.append(f'[source."{source_key}"]')
|
|
config_lines.append(f'git = "{source_info["url"]}"')
|
|
if source_info["type"] is not None:
|
|
config_lines.append(f'{source_info["type"]} = "{source_info["value"]}"')
|
|
config_lines.append('replace-with = "vendored-sources"')
|
|
|
|
elif source.startswith("registry+"):
|
|
|
|
filename = f"{pkg["name"]}-{pkg["version"]}.tar.gz"
|
|
tarball_path = vendor_staging_dir / "tarballs" / filename
|
|
|
|
extract_crate_tarball_contents(tarball_path, crate_out_dir)
|
|
|
|
# non-git based crates need the package checksum at minimum
|
|
with open(crate_out_dir / ".cargo-checksum.json", "w") as f:
|
|
json.dump({"files": {}, "package": pkg["checksum"]}, f)
|
|
|
|
else:
|
|
raise Exception(f"Can't process source: {source}.")
|
|
|
|
(out_dir / ".cargo").mkdir()
|
|
with open(out_dir / ".cargo" / "config.toml", "w") as config_file:
|
|
config_file.writelines(line + "\n" for line in config_lines)
|
|
|
|
|
|
def main() -> None:
|
|
subcommand = sys.argv[1]
|
|
|
|
subcommand_func_dict = {
|
|
"create-vendor-staging": lambda: create_vendor_staging(lockfile_path=Path(sys.argv[2]), out_dir=Path(sys.argv[3])),
|
|
"create-vendor": lambda: create_vendor(vendor_staging_dir=Path(sys.argv[2]), out_dir=Path(sys.argv[3]))
|
|
}
|
|
|
|
subcommand_func = subcommand_func_dict.get(subcommand)
|
|
|
|
if subcommand_func is None:
|
|
raise Exception(f"Unknown subcommand: '{subcommand}'. Must be one of {list(subcommand_func_dict.keys())}")
|
|
|
|
subcommand_func()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|