1use anyhow::{Result, anyhow};
2use colored::*;
3use dirs;
4use hex;
5use indicatif::{ProgressBar, ProgressStyle};
6use self_update::self_replace;
7use serde::Deserialize;
8use sha2::{Digest, Sha512};
9use std::env;
10use std::fs::File;
11use std::io::{self, Read, Write};
12use std::path::{Path, PathBuf};
13use tar::Archive;
14use tempfile::Builder;
15use zip::ZipArchive;
16use zstd::stream::read::Decoder as ZstdDecoder;
17
18const GITLAB_PROJECT_PATH: &str = "zillowe/zillwen/zusty/zoi";
19const GITLAB_PROJECT_ID: &str = "71087662";
20
21#[derive(Debug, Deserialize)]
22struct GitLabRelease {
23 tag_name: String,
24}
25
26fn get_latest_tag(branch_prefix: &str) -> Result<String> {
27 println!("Fetching latest release information from GitLab...");
28 let api_url = format!(
29 "https://gitlab.com/api/v4/projects/{}/releases",
30 GITLAB_PROJECT_ID
31 );
32 let client = reqwest::blocking::Client::builder()
33 .user_agent("Zoi-Upgrader")
34 .use_rustls_tls()
35 .build()?;
36 let releases: Vec<GitLabRelease> = client.get(&api_url).send()?.json()?;
37
38 let latest_tag = releases
39 .into_iter()
40 .find(|r| r.tag_name.starts_with(branch_prefix))
41 .map(|r| r.tag_name)
42 .ok_or_else(|| anyhow!("No release found with prefix '{}'", branch_prefix))?;
43
44 println!(
45 "Found latest tag for branch prefix '{}': {}",
46 branch_prefix,
47 latest_tag.green()
48 );
49 Ok(latest_tag)
50}
51
52fn download_file(url: &str, path: &Path) -> Result<()> {
53 let mut response = reqwest::blocking::get(url)?;
54 if !response.status().is_success() {
55 return Err(anyhow!(
56 "Failed to download file: HTTP {}",
57 response.status()
58 ));
59 }
60
61 let total_size = response.content_length().unwrap_or(0);
62 let pb = ProgressBar::new(total_size);
63 pb.set_style(ProgressStyle::default_bar()
64 .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec})")?
65 .progress_chars("#>- "));
66
67 let mut dest = File::create(path)?;
68 let mut buffer = [0; 8192];
69
70 loop {
71 let bytes_read = response.read(&mut buffer)?;
72 if bytes_read == 0 {
73 break;
74 }
75 dest.write_all(&buffer[..bytes_read])?;
76 pb.inc(bytes_read as u64);
77 }
78
79 pb.finish_with_message("Download complete.");
80 Ok(())
81}
82
83fn extract_archive(archive_path: &Path, target_dir: &Path) -> Result<()> {
84 println!("Extracting binary...");
85 let file = File::open(archive_path)?;
86
87 if archive_path.extension().and_then(|s| s.to_str()) == Some("zip") {
88 let mut archive = ZipArchive::new(file)?;
89 archive.extract(target_dir)?;
90 } else {
91 let tar = ZstdDecoder::new(file)?;
92 let mut archive = Archive::new(tar);
93 archive.unpack(target_dir)?;
94 }
95 Ok(())
96}
97
98fn verify_checksum(file_path: &Path, checksums_content: &str, filename: &str) -> Result<()> {
99 println!("Verifying checksum for {}...", filename);
100 let expected_hash = checksums_content
101 .lines()
102 .find(|line| line.contains(filename))
103 .and_then(|line| line.split_whitespace().next())
104 .ok_or(anyhow!("Checksum not found for {}.", filename))?;
105
106 let mut file = File::open(file_path)?;
107 let mut hasher = Sha512::new();
108 let mut buffer = [0; 8192];
109 loop {
110 let bytes_read = io::Read::read(&mut file, &mut buffer)?;
111 if bytes_read == 0 {
112 break;
113 }
114 hasher.update(&buffer[..bytes_read]);
115 }
116 let actual_hash = hex::encode(hasher.finalize());
117
118 if actual_hash != expected_hash {
119 return Err(anyhow!(
120 "Checksum mismatch for {}! The file may be corrupt.",
121 filename
122 ));
123 }
124 println!("Checksum verified successfully for {}.", filename.green());
125 Ok(())
126}
127
128fn get_platform_info() -> Result<(&'static str, &'static str)> {
129 let os = match env::consts::OS {
130 "linux" => "linux",
131 "macos" | "darwin" => "macos",
132 "windows" => "windows",
133 _ => return Err(anyhow!("Unsupported OS: {}", env::consts::OS)),
134 };
135 let arch = match env::consts::ARCH {
136 "x86_64" => "amd64",
137 "aarch64" => "arm64",
138 _ => return Err(anyhow!("Unsupported architecture: {}", env::consts::ARCH)),
139 };
140 Ok((os, arch))
141}
142
143fn fallback_full_upgrade(
144 base_url: &str,
145 checksums_content: &str,
146 os: &str,
147 arch: &str,
148) -> Result<(PathBuf, tempfile::TempDir)> {
149 let archive_ext = if os == "windows" { "zip" } else { "tar.zst" };
150 let archive_filename = format!("zoi-{os}-{arch}.{archive_ext}");
151 let download_url = format!("{base_url}/{archive_filename}");
152 let temp_dir = Builder::new().prefix("zoi-full-upgrade").tempdir()?;
153 let temp_archive_path = temp_dir.path().join(&archive_filename);
154
155 println!("Downloading Zoi from: {download_url}");
156 download_file(&download_url, &temp_archive_path)?;
157 verify_checksum(&temp_archive_path, checksums_content, &archive_filename)?;
158
159 extract_archive(&temp_archive_path, temp_dir.path())?;
160
161 let binary_filename = if os == "windows" { "zoi.exe" } else { "zoi" };
162 let new_binary_path = temp_dir.path().join(binary_filename);
163 if !new_binary_path.exists() {
164 return Err(anyhow!(
165 "Could not find executable in the extracted archive."
166 ));
167 }
168 Ok((new_binary_path, temp_dir))
169}
170
171fn try_delta_upgrade(
172 base_url: &str,
173 checksums_content: &str,
174 os: &str,
175 arch: &str,
176 current_version: &str,
177 latest_version: &str,
178) -> Result<(PathBuf, tempfile::TempDir)> {
179 let archive_basename = format!("zoi-{os}-{arch}");
180 let bsdiff_filename =
181 format!("{archive_basename}.from-v{current_version}-to-v{latest_version}.bsdiff");
182 let download_url = format!("{base_url}/{bsdiff_filename}");
183
184 if !checksums_content.contains(&bsdiff_filename) {
185 return Err(anyhow!("Delta patch not available for this upgrade path."));
186 }
187
188 let temp_dir = Builder::new().prefix("zoi-delta-upgrade").tempdir()?;
189 let temp_patch_path = temp_dir.path().join(&bsdiff_filename);
190
191 println!("{} Downloading delta patch...", "::".bold().blue());
192 download_file(&download_url, &temp_patch_path)?;
193 verify_checksum(&temp_patch_path, checksums_content, &bsdiff_filename)?;
194
195 println!("{} Applying delta patch...", "::".bold().blue());
196 let current_exe_path = env::current_exe()?;
197 let mut old_binary = Vec::new();
198 File::open(¤t_exe_path)?.read_to_end(&mut old_binary)?;
199
200 let mut patch_data = Vec::new();
201 File::open(&temp_patch_path)?.read_to_end(&mut patch_data)?;
202
203 let raw_patch = if patch_data.starts_with(&[0x28, 0xB5, 0x2F, 0xFD]) {
204 let mut decoder = ZstdDecoder::new(std::io::Cursor::new(&patch_data))?;
205 let mut buf = Vec::new();
206 decoder.read_to_end(&mut buf)?;
207 buf
208 } else {
209 patch_data
210 };
211
212 let mut cursor = std::io::Cursor::new(raw_patch);
213 let mut new_binary = Vec::new();
214 bsdiff::patch(&old_binary, &mut cursor, &mut new_binary)?;
215
216 let binary_filename = if os == "windows" { "zoi.exe" } else { "zoi" };
217 let new_binary_path = temp_dir.path().join(binary_filename);
218 std::fs::write(&new_binary_path, &new_binary)?;
219
220 Ok((new_binary_path, temp_dir))
221}
222
223pub fn run(
224 branch: &str,
225 status: &str,
226 number: &str,
227 force: bool,
228 tag: Option<String>,
229 custom_branch: Option<String>,
230) -> Result<()> {
231 if crate::offline::is_offline() {
232 return Err(anyhow!("Cannot upgrade Zoi: Zoi is in offline mode."));
233 }
234 let current_exe_path = env::current_exe()?;
235 let path_str = current_exe_path.to_string_lossy();
236
237 let is_cargo_install = dirs::home_dir()
238 .map(|home| current_exe_path.starts_with(home.join(".cargo").join("bin")))
239 .unwrap_or(false);
240
241 let pkg_manager = if path_str.contains("/Cellar/") {
242 Some("Homebrew")
243 } else if path_str.contains("scoop/apps/") {
244 Some("Scoop")
245 } else if path_str.starts_with("/usr/bin/") {
246 Some("a system package manager")
247 } else if is_cargo_install {
248 Some("Cargo")
249 } else {
250 None
251 };
252
253 if let Some(pm) = pkg_manager {
254 if !force {
255 eprintln!(
256 "{}{}{}",
257 "Warning: ".yellow().bold(),
258 "It looks like Zoi was installed via ".yellow(),
259 pm.yellow().bold()
260 );
261 eprintln!(
262 "{}",
263 "Using 'zoi upgrade' may conflict with your package manager.".yellow()
264 );
265 let upgrade_command = match pm {
266 "Homebrew" => "brew upgrade zoi",
267 "Scoop" => "scoop update zoi",
268 "Cargo" => "cargo install zoi-rs",
269 _ => "your package manager's upgrade command",
270 };
271 eprintln!(
272 "It is recommended to use '{}' to upgrade Zoi.",
273 upgrade_command.cyan()
274 );
275 eprintln!(
276 "To override this check and proceed anyway, run with the '{}' flag.",
277 "--force".cyan()
278 );
279 return Err(anyhow!("managed_by_package_manager"));
280 } else {
281 println!(
282 "{}{}",
283 "Warning: ".yellow().bold(),
284 "Forcing self-upgrade on a package-manager-controlled installation.".yellow()
285 );
286 }
287 }
288
289 let current_version = if status.is_empty() || status.eq_ignore_ascii_case("stable") {
290 number.to_string()
291 } else {
292 format!("{}-{}", number, status.to_lowercase())
293 };
294
295 let latest_tag = if let Some(tag_name) = tag {
296 println!("Upgrading to specified tag: {}", tag_name.green());
297 tag_name
298 } else {
299 let branch_prefix = if let Some(b) = custom_branch {
300 println!("Upgrading to latest release from branch: {}", b.green());
301 format!("{}-", b)
302 } else if branch.eq_ignore_ascii_case("public") {
303 "Pub-".to_string()
304 } else {
305 "Prod-".to_string()
306 };
307 get_latest_tag(&branch_prefix)?
308 };
309
310 let parts: Vec<&str> = latest_tag.split('-').collect();
311 let latest_version_num = parts
312 .last()
313 .ok_or(anyhow!("Could not get version number from tag"))?;
314
315 let latest_version_str = if parts.len() > 2 {
316 let prerelease = parts[1].to_lowercase();
317 format!("{}-{}", latest_version_num, prerelease)
318 } else {
319 latest_version_num.to_string()
320 };
321
322 if !force && !self_update::version::bump_is_greater(¤t_version, &latest_version_str)? {
323 println!(
324 "
325{}",
326 "You are already on the latest version!".green()
327 );
328 return Err(anyhow!("already_on_latest"));
329 }
330
331 let (os, arch) = get_platform_info()?;
332
333 let base_url =
334 format!("https://gitlab.com/{GITLAB_PROJECT_PATH}/-/releases/{latest_tag}/downloads");
335 let checksums_txt_url = format!("{base_url}/checksums.txt");
336
337 println!(
338 "Downloading archive and checksums from: {}",
339 checksums_txt_url
340 );
341 let checksums_txt_content = reqwest::blocking::get(&checksums_txt_url)?.text()?;
342
343 let (new_binary_path, _temp_dir_guard) = if !force {
344 match try_delta_upgrade(
345 &base_url,
346 &checksums_txt_content,
347 os,
348 arch,
349 ¤t_version,
350 &latest_version_str,
351 ) {
352 Ok(res) => res,
353 Err(e) => {
354 println!("Delta upgrade failed: {}. Falling back to full upgrade.", e);
355 fallback_full_upgrade(&base_url, &checksums_txt_content, os, arch)?
356 }
357 }
358 } else {
359 fallback_full_upgrade(&base_url, &checksums_txt_content, os, arch)?
360 };
361
362 println!("Replacing current executable...");
363 self_replace::self_replace(&new_binary_path)?;
364
365 Ok(())
366}