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 return self
222 .upgrade_resolved(
223 package,
224 ResolvedUpgradeTarget::Branch {
225 branch: branch.to_string(),
226 head_commit,
227 },
228 trust_mode,
229 download_progress,
230 message_callback,
231 )
232 .await
233 .map(Some);
234 }
235
236 let Some(target) = target else {
237 return Ok(None);
238 };
239
240 self.upgrade_resolved(
241 package,
242 target,
243 trust_mode,
244 download_progress,
245 message_callback,
246 )
247 .await
248 .map(Some)
249 }
250
251 pub async fn upgrade_resolved<F, H>(
252 &self,
253 package: &Package,
254 target: ResolvedUpgradeTarget,
255 trust_mode: TrustMode,
256 download_progress: &mut Option<F>,
257 message_callback: &mut Option<H>,
258 ) -> Result<Package>
259 where
260 F: FnMut(u64, u64),
261 H: FnMut(&str),
262 {
263 let mut no_progress: Option<fn(PackageProgressEvent)> = None;
264 self.upgrade_resolved_with_progress(
265 package,
266 target,
267 trust_mode,
268 download_progress,
269 message_callback,
270 &mut no_progress,
271 )
272 .await
273 }
274
275 pub async fn upgrade_resolved_with_progress<F, H, P>(
276 &self,
277 package: &Package,
278 target: ResolvedUpgradeTarget,
279 trust_mode: TrustMode,
280 download_progress: &mut Option<F>,
281 message_callback: &mut Option<H>,
282 progress_callback: &mut Option<P>,
283 ) -> Result<Package>
284 where
285 F: FnMut(u64, u64),
286 H: FnMut(&str),
287 P: FnMut(PackageProgressEvent),
288 {
289 if package.is_pinned {
290 bail!("Package '{}' is pinned", package.name);
291 }
292
293 let had_desktop_integration = package.icon_path.is_some();
294
295 message!(
296 message_callback,
297 "{}",
298 style(format!("Upgrading '{}' ...", package.name)).cyan()
299 );
300
301 let original_install_path = package
302 .install_path
303 .as_ref()
304 .ok_or_else(|| {
305 anyhow::anyhow!("Package '{}' has no install path recorded", package.name)
306 })?
307 .clone();
308 let backup_path = Self::backup_path(self.paths, &original_install_path)?;
309
310 Self::remove_path_if_exists(&backup_path)?;
311
312 progress!(
313 progress_callback,
314 PackageProgressEvent::Phase(PackagePhase::CreatingSnapshot)
315 );
316 fs::rename(&original_install_path, &backup_path).context(format!(
317 "Failed to back up '{}' to '{}'",
318 original_install_path.display(),
319 backup_path.display()
320 ))?;
321
322 progress!(
324 progress_callback,
325 PackageProgressEvent::Phase(PackagePhase::RemovingRuntimeLinks)
326 );
327 if let Err(e) = self
328 .remover
329 .remove_runtime_integrations(package, message_callback)
330 {
331 let _ = fs::rename(&backup_path, &original_install_path);
332 let _ = self
333 .remover
334 .restore_runtime_integrations(package, message_callback);
335 return Err(e).context(format!(
336 "Failed to remove runtime integration for '{}'",
337 package.name
338 ));
339 }
340
341 let install_result = if package.install_type == InstallType::Build {
343 progress!(
344 progress_callback,
345 PackageProgressEvent::Phase(PackagePhase::RebuildingFromSource)
346 );
347 let (version_tag, branch, branch_head_commit) = match &target {
348 ResolvedUpgradeTarget::Release(release) => (Some(release.tag.clone()), None, None),
349 ResolvedUpgradeTarget::Branch {
350 branch,
351 head_commit,
352 } => (None, Some(branch.clone()), Some(head_commit.clone())),
353 };
354 let worker = BuildWorker::new(self.provider_manager, self.paths);
355 let build_result = {
356 let mut build_line_callback = Some(|line: &str| {
357 let line = line.trim();
358 if !line.is_empty() {
359 progress!(
360 progress_callback,
361 PackageProgressEvent::Warning(line.to_string())
362 );
363 }
364 });
365 worker
366 .build(
367 BuildRequest {
368 name: package.name.clone(),
369 repo_slug: package.repo_slug.clone(),
370 provider: package.provider.clone(),
371 base_url: package.base_url.clone(),
372 version_tag,
373 branch,
374 requested_profile: None,
375 script_action: BuildScriptAction::Upgrade,
376 },
377 package.channel.clone(),
378 &mut build_line_callback,
379 )
380 .await
381 };
382
383 match build_result {
384 Ok(output) => {
385 let mut install_pkg = package.clone();
386 install_pkg.build_branch = output.branch.clone();
387 install_pkg.build_commit = output.commit.or(branch_head_commit.clone());
388 progress!(
389 progress_callback,
390 PackageProgressEvent::Phase(PackagePhase::InstallingPackage)
391 );
392 let mut install_message_callback = Some(|line: &str| {
393 let line = line.trim();
394 if !line.is_empty() {
395 progress!(
396 progress_callback,
397 PackageProgressEvent::Warning(line.to_string())
398 );
399 message!(message_callback, "{}", line);
400 }
401 });
402 self.installer.install_local_artifact_files(
403 install_pkg,
404 &output.artifact_path,
405 output.version,
406 &mut install_message_callback,
407 )
408 }
409 Err(e) => {
410 Err(e).context(format!("Failed to rebuild '{}' from source", package.name))
411 }
412 }
413 } else {
414 let ResolvedUpgradeTarget::Release(release) = &target else {
415 bail!(
416 "Resolved branch target cannot be used for release package '{}'",
417 package.name
418 );
419 };
420 self.installer
421 .install_package_files(
422 package.clone(),
423 release,
424 trust_mode,
425 &self.trusted_keys,
426 download_progress,
427 message_callback,
428 progress_callback,
429 )
430 .await
431 };
432 let mut updated_package = match install_result {
433 Ok(updated_package) => updated_package,
434 Err(install_err) => {
435 progress!(
436 progress_callback,
437 PackageProgressEvent::Phase(PackagePhase::RollingBack)
438 );
439
440 return self.rollback_failed_upgrade(
441 FailedUpgradeRollback {
442 previous_package: package,
443 partially_installed_package: None,
444 original_install_path: &original_install_path,
445 backup_path: &backup_path,
446 failure_context: "Failed to install new version",
447 },
448 install_err,
449 message_callback,
450 );
451 }
452 };
453
454 if had_desktop_integration {
456 progress!(
457 progress_callback,
458 PackageProgressEvent::Phase(PackagePhase::CreatingRuntimeLinks)
459 );
460
461 if let Err(err) = self
462 .add_desktop_integration(&mut updated_package, message_callback)
463 .await
464 {
465 progress!(
466 progress_callback,
467 PackageProgressEvent::Phase(PackagePhase::RollingBack)
468 );
469 return self.rollback_failed_upgrade(
470 FailedUpgradeRollback {
471 previous_package: package,
472 partially_installed_package: Some(&updated_package),
473 original_install_path: &original_install_path,
474 backup_path: &backup_path,
475 failure_context: "Failed to restore desktop integration",
476 },
477 err.context("Failed to restore desktop integration"),
478 message_callback,
479 );
480 }
481 }
482
483 if let Err(err) =
484 Self::capture_successful_upgrade_rollback(self.paths, package, &backup_path)
485 {
486 progress!(
487 progress_callback,
488 PackageProgressEvent::Warning(format!(
489 "Warning: failed to capture rollback for '{}': {}",
490 package.name, err
491 ))
492 );
493 Self::remove_path_if_exists(&backup_path)
494 .context(format!("Failed to remove backup for '{}'", package.name))?;
495 }
496
497 Ok(updated_package)
498 }
499
500 async fn add_desktop_integration<H>(
501 &self,
502 updated_package: &mut Package,
503 message_callback: &mut Option<H>,
504 ) -> Result<()>
505 where
506 H: FnMut(&str),
507 {
508 #[cfg(target_os = "linux")]
509 let appimage_extractor =
510 AppImageExtractor::new().context("Failed to initialize appimage extractor")?;
511
512 #[cfg(target_os = "linux")]
513 let icon_manager = IconManager::new(self.paths, &appimage_extractor);
514 #[cfg(not(target_os = "linux"))]
515 let icon_manager = IconManager::new(self.paths);
516
517 #[cfg(target_os = "linux")]
518 let desktop_manager = DesktopManager::new(self.paths, &appimage_extractor);
519 #[cfg(not(target_os = "linux"))]
520 let desktop_manager = DesktopManager::new(self.paths);
521
522 let install_path = updated_package.install_path.clone().ok_or_else(|| {
523 anyhow::anyhow!(
524 "Package '{}' has no install path after upgrade",
525 updated_package.name
526 )
527 })?;
528
529 let icon_path = icon_manager
530 .add_icon(
531 &updated_package.name,
532 &install_path,
533 &updated_package.filetype,
534 message_callback,
535 )
536 .await
537 .context(format!("Failed to add icon for '{}'", updated_package.name))?;
538
539 updated_package.icon_path = icon_path;
540
541 let desktop_entry = DesktopEntry::from_package(updated_package);
542
543 desktop_manager
544 .create_entry(
545 &install_path,
546 &updated_package.filetype,
547 desktop_entry,
548 message_callback,
549 )
550 .await
551 .context(format!(
552 "Failed to create desktop entry for '{}'",
553 updated_package.name
554 ))?;
555
556 Ok(())
557 }
558
559 fn rollback_failed_upgrade<H>(
560 &self,
561 rollback: FailedUpgradeRollback<'_>,
562 failure: anyhow::Error,
563 message_callback: &mut Option<H>,
564 ) -> Result<Package>
565 where
566 H: FnMut(&str),
567 {
568 let cleanup_result = if let Some(partial) = rollback.partially_installed_package {
569 self.remover.remove_package_files(partial, message_callback)
570 } else {
571 Self::remove_path_if_exists(rollback.original_install_path)
572 };
573
574 if let Err(cleanup_err) = cleanup_result {
575 return Err(anyhow::anyhow!(
576 "{} for '{}': {}. Rollback failed while removing partial install: {}",
577 rollback.failure_context,
578 rollback.previous_package.name,
579 failure,
580 cleanup_err
581 ));
582 }
583
584 fs::rename(rollback.backup_path, rollback.original_install_path).context(format!(
585 "{} for '{}': {}. Rollback failed while restoring backup",
586 rollback.failure_context, rollback.previous_package.name, failure
587 ))?;
588
589 self.remover
590 .restore_runtime_integrations(rollback.previous_package, message_callback)
591 .context(format!(
592 "{} for '{}': {}. Rollback failed while restoring runtime links",
593 rollback.failure_context, rollback.previous_package.name, failure
594 ))?;
595
596 Err(failure).context(format!(
597 "{} for '{}' (previous version restored)",
598 rollback.failure_context, rollback.previous_package.name
599 ))
600 }
601}
602
603#[cfg(test)]
604mod tests {
605 use super::{FailedUpgradeRollback, PackageUpgrader};
606 use crate::models::common::enums::{Channel, Filetype, Provider};
607 use crate::models::upstream::Package;
608 use crate::providers::provider_manager::ProviderManager;
609 use crate::services::packaging::{PackageInstaller, PackageRemover};
610 use crate::services::trust::TrustedSignatureKeys;
611 use crate::utils::{static_paths::UpstreamPaths, test_support};
612 use std::path::{Path, PathBuf};
613 use std::time::{SystemTime, UNIX_EPOCH};
614 use std::{fs, io};
615
616 fn temp_root(name: &str) -> PathBuf {
617 let nanos = SystemTime::now()
618 .duration_since(UNIX_EPOCH)
619 .map(|d| d.as_nanos())
620 .unwrap_or(0);
621 std::env::temp_dir().join(format!("upstream-upgrader-test-{name}-{nanos}"))
622 }
623
624 fn cleanup(path: &Path) -> io::Result<()> {
625 fs::remove_dir_all(path)
626 }
627
628 fn expected_symlink_path(paths: &UpstreamPaths, name: &str) -> PathBuf {
629 let base = paths.integration.symlinks_dir.join(name);
630 #[cfg(windows)]
631 {
632 return base.with_extension("exe");
633 }
634 #[cfg(not(windows))]
635 {
636 base
637 }
638 }
639
640 fn test_paths(root: &Path) -> crate::utils::static_paths::UpstreamPaths {
641 test_support::upstream_paths(root)
642 }
643
644 fn test_package(name: &str, install_path: PathBuf) -> Package {
645 let mut package = Package::with_defaults(
646 name.to_string(),
647 format!("owner/{name}"),
648 Filetype::Binary,
649 None,
650 None,
651 Channel::Stable,
652 Provider::Github,
653 None,
654 );
655 package.install_path = Some(install_path.clone());
656 package.exec_path = Some(install_path);
657 package
658 }
659
660 #[test]
661 fn remove_path_if_exists_handles_files_and_directories() {
662 let root = temp_root("remove");
663 let file = root.join("f.bin");
664 let dir = root.join("d");
665 fs::create_dir_all(&dir).expect("create dir");
666 fs::write(&file, b"content").expect("write file");
667
668 PackageUpgrader::remove_path_if_exists(&file).expect("remove file");
669 PackageUpgrader::remove_path_if_exists(&dir).expect("remove dir");
670 PackageUpgrader::remove_path_if_exists(&root.join("missing")).expect("ignore missing");
671
672 assert!(!file.exists());
673 assert!(!dir.exists());
674
675 cleanup(&root).expect("cleanup");
676 }
677
678 #[test]
679 fn rollback_failed_upgrade_removes_partial_install_and_restores_previous_binary() {
680 let root = temp_root("rollback-desktop-failure");
681 let paths = test_paths(&root);
682 fs::create_dir_all(&paths.install.binaries_dir).expect("create binaries dir");
683 fs::create_dir_all(&paths.install.tmp_dir).expect("create tmp dir");
684 fs::create_dir_all(&paths.integration.symlinks_dir).expect("create symlinks dir");
685
686 let install_path = paths.install.binaries_dir.join("tool");
687 let backup_path = paths.install.tmp_dir.join("tool.old");
688 fs::write(&install_path, b"new").expect("write partial new binary");
689 fs::write(&backup_path, b"old").expect("write backup binary");
690
691 let previous = test_package("tool", install_path.clone());
692 let partial = test_package("tool", install_path.clone());
693 let provider_manager = ProviderManager::new(None, None, None).expect("provider manager");
694 let installer = PackageInstaller::new(&provider_manager, &paths).expect("installer");
695 let remover = PackageRemover::new(&paths);
696 let upgrader = PackageUpgrader::new(
697 &provider_manager,
698 installer,
699 remover,
700 &paths,
701 TrustedSignatureKeys::default(),
702 );
703 let mut msg = Some(|_: &str| {});
704
705 let err = upgrader
706 .rollback_failed_upgrade(
707 FailedUpgradeRollback {
708 previous_package: &previous,
709 partially_installed_package: Some(&partial),
710 original_install_path: &install_path,
711 backup_path: &backup_path,
712 failure_context: "Failed to restore desktop integration",
713 },
714 anyhow::anyhow!("desktop failed"),
715 &mut msg,
716 )
717 .expect_err("rollback helper returns original failure");
718
719 assert!(err.to_string().contains("previous version restored"));
720 assert_eq!(
721 fs::read(&install_path).expect("read restored binary"),
722 b"old"
723 );
724 assert!(!backup_path.exists());
725 assert!(expected_symlink_path(&paths, "tool").exists());
726
727 cleanup(&root).expect("cleanup");
728 }
729}