From 1e6a374ce458d33589e9171ea9e73102b347510d Mon Sep 17 00:00:00 2001
From: Shantanu <shantanu@sfcompute.com>
Date: Thu, 24 Jul 2025 22:23:34 -0700
Subject: [PATCH] support macos universal binaries

---
 checksum/checksum.go           | 21 ++++++++++++---------
 release/asset/download.go      | 19 +++++++++++--------
 release/asset/download_test.go |  7 ++++++-
 upgrade.go                     |  6 +++---
 4 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/checksum/checksum.go b/checksum/checksum.go
index 9a2599c..741ed25 100644
--- a/checksum/checksum.go
+++ b/checksum/checksum.go
@@ -118,9 +118,9 @@ func WithOS(os string) ValidatorOption {
 	}
 }
 
-var fallbackArchMap = map[string]string{
-	"amd64": "x86_64",
-	"386":   "i386",
+var fallbackArchMap = map[string][]string{
+	"amd64": {"x86_64", "all"},
+	"386":   {"i386", "all"},
 }
 
 func WithArch(a string) ValidatorOption {
@@ -152,14 +152,17 @@ func (v *validator) IsCheckSumValid(ctx context.Context, binary string, info *In
 }
 
 func (v *validator) tryFallbackArch(binary string, info *Info, downloadedChecksum string) bool {
-	arch, ok := fallbackArchMap[v.arch]
+	archs, ok := fallbackArchMap[v.arch]
 	if !ok {
 		return false
 	}
-	key := fmt.Sprintf("%s_%s_%s", binary, v.os, arch)
-	expectedChecksum, ok := info.Checksums[key]
-	if !ok {
-		return false
+
+	for _, arch := range archs {
+		key := fmt.Sprintf("%s_%s_%s", binary, v.os, arch)
+		expectedChecksum, ok := info.Checksums[key]
+		if ok {
+			return expectedChecksum == downloadedChecksum
+		}
 	}
-	return expectedChecksum == downloadedChecksum
+	return false
 }
diff --git a/release/asset/download.go b/release/asset/download.go
index bfdb909..e1ca9ab 100644
--- a/release/asset/download.go
+++ b/release/asset/download.go
@@ -30,7 +30,7 @@ type Info struct {
 type downloader struct {
 	os                 string
 	arch               string
-	lookupArchFallback map[string]string
+	lookupArchFallback map[string][]string
 	executablePath     string
 }
 
@@ -50,7 +50,7 @@ func WithArch(arch string) AssetDownloadOpt {
 	}
 }
 
-func WithLookupArchFallback(lookupArchFallback map[string]string) AssetDownloadOpt {
+func WithLookupArchFallback(lookupArchFallback map[string][]string) AssetDownloadOpt {
 	return func(d *downloader) {
 		d.lookupArchFallback = lookupArchFallback
 	}
@@ -78,19 +78,22 @@ func (d *downloader) DownloadAsset(ctx context.Context, assets []release.Asset)
 		return d.downloadAsset(ctx, asset.BrowserDownloadURL)
 	}
 	// if asset not found, try a fallback. e.g amd64 -> x86_64
-	if d.lookupArchFallback == nil || len(d.lookupArchFallback) == 0 {
+	if len(d.lookupArchFallback) == 0 {
 		return nil, nil, fmt.Errorf("%w: os:%s arch:%s", ErrNoAsset, d.os, d.arch)
 	}
 
-	fallbackArch, ok := d.lookupArchFallback[d.arch]
+	fallbackArchs, ok := d.lookupArchFallback[d.arch]
 	if !ok {
 		return nil, nil, fmt.Errorf("%w: os:%s arch:%s", ErrNoAsset, d.os, d.arch)
 	}
 
-	fallbackSuffix := d.os + "_" + fallbackArch
-	asset, found = d.assetForSuffix(assets, fallbackSuffix)
-	if found {
-		return d.downloadAsset(ctx, asset.BrowserDownloadURL)
+	// Try to find an asset for each fallback architecture
+	for _, fallbackArch := range fallbackArchs {
+		fallbackSuffix := d.os + "_" + fallbackArch
+		asset, found = d.assetForSuffix(assets, fallbackSuffix)
+		if found {
+			return d.downloadAsset(ctx, asset.BrowserDownloadURL)
+		}
 	}
 	return nil, nil, fmt.Errorf("%w: os:%s arch:%s", ErrNoAsset, d.os, d.arch)
 }
diff --git a/release/asset/download_test.go b/release/asset/download_test.go
index cf53b5c..c24cd51 100644
--- a/release/asset/download_test.go
+++ b/release/asset/download_test.go
@@ -83,7 +83,12 @@ func TestAssetDownloader(t *testing.T) {
 			assert.Nil(t, cleanupFn)
 		})
 		t.Run("DownloadSucceedsWithFallback", func(t *testing.T) {
-			downloader := NewAssetDownloader(executablePath, WithOS("os"), WithArch("amd64"), WithLookupArchFallback(map[string]string{"amd64": "x86_64"}))
+			downloader := NewAssetDownloader(executablePath,
+				WithOS("os"),
+				WithArch("amd64"),
+				WithLookupArchFallback(
+					map[string][]string{"amd64": {"all", "x86_64"}},
+				))
 			asset, cleanupFn, err := downloader.DownloadAsset(ctx, []release.Asset{
 				{BrowserDownloadURL: srv.URL + "/download_os_x86_64"},
 			})
diff --git a/upgrade.go b/upgrade.go
index a3b34aa..f8ab5cd 100644
--- a/upgrade.go
+++ b/upgrade.go
@@ -57,9 +57,9 @@ func NewUpgrader(owner string, repo string, executablePath string, opts ...Opt)
 		owner:          owner,
 		executablePath: executablePath,
 		releaseGetter:  release.NewReleaseGetter(repo, owner),
-		assetDownloader: asset.NewAssetDownloader(executablePath, asset.WithLookupArchFallback(map[string]string{
-			"amd64": "x86_64",
-			"386":   "i86",
+		assetDownloader: asset.NewAssetDownloader(executablePath, asset.WithLookupArchFallback(map[string][]string{
+			"amd64": {"x86_64", "all"},
+			"386":   {"i86", "all"},
 		})),
 		checksumDownloader: checksum.NewCheckSumDownloader(),
 		checksumValidator:  checksum.NewCheckSumValidator(),