From 2512f51d97839b4a4ed6b7a4445b4bfc6285a1c8 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 20 Jul 2026 21:10:19 -0700 Subject: [PATCH 1/2] feat(dstackup): support configurable release API --- docs/testing-dstackup-image-pull.md | 39 ++++++ dstack/crates/dstackup/src/cli.rs | 6 + dstack/crates/dstackup/src/image.rs | 67 +++++++-- dstack/crates/dstackup/src/install.rs | 3 +- dstack/crates/dstackup/src/main.rs | 4 +- tools/mock-github-releases.py | 195 ++++++++++++++++++++++++++ 6 files changed, 297 insertions(+), 17 deletions(-) create mode 100644 docs/testing-dstackup-image-pull.md create mode 100755 tools/mock-github-releases.py diff --git a/docs/testing-dstackup-image-pull.md b/docs/testing-dstackup-image-pull.md new file mode 100644 index 000000000..85186967c --- /dev/null +++ b/docs/testing-dstackup-image-pull.md @@ -0,0 +1,39 @@ +# Test `dstackup image pull` with a local release API + +`--release-api-base-url` accepts plain HTTP URLs, including localhost. A small +GitHub Releases API overlay is included for local testing. Locally published +releases shadow GitHub; API requests that are not present locally are relayed to +`https://api.github.com`. + +Start it: + +```bash +./tools/mock-github-releases.py --port 8000 +``` + +Publish a release and copy a local tarball into the mock asset store. The mock +calculates and publishes its SHA-256 digest automatically: + +```bash +curl -fsS -X POST \ + http://localhost:8000/__admin/repos/Dstack-TEE/dstack/releases \ + -H 'content-type: application/json' \ + -d '{ + "tag_name": "guest-os-v99.0.0", + "assets": [{"local_path": "/tmp/dstack-99.0.0.tar.gz"}] + }' +``` + +Pull the locally published latest release: + +```bash +sudo dstackup image pull \ + --release-api-base-url http://localhost:8000/repos \ + --force +``` + +A release can instead reference an already hosted asset by supplying `name`, +`browser_download_url`, and optionally `digest` in the asset object. State is +in memory and resets when the mock exits; copied assets remain in `--asset-dir`. +GitHub API rate limits still apply to relayed requests. Set an `Authorization` +header on a direct request to the mock when testing authenticated relays. diff --git a/dstack/crates/dstackup/src/cli.rs b/dstack/crates/dstackup/src/cli.rs index 4d1b74689..0d4107c9e 100644 --- a/dstack/crates/dstackup/src/cli.rs +++ b/dstack/crates/dstackup/src/cli.rs @@ -12,6 +12,7 @@ pub(crate) const DEFAULT_AUTH_BIN: &str = "dstack-auth"; pub(crate) const DEFAULT_SUPERVISOR_BIN: &str = "supervisor"; pub(crate) const DEFAULT_SOURCE_REPO: &str = "https://github.com/Dstack-TEE/dstack"; pub(crate) const DEFAULT_SOURCE_REF: &str = "master"; +pub(crate) const DEFAULT_RELEASE_API_BASE_URL: &str = "https://api.github.com/repos"; #[derive(Parser)] #[command(name = "dstackup", version, about = "set up and manage a dstack host")] @@ -21,6 +22,11 @@ pub(crate) struct Cli { #[arg(long, global = true)] pub(crate) host: Option, + /// Base URL of the GitHub-compatible releases API. The repository path is + /// appended automatically (useful for an API proxy or a test server). + #[arg(long, global = true, default_value = DEFAULT_RELEASE_API_BASE_URL)] + pub(crate) release_api_base_url: String, + #[command(subcommand)] pub(crate) command: Command, } diff --git a/dstack/crates/dstackup/src/image.rs b/dstack/crates/dstackup/src/image.rs index 1874d31b9..191a25e7a 100644 --- a/dstack/crates/dstackup/src/image.rs +++ b/dstack/crates/dstackup/src/image.rs @@ -64,7 +64,7 @@ struct PullSpec { gpu: bool, } -pub(crate) async fn cmd_image(cmd: ImageCmd) -> Result<()> { +pub(crate) async fn cmd_image(cmd: ImageCmd, release_api_base_url: &str) -> Result<()> { match cmd { ImageCmd::Pull { version, @@ -75,7 +75,15 @@ pub(crate) async fn cmd_image(cmd: ImageCmd) -> Result<()> { } => { let image_dir = loc.dir(); validate_image_dir(&image_dir)?; - pull(version.as_deref(), gpu, &image_dir, force, insecure).await?; + pull( + version.as_deref(), + gpu, + &image_dir, + force, + insecure, + release_api_base_url, + ) + .await?; Ok(()) } ImageCmd::List { loc } => { @@ -98,6 +106,7 @@ pub(crate) async fn pull( image_dir: &str, force: bool, insecure: bool, + release_api_base_url: &str, ) -> Result { println!( "dstackup image pull — {} image", @@ -107,7 +116,7 @@ pub(crate) async fn pull( "unified" } ); - let release = fetch_release(version).await?; + let release = fetch_release(version, release_api_base_url).await?; let asset = pick_asset(&release.assets, gpu).with_context(|| { format!( @@ -376,6 +385,7 @@ pub(crate) async fn resolve_or_pull_image( requested: Option<&str>, require: bool, required_files: &[&str], + release_api_base_url: &str, ) -> Result> { if let Some(name) = requested { if !valid_image_name(name) { @@ -391,7 +401,15 @@ pub(crate) async fn resolve_or_pull_image( } if let Some(spec) = pull_spec(name) { println!(" [..] image {name} not found locally; downloading it"); - let pulled = pull(Some(&spec.version), spec.gpu, image_dir, false, false).await?; + let pulled = pull( + Some(&spec.version), + spec.gpu, + image_dir, + false, + false, + release_api_base_url, + ) + .await?; ensure_image_has_required_files(image_dir, &pulled, required_files)?; return Ok(Some(pulled)); } @@ -459,7 +477,7 @@ pub(crate) async fn resolve_or_pull_image( } else { println!(" [..] no local guest image found; downloading the latest cpu image"); } - let pulled = pull(None, false, image_dir, false, false).await?; + let pulled = pull(None, false, image_dir, false, false, release_api_base_url).await?; if Path::new(image_dir) .join(&pulled) @@ -599,10 +617,10 @@ fn missing_named_image_message(image_dir: &str, name: &str) -> String { /// released from this monorepo. Do not probe the new repository first for old /// versions: the version boundary is authoritative and avoids redundant or /// misleading requests. -async fn fetch_release(version: Option<&str>) -> Result { +async fn fetch_release(version: Option<&str>, release_api_base_url: &str) -> Result { let client = reqwest::Client::new(); if let Some(version) = version { - let (version, url, releases_url) = tagged_release_location(version)?; + let (version, url, releases_url) = tagged_release_location(version, release_api_base_url)?; return fetch_tagged_release(&client, &url, releases_url) .await? .with_context(|| { @@ -610,7 +628,8 @@ async fn fetch_release(version: Option<&str>) -> Result { }); } - let list_url = format!("https://api.github.com/repos/{REPO}/releases?per_page=100"); + let api_base = release_api_base_url.trim_end_matches('/'); + let list_url = format!("{api_base}/{REPO}/releases?per_page=100"); let releases: Vec = client .get(&list_url) .header("user-agent", "dstackup") @@ -630,13 +649,16 @@ async fn fetch_release(version: Option<&str>) -> Result { return Ok(release); } - let legacy_url = format!("https://api.github.com/repos/{LEGACY_REPO}/releases/latest"); + let legacy_url = format!("{api_base}/{LEGACY_REPO}/releases/latest"); fetch_tagged_release(&client, &legacy_url, LEGACY_RELEASES_URL) .await? .with_context(|| format!("no guest-OS release found; check {RELEASES_URL}")) } -fn tagged_release_location(version: &str) -> Result<(String, String, &'static str)> { +fn tagged_release_location( + version: &str, + release_api_base_url: &str, +) -> Result<(String, String, &'static str)> { let version = version .trim_start_matches(RELEASE_TAG_PREFIX) .trim_start_matches('v'); @@ -648,7 +670,10 @@ fn tagged_release_location(version: &str) -> Result<(String, String, &'static st }; Ok(( version.to_string(), - format!("https://api.github.com/repos/{repo}/releases/tags/{tag_prefix}{version}"), + format!( + "{}/{repo}/releases/tags/{tag_prefix}{version}", + release_api_base_url.trim_end_matches('/') + ), releases_url, )) } @@ -830,7 +855,8 @@ mod tests { #[test] fn routes_pinned_releases_at_the_monorepo_boundary() { for version in ["0.5.11", "v0.5.11", "guest-os-v0.5.11"] { - let (normalized, url, releases_url) = tagged_release_location(version).unwrap(); + let (normalized, url, releases_url) = + tagged_release_location(version, crate::cli::DEFAULT_RELEASE_API_BASE_URL).unwrap(); assert_eq!(normalized, "0.5.11"); assert_eq!( url, @@ -840,7 +866,8 @@ mod tests { } for version in ["0.6.0", "0.6.0.a2", "1.0.0"] { - let (normalized, url, releases_url) = tagged_release_location(version).unwrap(); + let (normalized, url, releases_url) = + tagged_release_location(version, crate::cli::DEFAULT_RELEASE_API_BASE_URL).unwrap(); assert_eq!(normalized, version); assert_eq!( url, @@ -853,10 +880,22 @@ mod tests { #[test] fn rejects_versions_without_a_numeric_core() { for version in ["0.6", "latest", "0.x.0", "0.6.x"] { - assert!(tagged_release_location(version).is_err(), "{version}"); + assert!( + tagged_release_location(version, crate::cli::DEFAULT_RELEASE_API_BASE_URL).is_err(), + "{version}" + ); } } + #[test] + fn release_api_base_url_is_configurable_and_trailing_slash_safe() { + let (_, url, _) = tagged_release_location("0.6.0", "http://127.0.0.1:1234/api/").unwrap(); + assert_eq!( + url, + "http://127.0.0.1:1234/api/Dstack-TEE/dstack/releases/tags/guest-os-v0.6.0" + ); + } + #[test] fn messages_mention_the_pull_command() { assert!(no_image_message("/d").contains("dstackup image pull")); diff --git a/dstack/crates/dstackup/src/install.rs b/dstack/crates/dstackup/src/install.rs index 469fc51ae..0a9c442d8 100644 --- a/dstack/crates/dstackup/src/install.rs +++ b/dstack/crates/dstackup/src/install.rs @@ -28,7 +28,7 @@ const DAEMON_BINARIES: &[(&str, &str)] = &[ ("supervisor", "supervisor"), ]; -pub(crate) async fn cmd_install(mut o: InstallOpts) -> Result<()> { +pub(crate) async fn cmd_install(mut o: InstallOpts, release_api_base_url: &str) -> Result<()> { // --expose is not safe yet: the rendered vmm.toml binds the VM-control // plane with neither TLS nor an auth token (the management RPCs are not // behind an auth guard), so exposing it would hand deploy/destroy to anyone @@ -99,6 +99,7 @@ pub(crate) async fn cmd_install(mut o: InstallOpts) -> Result<()> { o.image.as_deref(), !o.no_kms, required_image_files, + release_api_base_url, ) .await?; let os_image_hash = resolve_image_pin(&o, &images, platform)?; diff --git a/dstack/crates/dstackup/src/main.rs b/dstack/crates/dstackup/src/main.rs index 42082685c..9529231e5 100644 --- a/dstack/crates/dstackup/src/main.rs +++ b/dstack/crates/dstackup/src/main.rs @@ -38,8 +38,8 @@ async fn main() -> Result<()> { let platform = default_platform(prefix.as_deref()).or_else(host::Platform::detect); cmd_status(&host, platform).await } - Command::Install(opts) => install::cmd_install(opts).await, - Command::Image(cmd) => image::cmd_image(cmd).await, + Command::Install(opts) => install::cmd_install(opts, &cli.release_api_base_url).await, + Command::Image(cmd) => image::cmd_image(cmd, &cli.release_api_base_url).await, Command::Destroy { prefix, purge } => destroy::cmd_destroy(prefix.as_deref(), purge).await, } } diff --git a/tools/mock-github-releases.py b/tools/mock-github-releases.py new file mode 100755 index 000000000..0d7b240a5 --- /dev/null +++ b/tools/mock-github-releases.py @@ -0,0 +1,195 @@ +#!/usr/bin/env python3 +"""Local GitHub Releases API overlay for dstackup development. + +Locally published releases shadow GitHub. Requests not satisfied locally are +relayed to https://api.github.com. Release assets can be copied into --asset-dir +by specifying an asset's "local_path" in the publish request. +""" + +import argparse +import hashlib +import json +import mimetypes +import shutil +import sys +import urllib.error +import urllib.request +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from urllib.parse import unquote, urlsplit + +UPSTREAM = "https://api.github.com" + + +class State: + def __init__(self, asset_dir: Path): + self.asset_dir = asset_dir.resolve() + self.asset_dir.mkdir(parents=True, exist_ok=True) + self.releases = {} # repo -> tag -> release; insertion order is newest first + + def publish(self, repo, release, public_base): + tag = release.get("tag_name") + if not isinstance(tag, str) or not tag: + raise ValueError("tag_name must be a non-empty string") + assets = [] + for source in release.get("assets", []): + asset = dict(source) + local_path = asset.pop("local_path", None) + if local_path: + src = Path(local_path).resolve(strict=True) + name = asset.get("name") or src.name + if Path(name).name != name: + raise ValueError(f"unsafe asset name: {name!r}") + dst = self.asset_dir / name + shutil.copyfile(src, dst) + digest = hashlib.sha256(dst.read_bytes()).hexdigest() + asset.update( + name=name, + browser_download_url=f"{public_base}/__assets/{name}", + digest=f"sha256:{digest}", + ) + if not asset.get("name") or not asset.get("browser_download_url"): + raise ValueError("each asset needs name and browser_download_url, or local_path") + assets.append(asset) + release = dict(release, assets=assets) + releases = self.releases.setdefault(repo, {}) + releases.pop(tag, None) + self.releases[repo] = {tag: release, **releases} + return release + + +class Handler(BaseHTTPRequestHandler): + server_version = "dstack-release-mock/1" + + def json_response(self, status, value): + body = json.dumps(value, indent=2).encode() + self.send_response(status) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_POST(self): + # POST /__admin/repos/{owner}/{repo}/releases + parts = self.path.strip("/").split("/") + if len(parts) != 5 or parts[:2] != ["__admin", "repos"] or parts[4:] != ["releases"]: + self.send_error(404) + return + try: + size = int(self.headers.get("content-length", "0")) + release = json.loads(self.rfile.read(size)) + repo = "/".join(parts[2:4]) + host = self.headers.get("host", f"127.0.0.1:{self.server.server_port}") + public_base = f"http://{host}" + release = self.server.state.publish(repo, release, public_base) + self.json_response(201, release) + except (ValueError, OSError, json.JSONDecodeError) as error: + self.json_response(400, {"error": str(error)}) + + def do_GET(self): + path = urlsplit(self.path).path + if path.startswith("/__assets/"): + self.serve_asset(unquote(path.removeprefix("/__assets/"))) + return + # GitHub-compatible paths below /repos/{owner}/{repo}/releases... + parts = path.strip("/").split("/") + if len(parts) >= 4 and parts[0] == "repos": + repo = "/".join(parts[1:3]) + suffix = parts[3:] + local = self.server.state.releases.get(repo, {}) + if suffix == ["releases"] and local: + # Preserve local newest-first ordering, then add non-shadowed upstream releases. + upstream = self.fetch_upstream_json() + if isinstance(upstream, list): + local_tags = set(local) + upstream = [r for r in upstream if r.get("tag_name") not in local_tags] + else: + upstream = [] + self.json_response(200, list(local.values()) + upstream) + return + if suffix == ["releases", "latest"] and local: + self.json_response(200, next(iter(local.values()))) + return + if len(suffix) == 3 and suffix[:2] == ["releases", "tags"]: + release = local.get(unquote(suffix[2])) + if release: + self.json_response(200, release) + return + self.relay() + + def serve_asset(self, name): + if not name or Path(name).name != name: + self.send_error(400, "invalid asset name") + return + path = self.server.state.asset_dir / name + if not path.is_file(): + self.send_error(404) + return + self.send_response(200) + self.send_header("content-type", mimetypes.guess_type(name)[0] or "application/octet-stream") + self.send_header("content-length", str(path.stat().st_size)) + self.end_headers() + with path.open("rb") as source: + shutil.copyfileobj(source, self.wfile) + + def upstream_request(self): + url = UPSTREAM + self.path + headers = {"user-agent": "dstack-release-mock", "accept": "application/vnd.github+json"} + auth = self.headers.get("authorization") + if auth: + headers["authorization"] = auth + return urllib.request.Request(url, headers=headers) + + def fetch_upstream_json(self): + try: + with urllib.request.urlopen(self.upstream_request(), timeout=30) as response: + return json.load(response) + except (urllib.error.URLError, json.JSONDecodeError): + return [] + + def relay(self): + try: + with urllib.request.urlopen(self.upstream_request(), timeout=30) as response: + body = response.read() + self.send_response(response.status) + self.send_header("content-type", response.headers.get("content-type", "application/json")) + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + except urllib.error.HTTPError as error: + body = error.read() + self.send_response(error.code) + self.send_header("content-type", error.headers.get("content-type", "application/json")) + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + except urllib.error.URLError as error: + self.json_response(502, {"error": f"upstream request failed: {error}"}) + + def log_message(self, fmt, *args): + sys.stderr.write("%s - %s\n" % (self.address_string(), fmt % args)) + + +class Server(ThreadingHTTPServer): + def __init__(self, address, state): + super().__init__(address, Handler) + self.state = state + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--listen", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--asset-dir", type=Path, default=Path("/tmp/dstack-release-mock-assets")) + args = parser.parse_args() + server = Server((args.listen, args.port), State(args.asset_dir)) + print(f"release API: http://{args.listen}:{args.port}/repos", flush=True) + print(f"publish API: http://{args.listen}:{args.port}/__admin/repos/OWNER/REPO/releases", flush=True) + try: + server.serve_forever() + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + main() From 445cd1ad75061059091c6e3fc704fa00fcbae230 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 20 Jul 2026 21:18:31 -0700 Subject: [PATCH 2/2] fix: address release API review feedback --- dstack/crates/dstackup/src/cli.rs | 4 +- dstack/crates/dstackup/src/image.rs | 7 +-- tools/mock-github-releases.py | 81 ++++++++++++++++++++++++----- 3 files changed, 73 insertions(+), 19 deletions(-) diff --git a/dstack/crates/dstackup/src/cli.rs b/dstack/crates/dstackup/src/cli.rs index 0d4107c9e..4022aada7 100644 --- a/dstack/crates/dstackup/src/cli.rs +++ b/dstack/crates/dstackup/src/cli.rs @@ -22,8 +22,8 @@ pub(crate) struct Cli { #[arg(long, global = true)] pub(crate) host: Option, - /// Base URL of the GitHub-compatible releases API. The repository path is - /// appended automatically (useful for an API proxy or a test server). + /// Base URL of the GitHub-compatible releases API, including its `/repos` + /// prefix. The owner/repository path is appended automatically. #[arg(long, global = true, default_value = DEFAULT_RELEASE_API_BASE_URL)] pub(crate) release_api_base_url: String, diff --git a/dstack/crates/dstackup/src/image.rs b/dstack/crates/dstackup/src/image.rs index 191a25e7a..e598915ac 100644 --- a/dstack/crates/dstackup/src/image.rs +++ b/dstack/crates/dstackup/src/image.rs @@ -628,7 +628,7 @@ async fn fetch_release(version: Option<&str>, release_api_base_url: &str) -> Res }); } - let api_base = release_api_base_url.trim_end_matches('/'); + let api_base = release_api_base_url.trim().trim_end_matches('/'); let list_url = format!("{api_base}/{REPO}/releases?per_page=100"); let releases: Vec = client .get(&list_url) @@ -672,7 +672,7 @@ fn tagged_release_location( version.to_string(), format!( "{}/{repo}/releases/tags/{tag_prefix}{version}", - release_api_base_url.trim_end_matches('/') + release_api_base_url.trim().trim_end_matches('/') ), releases_url, )) @@ -889,7 +889,8 @@ mod tests { #[test] fn release_api_base_url_is_configurable_and_trailing_slash_safe() { - let (_, url, _) = tagged_release_location("0.6.0", "http://127.0.0.1:1234/api/").unwrap(); + let (_, url, _) = + tagged_release_location("0.6.0", " http://127.0.0.1:1234/api/ ").unwrap(); assert_eq!( url, "http://127.0.0.1:1234/api/Dstack-TEE/dstack/releases/tags/guest-os-v0.6.0" diff --git a/tools/mock-github-releases.py b/tools/mock-github-releases.py index 0d7b240a5..5132507eb 100755 --- a/tools/mock-github-releases.py +++ b/tools/mock-github-releases.py @@ -1,4 +1,8 @@ #!/usr/bin/env python3 +# SPDX-FileCopyrightText: © 2026 Phala Network +# +# SPDX-License-Identifier: Apache-2.0 + """Local GitHub Releases API overlay for dstackup development. Locally published releases shadow GitHub. Requests not satisfied locally are @@ -22,12 +26,16 @@ class State: + """Thread-safe-by-snapshot in-memory release overlay.""" + def __init__(self, asset_dir: Path): + """Create an empty overlay and its local asset directory.""" self.asset_dir = asset_dir.resolve() self.asset_dir.mkdir(parents=True, exist_ok=True) self.releases = {} # repo -> tag -> release; insertion order is newest first def publish(self, repo, release, public_base): + """Publish a release as the newest local release for a repository.""" tag = release.get("tag_name") if not isinstance(tag, str) or not tag: raise ValueError("tag_name must be a non-empty string") @@ -49,19 +57,28 @@ def publish(self, repo, release, public_base): digest=f"sha256:{digest}", ) if not asset.get("name") or not asset.get("browser_download_url"): - raise ValueError("each asset needs name and browser_download_url, or local_path") + raise ValueError( + "each asset needs name and browser_download_url, or local_path" + ) assets.append(asset) release = dict(release, assets=assets) - releases = self.releases.setdefault(repo, {}) - releases.pop(tag, None) - self.releases[repo] = {tag: release, **releases} + # Build a new snapshot rather than mutating one observed by a concurrent + # GET handler. Assignment is atomic under CPython's GIL. + releases = self.releases.get(repo, {}) + self.releases[repo] = { + tag: release, + **{old_tag: old for old_tag, old in releases.items() if old_tag != tag}, + } return release class Handler(BaseHTTPRequestHandler): + """Serve the local release overlay and relay misses to GitHub.""" + server_version = "dstack-release-mock/1" def json_response(self, status, value): + """Send a JSON response with a fixed content length.""" body = json.dumps(value, indent=2).encode() self.send_response(status) self.send_header("content-type", "application/json") @@ -70,9 +87,14 @@ def json_response(self, status, value): self.wfile.write(body) def do_POST(self): + """Publish a release through the local administration endpoint.""" # POST /__admin/repos/{owner}/{repo}/releases parts = self.path.strip("/").split("/") - if len(parts) != 5 or parts[:2] != ["__admin", "repos"] or parts[4:] != ["releases"]: + if ( + len(parts) != 5 + or parts[:2] != ["__admin", "repos"] + or parts[4:] != ["releases"] + ): self.send_error(404) return try: @@ -87,6 +109,7 @@ def do_POST(self): self.json_response(400, {"error": str(error)}) def do_GET(self): + """Serve an asset or a GitHub-compatible releases API request.""" path = urlsplit(self.path).path if path.startswith("/__assets/"): self.serve_asset(unquote(path.removeprefix("/__assets/"))) @@ -102,7 +125,9 @@ def do_GET(self): upstream = self.fetch_upstream_json() if isinstance(upstream, list): local_tags = set(local) - upstream = [r for r in upstream if r.get("tag_name") not in local_tags] + upstream = [ + r for r in upstream if r.get("tag_name") not in local_tags + ] else: upstream = [] self.json_response(200, list(local.values()) + upstream) @@ -118,6 +143,7 @@ def do_GET(self): self.relay() def serve_asset(self, name): + """Serve one basename-only file from the configured asset directory.""" if not name or Path(name).name != name: self.send_error(400, "invalid asset name") return @@ -126,40 +152,57 @@ def serve_asset(self, name): self.send_error(404) return self.send_response(200) - self.send_header("content-type", mimetypes.guess_type(name)[0] or "application/octet-stream") + self.send_header( + "content-type", mimetypes.guess_type(name)[0] or "application/octet-stream" + ) self.send_header("content-length", str(path.stat().st_size)) self.end_headers() with path.open("rb") as source: shutil.copyfileobj(source, self.wfile) def upstream_request(self): + """Build the corresponding authenticated GitHub API request.""" url = UPSTREAM + self.path - headers = {"user-agent": "dstack-release-mock", "accept": "application/vnd.github+json"} + headers = { + "user-agent": "dstack-release-mock", + "accept": "application/vnd.github+json", + } auth = self.headers.get("authorization") if auth: headers["authorization"] = auth return urllib.request.Request(url, headers=headers) def fetch_upstream_json(self): + """Fetch upstream JSON for merging, returning an empty list on failure.""" try: - with urllib.request.urlopen(self.upstream_request(), timeout=30) as response: + with urllib.request.urlopen( + self.upstream_request(), timeout=30 + ) as response: return json.load(response) except (urllib.error.URLError, json.JSONDecodeError): return [] def relay(self): + """Relay the current request and its HTTP status to GitHub.""" try: - with urllib.request.urlopen(self.upstream_request(), timeout=30) as response: + with urllib.request.urlopen( + self.upstream_request(), timeout=30 + ) as response: body = response.read() self.send_response(response.status) - self.send_header("content-type", response.headers.get("content-type", "application/json")) + self.send_header( + "content-type", + response.headers.get("content-type", "application/json"), + ) self.send_header("content-length", str(len(body))) self.end_headers() self.wfile.write(body) except urllib.error.HTTPError as error: body = error.read() self.send_response(error.code) - self.send_header("content-type", error.headers.get("content-type", "application/json")) + self.send_header( + "content-type", error.headers.get("content-type", "application/json") + ) self.send_header("content-length", str(len(body))) self.end_headers() self.wfile.write(body) @@ -167,24 +210,34 @@ def relay(self): self.json_response(502, {"error": f"upstream request failed: {error}"}) def log_message(self, fmt, *args): + """Write HTTP access logs to stderr.""" sys.stderr.write("%s - %s\n" % (self.address_string(), fmt % args)) class Server(ThreadingHTTPServer): + """Threading HTTP server carrying shared overlay state.""" + def __init__(self, address, state): + """Bind the server and attach its overlay state.""" super().__init__(address, Handler) self.state = state def main(): + """Run the local releases API overlay until interrupted.""" parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--listen", default="127.0.0.1") parser.add_argument("--port", type=int, default=8000) - parser.add_argument("--asset-dir", type=Path, default=Path("/tmp/dstack-release-mock-assets")) + parser.add_argument( + "--asset-dir", type=Path, default=Path("/tmp/dstack-release-mock-assets") + ) args = parser.parse_args() server = Server((args.listen, args.port), State(args.asset_dir)) print(f"release API: http://{args.listen}:{args.port}/repos", flush=True) - print(f"publish API: http://{args.listen}:{args.port}/__admin/repos/OWNER/REPO/releases", flush=True) + print( + f"publish API: http://{args.listen}:{args.port}/__admin/repos/OWNER/REPO/releases", + flush=True, + ) try: server.serve_forever() except KeyboardInterrupt: