From 02791a4d4232897645519c735298abb3d2ff728b Mon Sep 17 00:00:00 2001
From: Vladimir Stoilov <vladimir@safing.io>
Date: Thu, 10 Oct 2024 13:02:13 +0300
Subject: [PATCH] [WIP] Add fallback on corrupted install

---
 service/updates/module.go   | 29 +++++++++++++++++++++++------
 service/updates/registry.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/service/updates/module.go b/service/updates/module.go
index 2e583a84..f33c830c 100644
--- a/service/updates/module.go
+++ b/service/updates/module.go
@@ -13,9 +13,10 @@ import (
 )
 
 const (
-	updateTaskRepeatDuration      = 1 * time.Hour
-	updateAvailableNotificationID = "updates:update-available"
-	updateFailedNotificationID    = "updates:update-failed"
+	updateTaskRepeatDuration          = 1 * time.Hour
+	updateAvailableNotificationID     = "updates:update-available"
+	updateFailedNotificationID        = "updates:update-failed"
+	corruptInstallationNotificationID = "updates:corrupt-installation"
 
 	// ResourceUpdateEvent is emitted every time the
 	// updater successfully performed a resource update.
@@ -59,6 +60,8 @@ type Updates struct {
 	autoApply    bool
 	needsRestart bool
 
+	corruptedInstallation bool
+
 	isUpdateRunning *abool.AtomicBool
 
 	instance instance
@@ -96,7 +99,14 @@ func New(instance instance, name string, index UpdateIndex) (*Updates, error) {
 	var err error
 	module.registry, err = CreateRegistry(index)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create registry: %w", err)
+		// Installation is corrupt, set flag and fall back to folder scanning for artifacts discovery.
+		log.Criticalf("updates: failed to create registry: %s (falling back to folder scanning)", err)
+		module.corruptedInstallation = true
+
+		module.registry, err = CreateRegistryFromFolder(index)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	module.downloader = CreateDownloader(index)
@@ -194,14 +204,16 @@ func (u *Updates) applyUpdates(downloader Downloader, force bool) error {
 	currentBundle := u.registry.bundle
 	downloadBundle := downloader.bundle
 
-	if !force {
+	if !force && u.registry.version != nil {
 		if u.downloader.version.LessThanOrEqual(u.registry.version) {
 			// No new version, silently return.
 			return nil
 		}
 	}
+	if currentBundle != nil {
+		log.Infof("update: starting update: %s %s -> %s", currentBundle.Name, currentBundle.Version, downloadBundle.Version)
+	}
 
-	log.Infof("update: starting update: %s %s -> %s", currentBundle.Name, currentBundle.Version, downloadBundle.Version)
 	err := u.registry.performRecoverableUpgrade(downloader.dir, downloader.indexFile)
 	if err != nil {
 		// Notify the user that update failed.
@@ -247,6 +259,11 @@ func (u *Updates) Start() error {
 		_ = u.downloader.deleteUnfinishedDownloads()
 		return nil
 	})
+
+	if u.corruptedInstallation {
+		notifications.NotifyError(corruptInstallationNotificationID, "Corrupted installation. Reinstall the software.", "")
+	}
+
 	u.updateCheckWorkerMgr.Go()
 
 	return nil
diff --git a/service/updates/registry.go b/service/updates/registry.go
index 9b07a742..98df62a0 100644
--- a/service/updates/registry.go
+++ b/service/updates/registry.go
@@ -5,6 +5,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"strings"
 
 	semver "github.com/hashicorp/go-version"
 
@@ -54,6 +55,35 @@ func CreateRegistry(index UpdateIndex) (Registry, error) {
 	return registry, nil
 }
 
+func CreateRegistryFromFolder(index UpdateIndex) (Registry, error) {
+	registry := Registry{
+		dir:      index.Directory,
+		purgeDir: index.PurgeDirectory,
+		files:    make(map[string]File),
+	}
+
+	files, err := os.ReadDir(index.Directory)
+	if err != nil {
+		return Registry{}, nil
+	}
+	for _, file := range files {
+		// Skip dirs
+		if file.IsDir() {
+			continue
+		}
+
+		// Skip the uninstaller. (Windows)
+		if strings.Contains(strings.ToLower(file.Name()), "uninstall") {
+			continue
+		}
+
+		artifactPath := filepath.Join(registry.dir, file.Name())
+		registry.files[file.Name()] = File{id: file.Name(), path: artifactPath, version: "", sha256: ""}
+	}
+
+	return registry, nil
+}
+
 func (r *Registry) performUpgrade(downloadDir string, indexFile string) error {
 	// Make sure provided update is valid
 	indexFilepath := filepath.Join(downloadDir, indexFile)