1#[cfg(target_os = "linux")]
2use crate::services::integration::AppImageExtractor;
3use crate::{
4 models::{
5 common::{DesktopEntry, enums::TrustMode},
6 provider::Release,
7 upstream::{InstallType, Package},
8 },
9 providers::provider_manager::ProviderManager,
10 services::builder::{BuildRequest, scripts::BuildScriptAction, worker::BuildWorker},
11 services::{
12 integration::{DesktopManager, IconManager},
13 packaging::RollbackManager,
14 packaging::{PackageInstaller, PackagePhase, PackageProgressEvent, PackageRemover},
15 storage::rollback_storage::{RollbackSource, RollbackStorage},
16 trust::TrustedSignatureKeys,
17 },
18 utils::static_paths::UpstreamPaths,
19};
20
21use anyhow::{Context, Result, bail};
22use console::style;
23use std::{
24 fs,
25 path::{Path, PathBuf},
26};
27
28macro_rules! message {
29 ($cb:expr, $($arg:tt)*) => {{
30 if let Some(cb) = $cb.as_mut() {
31 cb(&format!($($arg)*));
32 }
33 }};
34}
35
36macro_rules! progress {
37 ($cb:expr, $event:expr) => {{
38 if let Some(cb) = $cb.as_mut() {
39 cb($event);
40 }
41 }};
42}
43
44pub struct PackageUpgrader<'a> {
45 provider_manager: &'a ProviderManager,
46 installer: PackageInstaller<'a>,
47 remover: PackageRemover<'a>,
48 paths: &'a UpstreamPaths,
49 trusted_keys: TrustedSignatureKeys,
50}
51
52#[derive(Clone)]
53pub enum ResolvedUpgradeTarget {
54 Release(Release),
55 Branch { branch: String, head_commit: String },
56}
57
58struct FailedUpgradeRollback<'a> {
59 previous_package: &'a Package,
60 partially_installed_package: Option<&'a Package>,
61 original_install_path: &'a Path,
62 backup_path: &'a Path,
63 failure_context: &'a str,
64}
65
66impl<'a> PackageUpgrader<'a> {
67 fn backup_path(paths: &UpstreamPaths, install_path: &Path) -> Result<PathBuf> {
68 let file_name = install_path.file_name().ok_or_else(|| {
69 anyhow::anyhow!("Install path '{}' has no filename", install_path.display())
70 })?;
71 fs::create_dir_all(&paths.install.tmp_dir).context(format!(
72 "Failed to create upgrade temp directory '{}'",
73 paths.install.tmp_dir.display()
74 ))?;
75 Ok(paths
76 .install
77 .tmp_dir
78 .join(format!("{}.old", file_name.to_string_lossy())))
79 }
80
81 fn remove_path_if_exists(path: &Path) -> Result<()> {
82 if path.is_dir() {
83 fs::remove_dir_all(path)
84 .context(format!("Failed to remove directory '{}'", path.display()))?;
85 } else if path.is_file() {
86 fs::remove_file(path).context(format!("Failed to remove file '{}'", path.display()))?;
87 }
88 Ok(())
89 }
90
91 fn capture_successful_upgrade_rollback(
92 paths: &UpstreamPaths,
93 package: &Package,
94 backup_path: &Path,
95 ) -> Result<()> {
96 let rollback_file = RollbackManager::rollback_file_path(paths);
97 let mut rollback_storage = RollbackStorage::new(&rollback_file)?;
98 RollbackManager::capture_backup_path(
99 paths,
100 &mut rollback_storage,
101 package,
102 backup_path,
103 RollbackSource::Upgrade,
104 &mut None::<fn(&str)>,
105 )
106 }
107
108 pub fn new(
109 provider_manager: &'a ProviderManager,
110 installer: PackageInstaller<'a>,
111 remover: PackageRemover<'a>,
112 paths: &'a UpstreamPaths,
113 trusted_keys: TrustedSignatureKeys,
114 ) -> Self {
115 Self {
116 provider_manager,
117 installer,
118 remover,
119 paths,
120 trusted_keys,
121 }
122 }
123
124 pub async fn upgrade<F, H>(
130 &self,
131 package: &Package,
132 force: bool,
133 trust_mode: TrustMode,
134 download_progress: &mut Option<F>,
135 message_callback: &mut Option<H>,
136 ) -> Result<Option<Package>>
137 where
138 F: FnMut(u64, u64),
139 H: FnMut(&str),
140 {
141 if package.is_pinned {
142 message!(
143 message_callback,
144 "Upgrade skipped: '{}' is pinned",
145 package.name
146 );
147 return Ok(None);
148 }
149
150 let target = if package.install_type == InstallType::Build && package.build_branch.is_some()
151 {
152 None
153 } else {
154 message!(message_callback, "Fetching latest release ...");
155 let release = if force {
156 self.provider_manager
157 .get_latest_release(
158 &package.repo_slug,
159 &package.provider,
160 &package.channel,
161 package.base_url.as_deref(),
162 )
163 .await
164 .context(format!(
165 "Failed to fetch latest release for '{}'",
166 package.name
167 ))?
168 } else {
169 let Some(latest_release) = self
170 .provider_manager
171 .check_for_updates(package)
172 .await
173 .context(format!(
174 "Failed to fetch latest release for '{}'",
175 package.name
176 ))?
177 else {
178 message!(message_callback, "'{}' is already up to date", package.name);
179 return Ok(None);
180 };
181
182 latest_release
183 };
184
185 if !force && !package.is_update_available(&release) {
186 message!(message_callback, "'{}' is already up to date", package.name);
187 return Ok(None);
188 }
189 Some(ResolvedUpgradeTarget::Release(release))
190 };
191
192 if package.install_type == InstallType::Build
193 && let Some(branch) = package.build_branch.as_deref()
194 {
195 let head_commit = self
196 .provider_manager
197 .get_branch_head_sha(
198 &package.repo_slug,
199 &package.provider,
200 branch,
201 package.base_url.as_deref(),
202 )
203 .await
204 .context(format!(
205 "Failed to fetch branch head for '{}' on '{}'",
206 branch, package.name
207 ))?;
208 let up_to_date = package
209 .build_commit
210 .as_deref()
211 .is_some_and(|saved| saved == head_commit);
212 if up_to_date && !force {
213 message!(
214 message_callback,
215 "'{}' is already up to date (branch '{}')",
216 package.name,
217 branch
218 );
219 return Ok(None);
220 }
221 let mut no_progress: Option<fn(PackageProgressEvent)> = None;
222 return self
223 .upgrade_resolved(
224 package,
225 ResolvedUpgradeTarget::Branch {
226 branch: branch.to_string(),
227 head_commit,
228 },
229 trust_mode,
230 download_progress,
231 message_callback,
232 &mut no_progress,
233 )
234 .await
235 .map(Some);
236 }
237
238 let Some(target) = target else {
239 return Ok(None);
240 };
241
242 let mut no_progress: Option<fn(PackageProgressEvent)> = None;
243 self.upgrade_resolved(
244 package,
245 target,
246 trust_mode,
247 download_progress,
248 message_callback,
249 &mut no_progress,
250 )
251 .await
252 .map(Some)
253 }
254
255 pub async fn upgrade_resolved<F, H, P>(
256 &self,
257 package: &Package,
258 target: ResolvedUpgradeTarget,
259 trust_mode: TrustMode,
260 download_progress: &mut Option<F>,
261 message_callback: &mut Option<H>,
262 progress_callback: &mut Option<P>,
263 ) -> Result<Package>
264 where
265 F: FnMut(u64, u64),
266 H: FnMut(&str),
267 P: FnMut(PackageProgressEvent),
268 {
269 if package.is_pinned {
270 bail!("Package '{}' is pinned", package.name);
271 }
272
273 let had_desktop_integration = package.icon_path.is_some();
274
275 message!(
276 message_callback,
277 "{}",
278 style(format!("Upgrading '{}' ...", package.name)).cyan()
279 );
280
281 let original_install_path = package
282 .install_path
283 .as_ref()
284 .ok_or_else(|| {
285 anyhow::anyhow!("Package '{}' has no install path recorded", package.name)
286 })?
287 .clone();
288 let backup_path = Self::backup_path(self.paths, &original_install_path)?;
289
290 Self::remove_path_if_exists(&backup_path)?;
291
292 progress!(
293 progress_callback,
294 PackageProgressEvent::Phase(PackagePhase::CreatingSnapshot)
295 );
296 fs::rename(&original_install_path, &backup_path).context(format!(
297 "Failed to back up '{}' to '{}'",
298 original_install_path.display(),
299 backup_path.display()
300 ))?;
301
302 progress!(
304 progress_callback,
305 PackageProgressEvent::Phase(PackagePhase::RemovingRuntimeLinks)
306 );
307 if let Err(e) = self
308 .remover
309 .remove_runtime_integrations(package, message_callback)
310 {
311 let _ = fs::rename(&backup_path, &original_install_path);
312 let _ = self
313 .remover
314 .restore_runtime_integrations(package, message_callback);
315 return Err(e).context(format!(
316 "Failed to remove runtime integration for '{}'",
317 package.name
318 ));
319 }
320
321 let install_result = if package.install_type == InstallType::Build {
323 progress!(
324 progress_callback,
325 PackageProgressEvent::Phase(PackagePhase::RebuildingFromSource)
326 );
327 let (version_tag, branch, branch_head_commit) = match &target {
328 ResolvedUpgradeTarget::Release(release) => (Some(release.tag.clone()), None, None),
329 ResolvedUpgradeTarget::Branch {
330 branch,
331 head_commit,
332 } => (None, Some(branch.clone()), Some(head_commit.clone())),
333 };
334 let worker = BuildWorker::new(self.provider_manager, self.paths);
335 let build_result = {
336 let mut build_line_callback = Some(|line: &str| {
337 let line = line.trim();
338 if !line.is_empty() {
339 progress!(
340 progress_callback,
341 PackageProgressEvent::Warning(line.to_string())
342 );
343 }
344 });
345 worker
346 .build(
347 BuildRequest {
348 name: package.name.clone(),
349 repo_slug: package.repo_slug.clone(),
350 provider: package.provider.clone(),
351 base_url: package.base_url.clone(),
352 version_tag,
353 branch,
354 requested_profile: None,
355 script_action: BuildScriptAction::Upgrade,
356 },
357 package.channel.clone(),
358 &mut build_line_callback,
359 )
360 .await
361 };
362
363 match build_result {
364 Ok(output) => {
365 let mut install_pkg = package.clone();
366 install_pkg.build_branch = output.branch.clone();
367 install_pkg.build_commit = output.commit.or(branch_head_commit.clone());
368 progress!(
369 progress_callback,
370 PackageProgressEvent::Phase(PackagePhase::InstallingPackage)
371 );
372 let mut install_message_callback = Some(|line: &str| {
373 let line = line.trim();
374 if !line.is_empty() {
375 progress!(
376 progress_callback,
377 PackageProgressEvent::Warning(line.to_string())
378 );
379 message!(message_callback, "{}", line);
380 }
381 });
382 self.installer.install_local_artifact_files(
383 install_pkg,
384 &output.artifact_path,
385 output.version,
386 &mut install_message_callback,
387 )
388 }
389 Err(e) => {
390 Err(e).context(format!("Failed to rebuild '{}' from source", package.name))
391 }
392 }
393 } else {
394 let ResolvedUpgradeTarget::Release(release) = &target else {
395 bail!(
396 "Resolved branch target cannot be used for release package '{}'",
397 package.name
398 );
399 };
400 self.installer
401 .install_package_files(
402 package.clone(),
403 release,
404 trust_mode,
405 &self.trusted_keys,
406 download_progress,
407 message_callback,
408 progress_callback,
409 )
410 .await
411 };
412 let mut updated_package = match install_result {
413 Ok(updated_package) => updated_package,
414 Err(install_err) => {
415 progress!(
416 progress_callback,
417 PackageProgressEvent::Phase(PackagePhase::RollingBack)
418 );
419
420 return self.rollback_failed_upgrade(
421 FailedUpgradeRollback {
422 previous_package: package,
423 partially_installed_package: None,
424 original_install_path: &original_install_path,
425 backup_path: &backup_path,
426 failure_context: "Failed to install new version",
427 },
428 install_err,
429 message_callback,
430 );
431 }
432 };
433
434 if had_desktop_integration {
436 progress!(
437 progress_callback,
438 PackageProgressEvent::Phase(PackagePhase::CreatingRuntimeLinks)
439 );
440
441 if let Err(err) = self
442 .add_desktop_integration(&mut updated_package, message_callback)
443 .await
444 {
445 progress!(
446 progress_callback,
447 PackageProgressEvent::Phase(PackagePhase::RollingBack)
448 );
449 return self.rollback_failed_upgrade(
450 FailedUpgradeRollback {
451 previous_package: package,
452 partially_installed_package: Some(&updated_package),
453 original_install_path: &original_install_path,
454 backup_path: &backup_path,
455 failure_context: "Failed to restore desktop integration",
456 },
457 err.context("Failed to restore desktop integration"),
458 message_callback,
459 );
460 }
461 }
462
463 if let Err(err) =
464 Self::capture_successful_upgrade_rollback(self.paths, package, &backup_path)
465 {
466 progress!(
467 progress_callback,
468 PackageProgressEvent::Warning(format!(
469 "Warning: failed to capture rollback for '{}': {}",
470 package.name, err
471 ))
472 );
473 Self::remove_path_if_exists(&backup_path)
474 .context(format!("Failed to remove backup for '{}'", package.name))?;
475 }
476
477 Ok(updated_package)
478 }
479
480 async fn add_desktop_integration<H>(
481 &self,
482 updated_package: &mut Package,
483 message_callback: &mut Option<H>,
484 ) -> Result<()>
485 where
486 H: FnMut(&str),
487 {
488 #[cfg(target_os = "linux")]
489 let appimage_extractor =
490 AppImageExtractor::new().context("Failed to initialize appimage extractor")?;
491
492 #[cfg(target_os = "linux")]
493 let icon_manager = IconManager::new(self.paths, &appimage_extractor);
494 #[cfg(not(target_os = "linux"))]
495 let icon_manager = IconManager::new(self.paths);
496
497 #[cfg(target_os = "linux")]
498 let desktop_manager = DesktopManager::new(self.paths, &appimage_extractor);
499 #[cfg(not(target_os = "linux"))]
500 let desktop_manager = DesktopManager::new(self.paths);
501
502 let install_path = updated_package.install_path.clone().ok_or_else(|| {
503 anyhow::anyhow!(
504 "Package '{}' has no install path after upgrade",
505 updated_package.name
506 )
507 })?;
508
509 let icon_path = icon_manager
510 .add_icon(
511 &updated_package.name,
512 &install_path,
513 &updated_package.filetype,
514 message_callback,
515 )
516 .await
517 .context(format!("Failed to add icon for '{}'", updated_package.name))?;
518
519 updated_package.icon_path = icon_path;
520
521 let desktop_entry = DesktopEntry::from_package(updated_package);
522
523 desktop_manager
524 .create_entry(
525 &install_path,
526 &updated_package.filetype,
527 desktop_entry,
528 message_callback,
529 )
530 .await
531 .context(format!(
532 "Failed to create desktop entry for '{}'",
533 updated_package.name
534 ))?;
535
536 Ok(())
537 }
538
539 fn rollback_failed_upgrade<H>(
540 &self,
541 rollback: FailedUpgradeRollback<'_>,
542 failure: anyhow::Error,
543 message_callback: &mut Option<H>,
544 ) -> Result<Package>
545 where
546 H: FnMut(&str),
547 {
548 let cleanup_result = if let Some(partial) = rollback.partially_installed_package {
549 self.remover.remove_package_files(partial, message_callback)
550 } else {
551 Self::remove_path_if_exists(rollback.original_install_path)
552 };
553
554 if let Err(cleanup_err) = cleanup_result {
555 return Err(anyhow::anyhow!(
556 "{} for '{}': {}. Rollback failed while removing partial install: {}",
557 rollback.failure_context,
558 rollback.previous_package.name,
559 failure,
560 cleanup_err
561 ));
562 }
563
564 fs::rename(rollback.backup_path, rollback.original_install_path).context(format!(
565 "{} for '{}': {}. Rollback failed while restoring backup",
566 rollback.failure_context, rollback.previous_package.name, failure
567 ))?;
568
569 self.remover
570 .restore_runtime_integrations(rollback.previous_package, message_callback)
571 .context(format!(
572 "{} for '{}': {}. Rollback failed while restoring runtime links",
573 rollback.failure_context, rollback.previous_package.name, failure
574 ))?;
575
576 Err(failure).context(format!(
577 "{} for '{}' (previous version restored)",
578 rollback.failure_context, rollback.previous_package.name
579 ))
580 }
581}
582
583#[cfg(test)]
584mod tests {
585 use super::{FailedUpgradeRollback, PackageUpgrader};
586 use crate::models::common::enums::{Channel, Filetype, Provider};
587 use crate::models::upstream::Package;
588 use crate::providers::provider_manager::ProviderManager;
589 use crate::services::packaging::{PackageInstaller, PackageRemover};
590 use crate::services::trust::TrustedSignatureKeys;
591 use crate::utils::{static_paths::UpstreamPaths, test_support};
592 use std::path::{Path, PathBuf};
593 use std::time::{SystemTime, UNIX_EPOCH};
594 use std::{fs, io};
595
596 fn temp_root(name: &str) -> PathBuf {
597 let nanos = SystemTime::now()
598 .duration_since(UNIX_EPOCH)
599 .map(|d| d.as_nanos())
600 .unwrap_or(0);
601 std::env::temp_dir().join(format!("upstream-upgrader-test-{name}-{nanos}"))
602 }
603
604 fn cleanup(path: &Path) -> io::Result<()> {
605 fs::remove_dir_all(path)
606 }
607
608 fn expected_symlink_path(paths: &UpstreamPaths, name: &str) -> PathBuf {
609 let base = paths.integration.symlinks_dir.join(name);
610 #[cfg(windows)]
611 {
612 return base.with_extension("exe");
613 }
614 #[cfg(not(windows))]
615 {
616 base
617 }
618 }
619
620 fn test_paths(root: &Path) -> crate::utils::static_paths::UpstreamPaths {
621 test_support::upstream_paths(root)
622 }
623
624 fn test_package(name: &str, install_path: PathBuf) -> Package {
625 let mut package = Package::with_defaults(
626 name.to_string(),
627 format!("owner/{name}"),
628 Filetype::Binary,
629 None,
630 None,
631 Channel::Stable,
632 Provider::Github,
633 None,
634 );
635 package.install_path = Some(install_path.clone());
636 package.exec_path = Some(install_path);
637 package
638 }
639
640 #[test]
641 fn remove_path_if_exists_handles_files_and_directories() {
642 let root = temp_root("remove");
643 let file = root.join("f.bin");
644 let dir = root.join("d");
645 fs::create_dir_all(&dir).expect("create dir");
646 fs::write(&file, b"content").expect("write file");
647
648 PackageUpgrader::remove_path_if_exists(&file).expect("remove file");
649 PackageUpgrader::remove_path_if_exists(&dir).expect("remove dir");
650 PackageUpgrader::remove_path_if_exists(&root.join("missing")).expect("ignore missing");
651
652 assert!(!file.exists());
653 assert!(!dir.exists());
654
655 cleanup(&root).expect("cleanup");
656 }
657
658 #[test]
659 fn rollback_failed_upgrade_removes_partial_install_and_restores_previous_binary() {
660 let root = temp_root("rollback-desktop-failure");
661 let paths = test_paths(&root);
662 fs::create_dir_all(&paths.install.binaries_dir).expect("create binaries dir");
663 fs::create_dir_all(&paths.install.tmp_dir).expect("create tmp dir");
664 fs::create_dir_all(&paths.integration.symlinks_dir).expect("create symlinks dir");
665
666 let install_path = paths.install.binaries_dir.join("tool");
667 let backup_path = paths.install.tmp_dir.join("tool.old");
668 fs::write(&install_path, b"new").expect("write partial new binary");
669 fs::write(&backup_path, b"old").expect("write backup binary");
670
671 let previous = test_package("tool", install_path.clone());
672 let partial = test_package("tool", install_path.clone());
673 let provider_manager =
674 ProviderManager::new(None, None, None, Default::default()).expect("provider manager");
675 let installer = PackageInstaller::new(&provider_manager, &paths).expect("installer");
676 let remover = PackageRemover::new(&paths);
677 let upgrader = PackageUpgrader::new(
678 &provider_manager,
679 installer,
680 remover,
681 &paths,
682 TrustedSignatureKeys::default(),
683 );
684 let mut msg = Some(|_: &str| {});
685
686 let err = upgrader
687 .rollback_failed_upgrade(
688 FailedUpgradeRollback {
689 previous_package: &previous,
690 partially_installed_package: Some(&partial),
691 original_install_path: &install_path,
692 backup_path: &backup_path,
693 failure_context: "Failed to restore desktop integration",
694 },
695 anyhow::anyhow!("desktop failed"),
696 &mut msg,
697 )
698 .expect_err("rollback helper returns original failure");
699
700 assert!(err.to_string().contains("previous version restored"));
701 assert_eq!(
702 fs::read(&install_path).expect("read restored binary"),
703 b"old"
704 );
705 assert!(!backup_path.exists());
706 assert!(expected_symlink_path(&paths, "tool").exists());
707
708 cleanup(&root).expect("cleanup");
709 }
710}