From 4e00ee12d17717e716eda54f91905651eff388d6 Mon Sep 17 00:00:00 2001 From: ZacxDev Date: Thu, 25 Jun 2026 15:14:27 -0500 Subject: [PATCH 1/2] feat(cli): cached non-blocking update notice + `civitai upgrade` self-update MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds two related features in one PR: PART 1 — daily, cached, non-blocking "new version available" notice - A root PersistentPostRun hook prints at most ONE dim stderr line after any successful command. It never does a synchronous network call: it reads a cache (~/.config/civitai/update-check.json) and, when stale, spawns a DETACHED `civitai __update-check` (hidden subcommand) that fetches the latest release and rewrites the cache. The current run uses the cached value; the refresh lands for next time. First run (no cache) only kicks off the refresh. - Refresh at most once / 24h (last_check); notice shown at most once / 24h (last_notified) even while behind. Corrupt/missing cache => empty, fail-silent. - Suppressed when: stderr is not a TTY, CI env is set, --no-update-check / CIVITAI_NO_UPDATE_CHECK, or the command is version/upgrade/completion/help/ __update-check/__complete*. Only fires when current parses and latest > current. - stderr only — never pollutes stdout, so pipes/scripts are unaffected. - Reuses the existing fetchLatestRelease / semver helpers from update_check.go; `version` keeps its own explicit synchronous check (no double-notify). PART 2 — `civitai upgrade` self-update - Resolves the latest release (unauthenticated GitHub, no token ever sent). Already >= latest and not --force => "already up to date" no-op. - Homebrew detection: if the resolved executable lives under a brew path (/Cellar/, /Caskroom/, /opt/homebrew, /usr/local/Homebrew|Cellar, /home/linuxbrew/.linuxbrew), prints the brew upgrade command instead of self-replacing (--force overrides). - Otherwise downloads the platform tarball + checksums.txt, VERIFIES the tarball SHA-256 against checksums.txt and ABORTS on mismatch (binary left untouched), extracts the binary, and atomically replaces the running executable via github.com/minio/selfupdate. Permission-denied => clear sudo/brew/go-install guidance, non-zero exit, no half-written binary. Detaching uses a small build-tagged platform split (unix Setpgid / windows CREATE_NEW_PROCESS_GROUP); the parent never Waits. Spawn + apply + executable-path + TTY are behind injectable seams so tests assert behavior without forking or replacing the test binary. Co-Authored-By: Claude Opus 4.8 (1M context) --- go.mod | 8 +- go.sum | 28 +- internal/cmd/root.go | 15 ++ internal/cmd/spawn_unix.go | 35 +++ internal/cmd/spawn_windows.go | 31 +++ internal/cmd/update_cache.go | 137 ++++++++++ internal/cmd/update_check.go | 7 + internal/cmd/update_notice.go | 210 +++++++++++++++ internal/cmd/update_notice_test.go | 393 ++++++++++++++++++++++++++++ internal/cmd/upgrade.go | 318 +++++++++++++++++++++++ internal/cmd/upgrade_test.go | 397 +++++++++++++++++++++++++++++ internal/cmd/version.go | 5 +- 12 files changed, 1577 insertions(+), 7 deletions(-) create mode 100644 internal/cmd/spawn_unix.go create mode 100644 internal/cmd/spawn_windows.go create mode 100644 internal/cmd/update_cache.go create mode 100644 internal/cmd/update_notice.go create mode 100644 internal/cmd/update_notice_test.go create mode 100644 internal/cmd/upgrade.go create mode 100644 internal/cmd/upgrade_test.go diff --git a/go.mod b/go.mod index 99ab5ad..dcd3dda 100644 --- a/go.mod +++ b/go.mod @@ -1,16 +1,19 @@ module github.com/civitai/cli -go 1.25 +go 1.25.0 require ( + github.com/minio/selfupdate v0.6.0 github.com/santhosh-tekuri/jsonschema/v6 v6.0.1 github.com/spf13/cobra v1.8.1 github.com/spf13/viper v1.19.0 + golang.org/x/term v0.44.0 golang.org/x/text v0.14.0 gopkg.in/yaml.v3 v3.0.1 ) require ( + aead.dev/minisign v0.2.0 // indirect github.com/fsnotify/fsnotify v1.7.0 // indirect github.com/hashicorp/hcl v1.0.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect @@ -26,7 +29,8 @@ require ( github.com/subosito/gotenv v1.6.0 // indirect go.uber.org/atomic v1.9.0 // indirect go.uber.org/multierr v1.9.0 // indirect + golang.org/x/crypto v0.21.0 // indirect golang.org/x/exp v0.0.0-20230905200255-921286631fa9 // indirect - golang.org/x/sys v0.18.0 // indirect + golang.org/x/sys v0.46.0 // indirect gopkg.in/ini.v1 v1.67.0 // indirect ) diff --git a/go.sum b/go.sum index 40c53e4..2a3eff3 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +aead.dev/minisign v0.2.0 h1:kAWrq/hBRu4AARY6AlciO83xhNnW9UaC8YipS2uhLPk= +aead.dev/minisign v0.2.0/go.mod h1:zdq6LdSd9TbuSxchxwhpA9zEb9YXcVGoE8JakuiGaIQ= github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= @@ -21,6 +23,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0VQdvPDY= github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0= +github.com/minio/selfupdate v0.6.0 h1:i76PgT0K5xO9+hjzKcacQtO7+MjJ4JKA8Ak8XQ9DDwU= +github.com/minio/selfupdate v0.6.0/go.mod h1:bO02GTIPCMQFTEvE5h4DjYB58bCoZ35XLeBf0buTDdM= github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY= github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo= github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM= @@ -65,12 +69,32 @@ go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE= go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc= go.uber.org/multierr v1.9.0 h1:7fIwc/ZtS0q++VgcfqFDxSBZVv/Xo49/SYnDFupUwlI= go.uber.org/multierr v1.9.0/go.mod h1:X2jQV1h+kxSjClGpnseKVIxpmcjrj7MNnI0bnlfKTVQ= +golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= +golang.org/x/crypto v0.0.0-20210220033148-5ea612d1eb83/go.mod h1:jdWPYTVW3xRLrWPugEBEK3UY2ZEsg3UU495nc5E+M+I= +golang.org/x/crypto v0.0.0-20211209193657-4570a0811e8b/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4= +golang.org/x/crypto v0.21.0 h1:X31++rzVUdKhX5sWmSOFZxx8UW/ldWx55cbf08iNAMA= +golang.org/x/crypto v0.21.0/go.mod h1:0BP7YvVV9gBbVKyeTG0Gyn+gZm94bibOW5BjDEYAOMs= golang.org/x/exp v0.0.0-20230905200255-921286631fa9 h1:GoHiUyI/Tp2nVkLI2mCxVkOjsbSXD66ic0XW0js0R9g= golang.org/x/exp v0.0.0-20230905200255-921286631fa9/go.mod h1:S2oDrQGGwySpoQPVqRShND87VCbxmc6bL1Yd2oYrm6k= -golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4= -golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= +golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= +golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210228012217-479acdf4ea46/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= +golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw= +golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= +golang.org/x/term v0.44.0 h1:0rLvDRCtNj0gZkyIXhCyOb2OAzEhLVqc4B+hrsBhrmc= +golang.org/x/term v0.44.0/go.mod h1:7ze4MdzUzLXpSAoFP1H0bOI9aXDqveSvatT5vKcFh2Y= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ= golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo= gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/internal/cmd/root.go b/internal/cmd/root.go index 335840d..bd7f9d3 100644 --- a/internal/cmd/root.go +++ b/internal/cmd/root.go @@ -108,6 +108,7 @@ func isDevDefault(v string) bool { // NewRootCmd builds the root command with all subcommands attached. func NewRootCmd() *cobra.Command { + var noUpdateCheck bool root := &cobra.Command{ Use: "civitai", Short: "Civitai CLI — author and ship App Blocks", @@ -133,15 +134,29 @@ Get started: SilenceUsage: true, SilenceErrors: true, Version: version, + // PersistentPostRun fires after ANY subcommand's RunE. It prints at most + // one cached "new version available" line to stderr and (when the cache + // is stale) kicks off a detached background refresh. It is best-effort + // and NEVER blocks the command — see maybeNotifyUpdate. + PersistentPostRun: func(cmd *cobra.Command, args []string) { + maybeNotifyUpdate(cmd.ErrOrStderr(), cmd.Name(), noUpdateCheck) + }, } root.SetVersionTemplate("civitai {{.Version}}\n") + // A persistent flag so every command honours --no-update-check, and the + // post-run hook can read its resolved value. + root.PersistentFlags().BoolVar(&noUpdateCheck, "no-update-check", false, + "skip the background check for a newer release (also via CIVITAI_NO_UPDATE_CHECK)") + root.AddCommand(newAppCmd()) root.AddCommand(newLoginCmd()) root.AddCommand(newWhoAmICmd()) root.AddCommand(newBuzzCmd()) root.AddCommand(newVersionCmd()) + root.AddCommand(newUpgradeCmd()) root.AddCommand(newCompletionCmd()) + root.AddCommand(newUpdateCheckCmd()) return root } diff --git a/internal/cmd/spawn_unix.go b/internal/cmd/spawn_unix.go new file mode 100644 index 0000000..1851c7c --- /dev/null +++ b/internal/cmd/spawn_unix.go @@ -0,0 +1,35 @@ +//go:build !windows + +package cmd + +import ( + "os" + "os/exec" + "syscall" +) + +// detachAndStart launches cmd as a detached background process and returns +// immediately WITHOUT waiting for it. On unix we put it in its own process +// group (Setpgid) so it survives the parent exiting and isn't killed by a +// signal sent to the parent's group. stdio is redirected to /dev/null so it +// can never write to the user's terminal. +// +// The returned error reflects only the Start() (fork/exec) failure; we never +// Wait on the child, so the parent returns instantly. +func detachAndStart(cmd *exec.Cmd) error { + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + devnull, err := os.OpenFile(os.DevNull, os.O_RDWR, 0) + if err == nil { + cmd.Stdin = devnull + cmd.Stdout = devnull + cmd.Stderr = devnull + } + startErr := cmd.Start() + // The child has inherited the fd; close our copy so we don't leak it. + if devnull != nil { + _ = devnull.Close() + } + // Deliberately NOT calling cmd.Wait() — fire and forget. The child reaps as + // an orphan once the parent exits. + return startErr +} diff --git a/internal/cmd/spawn_windows.go b/internal/cmd/spawn_windows.go new file mode 100644 index 0000000..b81b413 --- /dev/null +++ b/internal/cmd/spawn_windows.go @@ -0,0 +1,31 @@ +//go:build windows + +package cmd + +import ( + "os" + "os/exec" + "syscall" +) + +// detachAndStart launches cmd as a detached background process on Windows and +// returns immediately WITHOUT waiting. CREATE_NEW_PROCESS_GROUP detaches the +// child from the parent's console process group so it survives the parent +// exiting; stdio is redirected to NUL so it can never write to the console. +func detachAndStart(cmd *exec.Cmd) error { + cmd.SysProcAttr = &syscall.SysProcAttr{ + CreationFlags: syscall.CREATE_NEW_PROCESS_GROUP, + } + devnull, err := os.OpenFile(os.DevNull, os.O_RDWR, 0) + if err == nil { + cmd.Stdin = devnull + cmd.Stdout = devnull + cmd.Stderr = devnull + } + startErr := cmd.Start() + if devnull != nil { + _ = devnull.Close() + } + // Fire and forget — never Wait. + return startErr +} diff --git a/internal/cmd/update_cache.go b/internal/cmd/update_cache.go new file mode 100644 index 0000000..23a568d --- /dev/null +++ b/internal/cmd/update_cache.go @@ -0,0 +1,137 @@ +package cmd + +import ( + "encoding/json" + "os" + "path/filepath" + "time" + + "github.com/civitai/cli/internal/config" +) + +// updateCacheName is the cache file (sibling of config.yaml) that backs the +// non-blocking daily update notice. It records when we last checked GitHub, the +// latest version we saw, and when we last actually showed the notice. +const updateCacheName = "update-check.json" + +// updateCacheTTL is how often the background refresh is allowed to hit GitHub. +const updateCacheTTL = 24 * time.Hour + +// updateNotifyTTL throttles how often the notice is shown — at most once a day +// even while the user stays behind, so it's a gentle nudge, not nagware. +const updateNotifyTTL = 24 * time.Hour + +// updateCache is the on-disk JSON state for the cached update notice. +type updateCache struct { + LastCheck time.Time `json:"last_check"` + LatestVersion string `json:"latest_version"` + LastNotified time.Time `json:"last_notified"` +} + +// updateCachePath returns ~/.config/civitai/update-check.json (the same dir as +// config.yaml). Returns an error only if the config dir can't be resolved. +func updateCachePath() (string, error) { + dir, err := config.Dir() + if err != nil { + return "", err + } + return filepath.Join(dir, updateCacheName), nil +} + +// readUpdateCache loads the cache. A missing or corrupt file is treated as an +// empty cache with no error — the notice path must never fail a command. +func readUpdateCache(path string) updateCache { + b, err := os.ReadFile(path) + if err != nil { + return updateCache{} + } + var c updateCache + if err := json.Unmarshal(b, &c); err != nil { + return updateCache{} // corrupt => empty, fail-silent. + } + return c +} + +// writeUpdateCache persists the cache atomically (0600 temp + rename) into the +// config dir. The dir is created if needed. Errors are returned to the caller, +// but the only caller (the detached refresh) ignores them — a failed write just +// means the next run refreshes again. +func writeUpdateCache(path string, c updateCache) error { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o700); err != nil { + return err + } + b, err := json.Marshal(c) + if err != nil { + return err + } + tmp, err := os.CreateTemp(dir, ".update-check-*.json.tmp") + if err != nil { + return err + } + tmpName := tmp.Name() + defer func() { _ = os.Remove(tmpName) }() + if _, err := tmp.Write(b); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + return os.Rename(tmpName, path) +} + +// cacheIsFresh reports whether the cached check is younger than the TTL, so no +// background refresh is needed. A zero LastCheck (empty cache) is never fresh. +func cacheIsFresh(c updateCache, now time.Time) bool { + if c.LastCheck.IsZero() { + return false + } + return now.Sub(c.LastCheck) < updateCacheTTL +} + +// noticeDecision describes what the post-run hook should do, derived purely from +// the cache + current version + the clock (no I/O, so it's unit-testable). +type noticeDecision struct { + // show is true when a one-line "newer version" notice should be printed. + show bool + // notice is the text to print (only meaningful when show is true). + notice string + // refresh is true when a background GitHub refresh should be kicked off. + refresh bool +} + +// decideNotice computes whether to show the notice and/or refresh the cache, +// from the cached state and the resolved current version. Pure function. +// +// Rules: +// - refresh when the cache is stale (older than updateCacheTTL or empty). +// - show the notice only when: the cached latest parses, current parses, +// latest > current, AND we haven't notified within updateNotifyTTL. +func decideNotice(c updateCache, current string, now time.Time) noticeDecision { + d := noticeDecision{refresh: !cacheIsFresh(c, now)} + + if c.LatestVersion == "" || !isParseableVersion(c.LatestVersion) { + return d // no usable cached value yet (e.g. first run) => no notice. + } + // Only nag when we KNOW the current version and it's genuinely behind. + if !isParseableVersion(current) { + return d + } + if compareVersions(current, c.LatestVersion) != -1 { + return d // up to date or ahead. + } + // Throttle: at most one notice per updateNotifyTTL. + if !c.LastNotified.IsZero() && now.Sub(c.LastNotified) < updateNotifyTTL { + return d + } + d.show = true + d.notice = formatUpdateNotice(c.LatestVersion, current) + return d +} + +// formatUpdateNotice renders the single dim stderr line. +func formatUpdateNotice(latest, current string) string { + return "A new version of civitai is available: " + latest + + " (you have " + current + "). Run 'civitai upgrade' to update." +} diff --git a/internal/cmd/update_check.go b/internal/cmd/update_check.go index a39eb9e..5c63879 100644 --- a/internal/cmd/update_check.go +++ b/internal/cmd/update_check.go @@ -23,6 +23,13 @@ var latestReleaseURL = "https://api.github.com/repos/civitai/cli/releases/latest // version` never hangs on a slow/offline network. const updateCheckTimeout = 2500 * time.Millisecond +// contextWithUpdateTimeout returns a context bounded by updateCheckTimeout for +// the unauthenticated GitHub round-trip. Shared by the synchronous `version` +// check and the detached background refresh. +func contextWithUpdateTimeout() (context.Context, context.CancelFunc) { + return context.WithTimeout(context.Background(), updateCheckTimeout) +} + // updateCheckDisabled reports whether the GitHub update check should be skipped: // the --no-update-check flag (noFlag) or a non-empty CIVITAI_NO_UPDATE_CHECK env // var both opt out. diff --git a/internal/cmd/update_notice.go b/internal/cmd/update_notice.go new file mode 100644 index 0000000..e1526c3 --- /dev/null +++ b/internal/cmd/update_notice.go @@ -0,0 +1,210 @@ +package cmd + +import ( + "io" + "os" + "os/exec" + "time" + + "github.com/spf13/cobra" + "golang.org/x/term" +) + +// updateCheckHiddenCmd is the name of the hidden subcommand that performs the +// background GitHub refresh and writes the cache. The post-run hook spawns a +// detached copy of ourselves running this command. +const updateCheckHiddenCmd = "__update-check" + +// noticeExcludedCommands are commands for which the cached update notice (and +// the background refresh spawn) are suppressed: +// - version: does its OWN synchronous check; double-notifying is noise. +// - upgrade: the whole point of running it; a notice is redundant. +// - __update-check: the refresh worker itself; must stay silent. +// - completion: emits shell-eval'd script to stdout — keep it pristine. +// - help: help output shouldn't carry a network-derived footer. +var noticeExcludedCommands = map[string]bool{ + "version": true, + "upgrade": true, + updateCheckHiddenCmd: true, + "completion": true, + "help": true, +} + +// ciEnvVars are environment variables whose presence indicates a CI/automation +// context where an interactive update notice is unwanted. +var ciEnvVars = []string{ + "CI", + "CONTINUOUS_INTEGRATION", + "GITHUB_ACTIONS", + "GITLAB_CI", + "BUILDKITE", + "CIRCLECI", + "TF_BUILD", // Azure Pipelines +} + +// stderrIsTerminal is a seam so tests can force the TTY check. By default it +// checks whether os.Stderr is a real terminal. +var stderrIsTerminal = func() bool { + return term.IsTerminal(int(os.Stderr.Fd())) +} + +// requestRefresh is a seam: in production it spawns the detached refresh +// process; tests override it to record that a refresh was requested without +// actually forking. It must never block. +var requestRefresh = spawnDetachedRefresh + +// runningInCI reports whether any known CI env var is set. +func runningInCI() bool { + for _, k := range ciEnvVars { + if os.Getenv(k) != "" { + return true + } + } + return false +} + +// invokedCommandName returns the name of the leaf command cobra is about to run +// (or just ran), used for the exclusion check. For the completion machinery the +// command name starts with "__complete"; we treat any such name as excluded. +func noticeSuppressedForCommand(name string) bool { + if name == "" { + return true + } + if len(name) >= 2 && name[0] == '_' && name[1] == '_' { + // __complete, __completeNoDesc, __update-check, etc. — all internal. + return true + } + return noticeExcludedCommands[name] +} + +// updateNoticeGated reports whether the notice+refresh should be entirely +// suppressed for this invocation. It does NOT touch the network or disk. +// +// Suppressed when ANY of: +// - the user opted out (--no-update-check / CIVITAI_NO_UPDATE_CHECK), OR +// - stderr is not a TTY, OR +// - we're in CI, OR +// - the invoked command is excluded (version/upgrade/completion/help/internal). +func updateNoticeGated(cmdName string, noFlag bool) bool { + if updateCheckDisabled(noFlag) { + return true + } + if noticeSuppressedForCommand(cmdName) { + return true + } + if runningInCI() { + return true + } + if !stderrIsTerminal() { + return true + } + return false +} + +// maybeNotifyUpdate is the PersistentPostRun body. It is best-effort and MUST +// NOT block: it reads the cache, prints the (possibly stale) cached notice if +// warranted, kicks off a detached background refresh when the cache is stale, +// and persists the last-notified timestamp. Any error is swallowed. +// +// errW is the command's stderr writer; cmdName is the leaf command's name; +// noFlag is the resolved --no-update-check value. +func maybeNotifyUpdate(errW io.Writer, cmdName string, noFlag bool) { + if updateNoticeGated(cmdName, noFlag) { + return + } + + path, err := updateCachePath() + if err != nil { + return + } + cache := readUpdateCache(path) + now := time.Now() + d := decideNotice(cache, version, now) + + if d.refresh { + // Fire-and-forget; never wait. A failure here is irrelevant — next run + // just tries again. + _ = requestRefresh() + } + + if d.show { + // stderr ONLY — never stdout, so pipes/scripts are unaffected. + printDim(errW, d.notice) + // Record that we notified so we don't repeat within updateNotifyTTL. + cache.LastNotified = now + _ = writeUpdateCache(path, cache) + } +} + +// printDim writes s as a dim line to w, framed by a leading newline so it +// visually separates from the command's own output. ANSI dim is only emitted +// when w is the real terminal (we already gate on stderr being a TTY before +// calling, but keep the escape minimal and safe). +func printDim(w io.Writer, s string) { + const ( + dim = "\033[2m" + reset = "\033[0m" + ) + _, _ = io.WriteString(w, "\n"+dim+s+reset+"\n") +} + +// spawnDetachedRefresh launches `civitai __update-check` as a detached +// background process that fetches the latest release and rewrites the cache, +// then returns immediately. The current process never waits on it. +func spawnDetachedRefresh() error { + exe, err := os.Executable() + if err != nil { + return err + } + cmd := exec.Command(exe, updateCheckHiddenCmd) + // Propagate the test/override URL into the child so it checks the same + // endpoint. (In production this env var is unset and the child uses the + // default GitHub URL.) + cmd.Env = os.Environ() + return detachAndStart(cmd) +} + +// newUpdateCheckCmd builds the hidden background-refresh subcommand. It produces +// NO output, calls fetchLatestRelease, and writes the cache. It always exits 0 +// (best-effort) so a detached failure never surfaces anywhere. +func newUpdateCheckCmd() *cobra.Command { + return &cobra.Command{ + Use: updateCheckHiddenCmd, + Hidden: true, + Args: cobra.NoArgs, + // No RunE error is ever returned — keep it utterly silent. + Run: func(cmd *cobra.Command, args []string) { + refreshUpdateCache() + }, + } +} + +// refreshUpdateCache does the actual GitHub fetch + cache write. Exported via a +// var seam (updateCheckURL) so the hidden command resolves the same endpoint +// the rest of the code uses. Best-effort: errors are swallowed. +func refreshUpdateCache() { + ctx, cancel := contextWithUpdateTimeout() + defer cancel() + + latest, err := fetchLatestRelease(ctx, latestReleaseURL) + if err != nil || latest == "" || !isParseableVersion(latest) { + // Even on failure, stamp last_check so we don't hammer GitHub every + // single command while offline — back off for the full TTL. + path, perr := updateCachePath() + if perr == nil { + c := readUpdateCache(path) + c.LastCheck = time.Now() + _ = writeUpdateCache(path, c) + } + return + } + + path, err := updateCachePath() + if err != nil { + return + } + c := readUpdateCache(path) + c.LastCheck = time.Now() + c.LatestVersion = latest + _ = writeUpdateCache(path, c) +} diff --git a/internal/cmd/update_notice_test.go b/internal/cmd/update_notice_test.go new file mode 100644 index 0000000..bf075af --- /dev/null +++ b/internal/cmd/update_notice_test.go @@ -0,0 +1,393 @@ +package cmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" +) + +// withParseableVersion pins the package `version` var to a known parseable +// value for the duration of the test (the notice only fires when current +// parses). Restores the original afterward. +func withParseableVersion(t *testing.T, v string) { + t.Helper() + orig := version + version = v + t.Cleanup(func() { version = orig }) +} + +// forceTTY makes stderrIsTerminal report the given value for the test. +func forceTTY(t *testing.T, isTTY bool) { + t.Helper() + orig := stderrIsTerminal + stderrIsTerminal = func() bool { return isTTY } + t.Cleanup(func() { stderrIsTerminal = orig }) +} + +// captureRefresh replaces requestRefresh with a recorder and returns a pointer +// to the call count. No process is forked. +func captureRefresh(t *testing.T) *int { + t.Helper() + orig := requestRefresh + var n int + requestRefresh = func() error { n++; return nil } + t.Cleanup(func() { requestRefresh = orig }) + return &n +} + +// useTempConfigDir points config.Dir() at a temp dir via XDG_CONFIG_HOME and +// clears CI/opt-out env so the gating doesn't accidentally suppress. +func useTempConfigDir(t *testing.T) string { + t.Helper() + dir := t.TempDir() + t.Setenv("XDG_CONFIG_HOME", dir) + t.Setenv("CIVITAI_NO_UPDATE_CHECK", "") + for _, k := range ciEnvVars { + t.Setenv(k, "") + } + return dir +} + +func TestCacheIsFresh(t *testing.T) { + now := time.Now() + if cacheIsFresh(updateCache{}, now) { + t.Error("empty cache (zero LastCheck) must not be fresh") + } + if !cacheIsFresh(updateCache{LastCheck: now.Add(-1 * time.Hour)}, now) { + t.Error("1h-old check should be fresh") + } + if cacheIsFresh(updateCache{LastCheck: now.Add(-25 * time.Hour)}, now) { + t.Error("25h-old check should be stale") + } +} + +func TestDecideNotice_FreshCacheNoRefresh(t *testing.T) { + now := time.Now() + c := updateCache{ + LastCheck: now.Add(-1 * time.Hour), // fresh + LatestVersion: "v0.1.12", + } + d := decideNotice(c, "v0.1.11", now) + if d.refresh { + t.Error("fresh cache must NOT request a refresh") + } + if !d.show { + t.Error("behind + fresh + never-notified should show the notice") + } + if !strings.Contains(d.notice, "v0.1.12") || !strings.Contains(d.notice, "v0.1.11") { + t.Errorf("notice should mention both versions: %q", d.notice) + } + if !strings.Contains(d.notice, "civitai upgrade") { + t.Errorf("notice should mention 'civitai upgrade': %q", d.notice) + } +} + +func TestDecideNotice_StaleCacheRequestsRefresh(t *testing.T) { + now := time.Now() + c := updateCache{ + LastCheck: now.Add(-48 * time.Hour), // stale + LatestVersion: "v0.1.12", + } + d := decideNotice(c, "v0.1.11", now) + if !d.refresh { + t.Error("stale cache MUST request a refresh") + } + // Stale but with a usable cached value => still shows the (stale) notice. + if !d.show { + t.Error("stale cache with a usable cached value should still show the notice") + } +} + +func TestDecideNotice_EmptyCacheFirstRun(t *testing.T) { + now := time.Now() + d := decideNotice(updateCache{}, "v0.1.11", now) + if !d.refresh { + t.Error("empty cache should kick off the first refresh") + } + if d.show { + t.Error("first run with no cached latest must NOT show a notice") + } +} + +func TestDecideNotice_UpToDateNoNotice(t *testing.T) { + now := time.Now() + c := updateCache{LastCheck: now, LatestVersion: "v0.1.11"} + if d := decideNotice(c, "v0.1.11", now); d.show { + t.Error("equal versions must not show a notice") + } + // Current ahead of latest (dev build): no notice. + c2 := updateCache{LastCheck: now, LatestVersion: "v0.1.11"} + if d := decideNotice(c2, "v0.2.0", now); d.show { + t.Error("current ahead of latest must not show a notice") + } +} + +func TestDecideNotice_UnparseableCurrentNoNotice(t *testing.T) { + now := time.Now() + c := updateCache{LastCheck: now, LatestVersion: "v0.1.12"} + if d := decideNotice(c, "dev", now); d.show { + t.Error("unparseable current (dev) must not be nagged") + } +} + +func TestDecideNotice_UnparseableCachedLatestNoNotice(t *testing.T) { + now := time.Now() + c := updateCache{LastCheck: now, LatestVersion: "garbage"} + if d := decideNotice(c, "v0.1.11", now); d.show { + t.Error("unparseable cached latest must not show a notice") + } +} + +func TestDecideNotice_NotifyThrottle(t *testing.T) { + now := time.Now() + // Behind, but notified 1h ago => throttled (<=1/24h). + c := updateCache{ + LastCheck: now, + LatestVersion: "v0.1.12", + LastNotified: now.Add(-1 * time.Hour), + } + if d := decideNotice(c, "v0.1.11", now); d.show { + t.Error("notice within 24h of last_notified must be throttled") + } + // Notified 25h ago => allowed again. + c.LastNotified = now.Add(-25 * time.Hour) + if d := decideNotice(c, "v0.1.11", now); !d.show { + t.Error("notice should be allowed once >24h since last_notified") + } +} + +func TestReadUpdateCache_CorruptIsEmptySilent(t *testing.T) { + dir := t.TempDir() + path := dir + "/update-check.json" + if err := writeFileHelper(path, "{ not valid json"); err != nil { + t.Fatal(err) + } + c := readUpdateCache(path) // must not panic/error + if c.LatestVersion != "" || !c.LastCheck.IsZero() { + t.Errorf("corrupt cache should read as empty, got %+v", c) + } +} + +func TestReadUpdateCache_MissingIsEmpty(t *testing.T) { + c := readUpdateCache(t.TempDir() + "/does-not-exist.json") + if c.LatestVersion != "" { + t.Errorf("missing cache should read as empty, got %+v", c) + } +} + +func TestWriteReadUpdateCacheRoundTrip(t *testing.T) { + dir := t.TempDir() + path := dir + "/update-check.json" + now := time.Now().Truncate(time.Second) + in := updateCache{LastCheck: now, LatestVersion: "v0.1.12", LastNotified: now} + if err := writeUpdateCache(path, in); err != nil { + t.Fatalf("write: %v", err) + } + out := readUpdateCache(path) + if out.LatestVersion != in.LatestVersion || !out.LastCheck.Equal(in.LastCheck) { + t.Errorf("round-trip mismatch: in=%+v out=%+v", in, out) + } +} + +// TestMaybeNotify_FreshCacheNoSpawnUsesCached: a fresh cache => no spawn, but +// the cached notice still prints to the provided (stderr) writer. +func TestMaybeNotify_FreshCacheNoSpawnUsesCached(t *testing.T) { + useTempConfigDir(t) + withParseableVersion(t, "v0.1.11") + forceTTY(t, true) + n := captureRefresh(t) + + path, _ := updateCachePath() + now := time.Now() + if err := writeUpdateCache(path, updateCache{LastCheck: now, LatestVersion: "v0.1.12"}); err != nil { + t.Fatal(err) + } + + var errb bytes.Buffer + maybeNotifyUpdate(&errb, "whoami", false) + + if *n != 0 { + t.Errorf("fresh cache should NOT spawn a refresh, got %d", *n) + } + if !strings.Contains(errb.String(), "v0.1.12") { + t.Errorf("fresh cache should print the cached notice: %q", errb.String()) + } +} + +// TestMaybeNotify_StaleCacheRequestsRefresh: a stale cache => a refresh is +// requested (asserted via the injected recorder; no fork). +func TestMaybeNotify_StaleCacheRequestsRefresh(t *testing.T) { + useTempConfigDir(t) + withParseableVersion(t, "v0.1.11") + forceTTY(t, true) + n := captureRefresh(t) + + path, _ := updateCachePath() + if err := writeUpdateCache(path, updateCache{ + LastCheck: time.Now().Add(-48 * time.Hour), + LatestVersion: "v0.1.12", + }); err != nil { + t.Fatal(err) + } + + var errb bytes.Buffer + maybeNotifyUpdate(&errb, "whoami", false) + + if *n != 1 { + t.Errorf("stale cache should request exactly 1 refresh, got %d", *n) + } +} + +// TestMaybeNotify_PersistsLastNotified proves the once-a-day throttle survives +// across invocations: the first call prints + records last_notified; the second +// (same day) stays silent. +func TestMaybeNotify_PersistsLastNotified(t *testing.T) { + useTempConfigDir(t) + withParseableVersion(t, "v0.1.11") + forceTTY(t, true) + captureRefresh(t) + + path, _ := updateCachePath() + if err := writeUpdateCache(path, updateCache{LastCheck: time.Now(), LatestVersion: "v0.1.12"}); err != nil { + t.Fatal(err) + } + + var first, second bytes.Buffer + maybeNotifyUpdate(&first, "whoami", false) + maybeNotifyUpdate(&second, "whoami", false) + + if !strings.Contains(first.String(), "v0.1.12") { + t.Errorf("first call should notify: %q", first.String()) + } + if second.Len() != 0 { + t.Errorf("second same-day call should be throttled silent, got: %q", second.String()) + } +} + +// Gating matrix: each suppression condition => no notice AND no spawn. +func TestMaybeNotify_GatingSuppresses(t *testing.T) { + type setup func(t *testing.T) + cases := []struct { + name string + cmdName string + noFlag bool + prep setup + }{ + {name: "CI env set", cmdName: "whoami", prep: func(t *testing.T) { t.Setenv("CI", "true") }}, + {name: "non-TTY", cmdName: "whoami", prep: func(t *testing.T) { forceTTY(t, false) }}, + {name: "opt-out flag", cmdName: "whoami", noFlag: true, prep: func(t *testing.T) {}}, + {name: "opt-out env", cmdName: "whoami", prep: func(t *testing.T) { t.Setenv("CIVITAI_NO_UPDATE_CHECK", "1") }}, + {name: "excluded version", cmdName: "version", prep: func(t *testing.T) {}}, + {name: "excluded upgrade", cmdName: "upgrade", prep: func(t *testing.T) {}}, + {name: "excluded completion", cmdName: "completion", prep: func(t *testing.T) {}}, + {name: "excluded help", cmdName: "help", prep: func(t *testing.T) {}}, + {name: "internal __complete", cmdName: "__complete", prep: func(t *testing.T) {}}, + {name: "internal __update-check", cmdName: "__update-check", prep: func(t *testing.T) {}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + useTempConfigDir(t) + withParseableVersion(t, "v0.1.11") + forceTTY(t, true) // default TTY; some cases override to false + n := captureRefresh(t) + tc.prep(t) + + // Behind, fresh cache => would notify if not gated. + path, _ := updateCachePath() + if err := writeUpdateCache(path, updateCache{LastCheck: time.Now(), LatestVersion: "v0.1.12"}); err != nil { + t.Fatal(err) + } + + var errb bytes.Buffer + maybeNotifyUpdate(&errb, tc.cmdName, tc.noFlag) + + if errb.Len() != 0 { + t.Errorf("%s: expected no notice, got %q", tc.name, errb.String()) + } + if *n != 0 { + t.Errorf("%s: expected no refresh spawn, got %d", tc.name, *n) + } + }) + } +} + +// TestNoticeGoesToStderrNotStdout drives the full root command with a stale, +// behind cache + forced TTY and asserts the notice lands on STDERR, never +// STDOUT (so pipes/scripts are unaffected). It uses a SUCCESSFUL whoami (cobra +// only runs PersistentPostRun when the command's RunE succeeds — the notice is +// deliberately never appended to a failed command). +func TestNoticeGoesToStderrNotStdout(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{"username":"bob","id":99}`)) + })) + t.Cleanup(srv.Close) + + useTempConfigDir(t) + withParseableVersion(t, "v0.1.11") + forceTTY(t, true) + captureRefresh(t) + t.Setenv("CIVITAI_TOKEN", "tok-1") + t.Setenv("CIVITAI_BASE_URL", srv.URL) + + path, _ := updateCachePath() + if err := writeUpdateCache(path, updateCache{LastCheck: time.Now(), LatestVersion: "v0.1.12"}); err != nil { + t.Fatal(err) + } + + out, errOut, err := run(t, "whoami") + if err != nil { + t.Fatalf("whoami should succeed: %v", err) + } + + if strings.Contains(out, "v0.1.12") { + t.Errorf("update notice must NOT pollute stdout: %q", out) + } + if !strings.Contains(errOut, "v0.1.12") { + t.Errorf("update notice should be on stderr: %q", errOut) + } +} + +func TestRefreshUpdateCache_WritesLatest(t *testing.T) { + useTempConfigDir(t) + srv, hits := newReleaseServer(t, "v0.1.12", 200) + pointAtServer(t, srv.URL) + + refreshUpdateCache() + + if *hits != 1 { + t.Errorf("expected 1 GitHub hit, got %d", *hits) + } + path, _ := updateCachePath() + c := readUpdateCache(path) + if c.LatestVersion != "v0.1.12" { + t.Errorf("refresh should cache the latest version, got %q", c.LatestVersion) + } + if c.LastCheck.IsZero() { + t.Error("refresh should stamp last_check") + } +} + +func TestRefreshUpdateCache_FailureStillStampsLastCheck(t *testing.T) { + useTempConfigDir(t) + srv, _ := newReleaseServer(t, "", 500) + pointAtServer(t, srv.URL) + + refreshUpdateCache() + + path, _ := updateCachePath() + c := readUpdateCache(path) + if c.LastCheck.IsZero() { + t.Error("a failed refresh should still stamp last_check to back off") + } + if c.LatestVersion != "" { + t.Errorf("a failed refresh should not set a latest version, got %q", c.LatestVersion) + } +} + +func writeFileHelper(path, content string) error { + return os.WriteFile(path, []byte(content), 0o600) +} diff --git a/internal/cmd/upgrade.go b/internal/cmd/upgrade.go new file mode 100644 index 0000000..2c3192f --- /dev/null +++ b/internal/cmd/upgrade.go @@ -0,0 +1,318 @@ +package cmd + +import ( + "archive/tar" + "compress/gzip" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "runtime" + "strings" + "time" + + "github.com/minio/selfupdate" + "github.com/spf13/cobra" +) + +// upgradeTimeout bounds the whole self-update (release lookup + downloads). +// Downloads are a few MB so this is generous. +const upgradeTimeout = 60 * time.Second + +// maxDownloadBytes caps each download so a misbehaving/hostile endpoint can't +// make us read forever. The largest current asset is ~5MB; 64MB is ample. +const maxDownloadBytes = 64 << 20 + +// releaseAsset is a single downloadable asset. +type releaseAsset struct { + Name string `json:"name"` + DownloadURL string `json:"browser_download_url"` +} + +// fullRelease is the release JSON shape with assets we actually unmarshal. +type fullRelease struct { + TagName string `json:"tag_name"` + Assets []releaseAsset `json:"assets"` +} + +// Seams for tests. +var ( + // osExecutable resolves the running binary path; overridable in tests. + osExecutable = os.Executable + // evalSymlinks resolves symlinks; overridable in tests. + evalSymlinks = filepath.EvalSymlinks + // applyUpdate performs the atomic binary replacement. Default uses + // minio/selfupdate; tests override it to write to a temp path. + applyUpdate = func(newBinary io.Reader, targetPath string) error { + return selfupdate.Apply(newBinary, selfupdate.Options{TargetPath: targetPath}) + } +) + +// brewPathMarkers are path fragments that indicate a Homebrew-managed install. +// If the resolved executable lives under any of these, we delegate to brew +// instead of self-replacing (unless --force). +var brewPathMarkers = []string{ + "/Cellar/", + "/Caskroom/", + "/opt/homebrew", + "/usr/local/Homebrew", + "/usr/local/Cellar", + "/home/linuxbrew/.linuxbrew", +} + +func newUpgradeCmd() *cobra.Command { + var force bool + cmd := &cobra.Command{ + Use: "upgrade", + Short: "Update the civitai CLI to the latest release", + Long: `Download and install the latest civitai release, replacing this binary. + +The latest release is resolved from the public GitHub releases API (no token is +ever sent). The downloaded tarball is verified against its SHA-256 checksum +before anything is replaced — a mismatch aborts the upgrade and leaves the +current binary untouched. + +If this binary was installed via Homebrew, upgrade delegates to: + brew upgrade civitai/tap/civitai +(use --force to self-replace anyway).`, + Example: ` civitai upgrade + civitai upgrade --force`, + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, args []string) error { + noUpdateCheck, _ := cmd.Flags().GetBool("no-update-check") + return runUpgrade(cmd.OutOrStdout(), force, noUpdateCheck) + }, + } + cmd.Flags().BoolVar(&force, "force", false, + "reinstall even if already up to date, and self-replace a Homebrew install") + return cmd +} + +// runUpgrade orchestrates the self-update. out is stdout for user messages. +func runUpgrade(out io.Writer, force, noUpdateCheck bool) error { + ctx, cancel := context.WithTimeout(context.Background(), upgradeTimeout) + defer cancel() + + rel, err := fetchRelease(ctx, latestReleaseURL) + if err != nil { + return fmt.Errorf("look up latest release: %w", err) + } + latest := rel.TagName + if !isParseableVersion(latest) { + return fmt.Errorf("latest release tag %q is not a recognizable version", latest) + } + + // Already up to date? (current >= latest). --force overrides. + if !force && compareVersions(version, latest) >= 0 && isParseableVersion(version) { + fmt.Fprintf(out, "civitai is already up to date (%s).\n", version) + return nil + } + + // Resolve the running executable, following symlinks (Homebrew installs a + // symlink in the bin dir pointing into the Cellar). + exe, err := osExecutable() + if err != nil { + return fmt.Errorf("locate current executable: %w", err) + } + resolved := exe + if r, rerr := evalSymlinks(exe); rerr == nil { + resolved = r + } + + // Homebrew delegation (unless --force). + if !force && isHomebrewPath(resolved) { + fmt.Fprintln(out, "civitai was installed via Homebrew. Upgrade with:") + fmt.Fprintln(out, " brew upgrade civitai/tap/civitai") + return nil + } + + // Build the asset names for this platform. + verNoV := strings.TrimPrefix(latest, "v") + tarName := fmt.Sprintf("civitai_%s_%s_%s.tar.gz", verNoV, runtime.GOOS, runtime.GOARCH) + + tarURL := findAssetURL(rel, tarName) + if tarURL == "" { + return fmt.Errorf("no release asset %q for %s/%s — upgrade manually from %s", + tarName, runtime.GOOS, runtime.GOARCH, "https://github.com/civitai/cli/releases/latest") + } + sumsURL := findAssetURL(rel, "checksums.txt") + if sumsURL == "" { + return errors.New("release is missing checksums.txt — refusing to upgrade without integrity verification") + } + + // Download the checksums and the tarball. + sums, err := download(ctx, sumsURL) + if err != nil { + return fmt.Errorf("download checksums: %w", err) + } + wantSum, ok := checksumFor(string(sums), tarName) + if !ok { + return fmt.Errorf("checksums.txt has no entry for %s — aborting", tarName) + } + + tarball, err := download(ctx, tarURL) + if err != nil { + return fmt.Errorf("download %s: %w", tarName, err) + } + + // MANDATORY: verify SHA-256 BEFORE we touch the binary on disk. + gotSum := sha256Hex(tarball) + if gotSum != wantSum { + return fmt.Errorf("checksum mismatch for %s: got %s, want %s — aborting upgrade (binary NOT replaced)", + tarName, gotSum, wantSum) + } + + // Extract the civitai binary from the verified tarball. + bin, err := extractBinaryFromTarGz(tarball, "civitai") + if err != nil { + return fmt.Errorf("extract binary: %w", err) + } + + // Atomically replace the running executable. + if err := applyUpdate(strings.NewReader(string(bin)), resolved); err != nil { + if isPermissionError(err) { + return fmt.Errorf("cannot replace %s: permission denied.\n"+ + "Try one of:\n"+ + " sudo civitai upgrade\n"+ + " reinstall via Homebrew (brew upgrade civitai/tap/civitai)\n"+ + " or: go install github.com/civitai/cli/cmd/civitai@latest\noriginal error: %w", + resolved, err) + } + return fmt.Errorf("install update: %w", err) + } + + fmt.Fprintf(out, "Upgraded civitai %s → %s.\n", version, latest) + return nil +} + +// isHomebrewPath reports whether path lives under a known Homebrew location. +func isHomebrewPath(path string) bool { + for _, m := range brewPathMarkers { + if strings.Contains(path, m) { + return true + } + } + return false +} + +// fetchRelease does the unauthenticated GitHub call and returns the full release +// (tag + assets). No Authorization header is ever attached. +func fetchRelease(ctx context.Context, url string) (fullRelease, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return fullRelease{}, err + } + req.Header.Set("Accept", "application/vnd.github+json") + // NOTE: intentionally NO Authorization header — public, token-free. + resp, err := http.DefaultClient.Do(req) + if err != nil { + return fullRelease{}, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return fullRelease{}, fmt.Errorf("github releases: status %d", resp.StatusCode) + } + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return fullRelease{}, err + } + var rel fullRelease + if err := json.Unmarshal(body, &rel); err != nil { + return fullRelease{}, err + } + return rel, nil +} + +// findAssetURL returns the download URL for the named asset, or "". +func findAssetURL(rel fullRelease, name string) string { + for _, a := range rel.Assets { + if a.Name == name { + return a.DownloadURL + } + } + return "" +} + +// download GETs url and returns the body (capped). It explicitly attaches NO +// Authorization header — release downloads are unauthenticated HTTPS. +func download(ctx context.Context, url string) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, err + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("status %d", resp.StatusCode) + } + return io.ReadAll(io.LimitReader(resp.Body, maxDownloadBytes)) +} + +// checksumFor parses a `sha256 filename` checksums.txt body and returns the +// hex digest for the given filename. +func checksumFor(sums, filename string) (string, bool) { + for _, line := range strings.Split(sums, "\n") { + line = strings.TrimSpace(line) + if line == "" { + continue + } + fields := strings.Fields(line) + if len(fields) != 2 { + continue + } + // goreleaser writes " " (filename may carry a leading *). + name := strings.TrimPrefix(fields[1], "*") + if name == filename { + return strings.ToLower(fields[0]), true + } + } + return "", false +} + +// sha256Hex returns the lowercase hex SHA-256 of b. +func sha256Hex(b []byte) string { + sum := sha256.Sum256(b) + return hex.EncodeToString(sum[:]) +} + +// extractBinaryFromTarGz pulls the named file out of a gzip'd tar archive. The +// goreleaser archive contains the binary at the top level (plus README/LICENSE). +func extractBinaryFromTarGz(data []byte, name string) ([]byte, error) { + gz, err := gzip.NewReader(strings.NewReader(string(data))) + if err != nil { + return nil, err + } + defer gz.Close() + tr := tar.NewReader(gz) + for { + hdr, err := tr.Next() + if err == io.EOF { + break + } + if err != nil { + return nil, err + } + if hdr.Typeflag != tar.TypeReg { + continue + } + // Match by base name so a possible directory prefix doesn't matter. + if filepath.Base(hdr.Name) == name { + return io.ReadAll(io.LimitReader(tr, maxDownloadBytes)) + } + } + return nil, fmt.Errorf("binary %q not found in archive", name) +} + +// isPermissionError reports whether err is (or wraps) a permission-denied error. +func isPermissionError(err error) bool { + return errors.Is(err, os.ErrPermission) +} diff --git a/internal/cmd/upgrade_test.go b/internal/cmd/upgrade_test.go new file mode 100644 index 0000000..98386a9 --- /dev/null +++ b/internal/cmd/upgrade_test.go @@ -0,0 +1,397 @@ +package cmd + +import ( + "archive/tar" + "bytes" + "compress/gzip" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "runtime" + "strings" + "testing" +) + +// makeTarGz builds an in-memory gzip'd tar containing a single regular file +// `name` with the given content. +func makeTarGz(t *testing.T, name string, content []byte) []byte { + t.Helper() + var buf bytes.Buffer + gz := gzip.NewWriter(&buf) + tw := tar.NewWriter(gz) + hdr := &tar.Header{Name: name, Mode: 0o755, Size: int64(len(content)), Typeflag: tar.TypeReg} + if err := tw.WriteHeader(hdr); err != nil { + t.Fatal(err) + } + if _, err := tw.Write(content); err != nil { + t.Fatal(err) + } + if err := tw.Close(); err != nil { + t.Fatal(err) + } + if err := gz.Close(); err != nil { + t.Fatal(err) + } + return buf.Bytes() +} + +func sha256Of(b []byte) string { + s := sha256.Sum256(b) + return hex.EncodeToString(s[:]) +} + +// upgradeServer stands up an httptest server emulating the GitHub release API + +// the asset downloads (tarball + checksums.txt). tarChecksumOverride, if +// non-empty, is written into checksums.txt instead of the real digest (to test +// the mismatch-abort path). It records whether any download carried an +// Authorization header. +type upgradeServer struct { + srv *httptest.Server + tarName string + sawAuth bool + tarBytes []byte + tarHits int + checksumHits int +} + +func newUpgradeServer(t *testing.T, tag string, binContent []byte, tarChecksumOverride string) *upgradeServer { + t.Helper() + us := &upgradeServer{} + verNoV := strings.TrimPrefix(tag, "v") + us.tarName = fmt.Sprintf("civitai_%s_%s_%s.tar.gz", verNoV, runtime.GOOS, runtime.GOARCH) + us.tarBytes = makeTarGz(t, "civitai", binContent) + + sum := sha256Of(us.tarBytes) + if tarChecksumOverride != "" { + sum = tarChecksumOverride + } + checksums := fmt.Sprintf("%s %s\n", sum, us.tarName) + + mux := http.NewServeMux() + // The release JSON is served at "/" (latestReleaseURL is pointed here). + mux.HandleFunc("/release", func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "" { + t.Error("release lookup sent an Authorization header") + } + base := us.srv.URL + fmt.Fprintf(w, `{"tag_name":%q,"assets":[ + {"name":%q,"browser_download_url":%q}, + {"name":"checksums.txt","browser_download_url":%q} + ]}`, tag, us.tarName, base+"/dl/"+us.tarName, base+"/dl/checksums.txt") + }) + mux.HandleFunc("/dl/checksums.txt", func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "" { + us.sawAuth = true + } + us.checksumHits++ + _, _ = io.WriteString(w, checksums) + }) + mux.HandleFunc("/dl/"+us.tarName, func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "" { + us.sawAuth = true + } + us.tarHits++ + _, _ = w.Write(us.tarBytes) + }) + us.srv = httptest.NewServer(mux) + t.Cleanup(us.srv.Close) + return us +} + +func (us *upgradeServer) releaseURL() string { return us.srv.URL + "/release" } + +// withExecutable pins osExecutable + evalSymlinks to a fixed path for a test. +func withExecutable(t *testing.T, path string) { + t.Helper() + oe, es := osExecutable, evalSymlinks + osExecutable = func() (string, error) { return path, nil } + evalSymlinks = func(p string) (string, error) { return p, nil } + t.Cleanup(func() { osExecutable, evalSymlinks = oe, es }) +} + +// captureApply replaces applyUpdate with a recorder writing the new bytes to the +// target path (emulating the atomic replace without the selfupdate machinery, +// so the test runs as non-root on any FS). Returns the recorded target. +func captureApply(t *testing.T) *string { + t.Helper() + orig := applyUpdate + var target string + applyUpdate = func(r io.Reader, tp string) error { + target = tp + b, err := io.ReadAll(r) + if err != nil { + return err + } + return os.WriteFile(tp, b, 0o755) + } + t.Cleanup(func() { applyUpdate = orig }) + return &target +} + +func TestIsHomebrewPath(t *testing.T) { + brew := []string{ + "/opt/homebrew/Cellar/civitai/0.1.11/bin/civitai", + "/usr/local/Cellar/civitai/0.1.11/bin/civitai", + "/home/linuxbrew/.linuxbrew/Cellar/civitai/0.1.11/bin/civitai", + "/opt/homebrew/bin/civitai", + "/Users/x/Library/Caskroom/civitai/0.1.11/civitai", + } + for _, p := range brew { + if !isHomebrewPath(p) { + t.Errorf("expected %q to be detected as a Homebrew path", p) + } + } + notBrew := []string{ + "/usr/local/bin/civitai", + "/home/zach/go/bin/civitai", + "/tmp/civitai", + } + for _, p := range notBrew { + if isHomebrewPath(p) { + t.Errorf("expected %q NOT to be a Homebrew path", p) + } + } +} + +func TestChecksumFor(t *testing.T) { + body := "abc123 civitai_0.1.11_linux_amd64.tar.gz\ndef456 other.zip\n" + got, ok := checksumFor(body, "civitai_0.1.11_linux_amd64.tar.gz") + if !ok || got != "abc123" { + t.Errorf("checksumFor = %q,%v want abc123,true", got, ok) + } + if _, ok := checksumFor(body, "missing.tar.gz"); ok { + t.Error("checksumFor should return false for a missing filename") + } + // BSD-style leading * marker. + star := "deadbeef *civitai_x.tar.gz\n" + if g, ok := checksumFor(star, "civitai_x.tar.gz"); !ok || g != "deadbeef" { + t.Errorf("checksumFor should strip a leading *: got %q,%v", g, ok) + } +} + +func TestExtractBinaryFromTarGz(t *testing.T) { + want := []byte("#!/fake/civitai\x00binary") + data := makeTarGz(t, "civitai", want) + got, err := extractBinaryFromTarGz(data, "civitai") + if err != nil { + t.Fatalf("extract: %v", err) + } + if !bytes.Equal(got, want) { + t.Errorf("extracted bytes mismatch") + } + if _, err := extractBinaryFromTarGz(data, "nope"); err == nil { + t.Error("expected error extracting a missing file") + } +} + +func TestUpgrade_AlreadyLatestNoOp(t *testing.T) { + withParseableVersion(t, "v0.1.11") + us := newUpgradeServer(t, "v0.1.11", []byte("new"), "") + pointAtServer(t, us.releaseURL()) + target := captureApply(t) + + var out bytes.Buffer + if err := runUpgrade(&out, false, false); err != nil { + t.Fatalf("upgrade: %v", err) + } + if !strings.Contains(out.String(), "already up to date") { + t.Errorf("expected up-to-date message: %q", out.String()) + } + if *target != "" { + t.Error("up-to-date should NOT call applyUpdate") + } + if us.tarHits != 0 { + t.Errorf("up-to-date should not download the tarball, got %d hits", us.tarHits) + } +} + +func TestUpgrade_ForceProceedsWhenLatest(t *testing.T) { + withParseableVersion(t, "v0.1.11") + binContent := []byte("FORCED-NEW-BINARY") + us := newUpgradeServer(t, "v0.1.11", binContent, "") + pointAtServer(t, us.releaseURL()) + target := captureApply(t) + + exe := filepath.Join(t.TempDir(), "civitai") + if err := os.WriteFile(exe, []byte("OLD"), 0o755); err != nil { + t.Fatal(err) + } + withExecutable(t, exe) + + var out bytes.Buffer + if err := runUpgrade(&out, true /*force*/, false); err != nil { + t.Fatalf("upgrade --force: %v", err) + } + if *target != exe { + t.Errorf("force should replace the resolved exe %q, got %q", exe, *target) + } + got, _ := os.ReadFile(exe) + if !bytes.Equal(got, binContent) { + t.Errorf("force should have swapped the binary content, got %q", got) + } +} + +func TestUpgrade_BrewDetectionDelegates(t *testing.T) { + withParseableVersion(t, "v0.1.10") // behind, so it would normally proceed + us := newUpgradeServer(t, "v0.1.11", []byte("new"), "") + pointAtServer(t, us.releaseURL()) + target := captureApply(t) + withExecutable(t, "/opt/homebrew/Cellar/civitai/0.1.10/bin/civitai") + + var out bytes.Buffer + if err := runUpgrade(&out, false, false); err != nil { + t.Fatalf("upgrade: %v", err) + } + if !strings.Contains(out.String(), "brew upgrade civitai/tap/civitai") { + t.Errorf("brew install should print the brew command: %q", out.String()) + } + if *target != "" { + t.Error("brew path should NOT self-replace") + } + if us.tarHits != 0 || us.checksumHits != 0 { + t.Errorf("brew path should not download anything, tar=%d sums=%d", us.tarHits, us.checksumHits) + } +} + +func TestUpgrade_BrewForceOverrides(t *testing.T) { + withParseableVersion(t, "v0.1.10") + binContent := []byte("BREW-FORCE-REPLACED") + us := newUpgradeServer(t, "v0.1.11", binContent, "") + pointAtServer(t, us.releaseURL()) + target := captureApply(t) + + // A brew path, but --force should self-replace anyway. + exe := filepath.Join(t.TempDir(), "Cellar", "civitai") + if err := os.MkdirAll(filepath.Dir(exe), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(exe, []byte("OLD"), 0o755); err != nil { + t.Fatal(err) + } + withExecutable(t, exe) + + var out bytes.Buffer + if err := runUpgrade(&out, true, false); err != nil { + t.Fatalf("upgrade --force on brew: %v", err) + } + if *target != exe { + t.Errorf("force should override brew delegation and replace %q, got %q", exe, *target) + } +} + +func TestUpgrade_ChecksumMismatchAborts(t *testing.T) { + withParseableVersion(t, "v0.1.10") + // Override the checksum to a bogus digest => must abort BEFORE replacing. + us := newUpgradeServer(t, "v0.1.11", []byte("new"), strings.Repeat("0", 64)) + pointAtServer(t, us.releaseURL()) + target := captureApply(t) + + exe := filepath.Join(t.TempDir(), "civitai") + original := []byte("ORIGINAL-UNTOUCHED") + if err := os.WriteFile(exe, original, 0o755); err != nil { + t.Fatal(err) + } + withExecutable(t, exe) + + var out bytes.Buffer + err := runUpgrade(&out, false, false) + if err == nil { + t.Fatal("expected checksum mismatch to abort with an error") + } + if !strings.Contains(err.Error(), "checksum mismatch") { + t.Errorf("error should mention checksum mismatch: %v", err) + } + if *target != "" { + t.Error("checksum mismatch must NOT call applyUpdate") + } + got, _ := os.ReadFile(exe) + if !bytes.Equal(got, original) { + t.Errorf("original binary must be untouched on mismatch, got %q", got) + } +} + +func TestUpgrade_HappyPathReplacesBinary(t *testing.T) { + withParseableVersion(t, "v0.1.10") + newBin := []byte("BRAND-NEW-CIVITAI-BINARY-BYTES") + us := newUpgradeServer(t, "v0.1.11", newBin, "") + pointAtServer(t, us.releaseURL()) + target := captureApply(t) + + exe := filepath.Join(t.TempDir(), "civitai") + if err := os.WriteFile(exe, []byte("OLD-BINARY"), 0o755); err != nil { + t.Fatal(err) + } + withExecutable(t, exe) + + var out bytes.Buffer + if err := runUpgrade(&out, false, false); err != nil { + t.Fatalf("upgrade happy path: %v", err) + } + if !strings.Contains(out.String(), "Upgraded civitai v0.1.10 → v0.1.11") { + t.Errorf("expected success message: %q", out.String()) + } + if *target != exe { + t.Errorf("applyUpdate target should be the resolved exe %q, got %q", exe, *target) + } + got, _ := os.ReadFile(exe) + if !bytes.Equal(got, newBin) { + t.Errorf("binary should be swapped to the new bytes, got %q", got) + } + // Verified-then-applied: both assets downloaded exactly once. + if us.tarHits != 1 || us.checksumHits != 1 { + t.Errorf("happy path should download tar+checksums once each, got tar=%d sums=%d", us.tarHits, us.checksumHits) + } + // SECURITY: no Authorization header on any download. + if us.sawAuth { + t.Error("upgrade must NOT send an Authorization header to the download host") + } +} + +func TestUpgrade_PermissionDeniedClearError(t *testing.T) { + withParseableVersion(t, "v0.1.10") + us := newUpgradeServer(t, "v0.1.11", []byte("new"), "") + pointAtServer(t, us.releaseURL()) + + // Force applyUpdate to return a permission error. + orig := applyUpdate + applyUpdate = func(r io.Reader, tp string) error { return os.ErrPermission } + t.Cleanup(func() { applyUpdate = orig }) + + exe := filepath.Join(t.TempDir(), "civitai") + if err := os.WriteFile(exe, []byte("OLD"), 0o755); err != nil { + t.Fatal(err) + } + withExecutable(t, exe) + + var out bytes.Buffer + err := runUpgrade(&out, false, false) + if err == nil { + t.Fatal("expected a permission error") + } + if !strings.Contains(err.Error(), "permission denied") || !strings.Contains(err.Error(), "sudo") { + t.Errorf("permission error should suggest sudo: %v", err) + } +} + +func TestUpgrade_MissingAssetForPlatform(t *testing.T) { + withParseableVersion(t, "v0.1.10") + // Serve a release whose only asset is checksums.txt — no platform tarball. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprintf(w, `{"tag_name":"v0.1.11","assets":[ + {"name":"checksums.txt","browser_download_url":"%s/sums"}]}`, "http://example.invalid") + })) + t.Cleanup(srv.Close) + pointAtServer(t, srv.URL) + withExecutable(t, filepath.Join(t.TempDir(), "civitai")) + + var out bytes.Buffer + err := runUpgrade(&out, false, false) + if err == nil || !strings.Contains(err.Error(), "no release asset") { + t.Errorf("expected a missing-asset error, got: %v", err) + } +} diff --git a/internal/cmd/version.go b/internal/cmd/version.go index a0c7be7..bf12732 100644 --- a/internal/cmd/version.go +++ b/internal/cmd/version.go @@ -8,7 +8,6 @@ import ( ) func newVersionCmd() *cobra.Command { - var noUpdateCheck bool cmd := &cobra.Command{ Use: "version", Short: "Print the CLI version, commit, and build date", @@ -32,13 +31,13 @@ token. Skip it with --no-update-check or by setting CIVITAI_NO_UPDATE_CHECK.`, fmt.Fprintf(out, " built: %s\n", date) fmt.Fprintf(out, " go: %s %s/%s\n", runtime.Version(), runtime.GOOS, runtime.GOARCH) + // --no-update-check is now a persistent flag inherited from root. + noUpdateCheck, _ := cmd.Flags().GetBool("no-update-check") if !updateCheckDisabled(noUpdateCheck) { printUpdateNotice(out, version) } return nil }, } - cmd.Flags().BoolVar(&noUpdateCheck, "no-update-check", false, - "skip the GitHub check for a newer release (also via CIVITAI_NO_UPDATE_CHECK)") return cmd } From da93a6942c0ac01f0f3131389a919790504fbea5 Mon Sep 17 00:00:00 2001 From: ZacxDev Date: Thu, 25 Jun 2026 15:25:03 -0500 Subject: [PATCH 2/2] fix(upgrade): enforce https + GitHub host on asset downloads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Audit 🟡#1 (https/host enforcement). Asset download URLs come straight out of the GitHub release JSON and were fetched with http.DefaultClient, which follows redirects including https->http downgrades with no scheme/host check. An http:// checksums.txt + http:// tarball pair would make the SHA-256 gate self-referential (both halves attacker-controlled). - validateAssetURL: require scheme==https AND host in an allowlist {github.com, objects.githubusercontent.com, release-assets.githubusercontent.com}; reject (abort, no download) otherwise. Applied to BOTH the tarball and checksums.txt URLs before any bytes are read. - assetDownloadClient: dedicated *http.Client with CheckRedirect that re-validates every hop, so an https->http downgrade redirect is rejected rather than followed. Context timeouts + body-size caps kept. - Checksum-verify-before-replace gate unchanged (defense-in-depth on transport). Tests: validateAssetURL allow/reject table; an http:// asset URL and an off-host https asset URL each abort before download with the binary untouched; the asset client rejects an https->http redirect. httptest fixtures now serve TLS and inject the loopback host via a test seam so the production allowlist is never weakened to pass tests. Co-Authored-By: Claude Opus 4.8 (1M context) --- internal/cmd/upgrade.go | 75 +++++++++++++++- internal/cmd/upgrade_test.go | 168 ++++++++++++++++++++++++++++++++++- 2 files changed, 240 insertions(+), 3 deletions(-) diff --git a/internal/cmd/upgrade.go b/internal/cmd/upgrade.go index 2c3192f..f1f68b2 100644 --- a/internal/cmd/upgrade.go +++ b/internal/cmd/upgrade.go @@ -11,6 +11,7 @@ import ( "fmt" "io" "net/http" + "net/url" "os" "path/filepath" "runtime" @@ -29,6 +30,66 @@ const upgradeTimeout = 60 * time.Second // make us read forever. The largest current asset is ~5MB; 64MB is ample. const maxDownloadBytes = 64 << 20 +// assetHostAllowlist is the set of hosts GitHub serves release assets (and the +// hosts it redirects browser_download_url through). A release's asset URLs are +// attacker-influenced data (they come straight out of the release JSON), so we +// pin both the scheme (https) AND the host before fetching either the tarball +// or checksums.txt. Without this, a release that advertised an http:// pair +// (checksums.txt + tarball) would make the SHA-256 gate self-referential — +// both halves attacker-controlled — and http.DefaultClient would happily follow +// an https->http downgrade redirect. +// +// - github.com — the canonical browser_download_url host +// - objects.githubusercontent.com — the S3-backed CDN github.com 302s to +// - release-assets.githubusercontent.com — newer release-asset redirect target +// +// It is a var (not a const map) so tests can inject the httptest loopback host +// via withAssetHost without weakening the production default. +var assetHostAllowlist = map[string]bool{ + "github.com": true, + "objects.githubusercontent.com": true, + "release-assets.githubusercontent.com": true, +} + +// validateAssetURL rejects any asset URL that is not https on an allowlisted +// GitHub host. It returns a non-nil error (aborting the upgrade, before any +// bytes are read) for anything else. +func validateAssetURL(rawURL string) error { + u, err := url.Parse(rawURL) + if err != nil { + return fmt.Errorf("invalid asset URL %q: %w", rawURL, err) + } + if u.Scheme != "https" { + return fmt.Errorf("refusing asset URL %q: scheme %q is not https", rawURL, u.Scheme) + } + if !assetHostAllowlist[u.Hostname()] { + return fmt.Errorf("refusing asset URL %q: host %q is not an allowed GitHub release host", rawURL, u.Hostname()) + } + return nil +} + +// assetDownloadClient is the HTTP client used for the tarball + checksums.txt +// downloads. Unlike http.DefaultClient it refuses to follow a redirect to a +// non-https URL or to a host outside assetHostAllowlist — closing the +// https->http downgrade (and host-pivot) foot-gun on every hop, not just the +// initial URL. Timeouts/body caps stay on the request context + LimitReader. +var assetDownloadClient = &http.Client{ + CheckRedirect: func(req *http.Request, via []*http.Request) error { + if len(via) >= 10 { + return errors.New("stopped after 10 redirects") + } + if err := validateAssetURL(req.URL.String()); err != nil { + return fmt.Errorf("insecure redirect: %w", err) + } + return nil + }, +} + +// releaseClient fetches the release JSON. It is http.DefaultClient in +// production; tests override its transport to trust an httptest TLS server (the +// asset URLs are now scheme-pinned to https, so the test fixtures serve TLS). +var releaseClient = http.DefaultClient + // releaseAsset is a single downloadable asset. type releaseAsset struct { Name string `json:"name"` @@ -146,6 +207,16 @@ func runUpgrade(out io.Writer, force, noUpdateCheck bool) error { return errors.New("release is missing checksums.txt — refusing to upgrade without integrity verification") } + // Defense-in-depth on the transport: the asset URLs come from the release + // JSON, so pin scheme+host BEFORE fetching anything. An http:// or + // non-GitHub asset URL aborts the upgrade here — the binary is never touched. + if err := validateAssetURL(tarURL); err != nil { + return err + } + if err := validateAssetURL(sumsURL); err != nil { + return err + } + // Download the checksums and the tarball. sums, err := download(ctx, sumsURL) if err != nil { @@ -210,7 +281,7 @@ func fetchRelease(ctx context.Context, url string) (fullRelease, error) { } req.Header.Set("Accept", "application/vnd.github+json") // NOTE: intentionally NO Authorization header — public, token-free. - resp, err := http.DefaultClient.Do(req) + resp, err := releaseClient.Do(req) if err != nil { return fullRelease{}, err } @@ -246,7 +317,7 @@ func download(ctx context.Context, url string) ([]byte, error) { if err != nil { return nil, err } - resp, err := http.DefaultClient.Do(req) + resp, err := assetDownloadClient.Do(req) if err != nil { return nil, err } diff --git a/internal/cmd/upgrade_test.go b/internal/cmd/upgrade_test.go index 98386a9..1e084e1 100644 --- a/internal/cmd/upgrade_test.go +++ b/internal/cmd/upgrade_test.go @@ -10,6 +10,7 @@ import ( "io" "net/http" "net/http/httptest" + "net/url" "os" "path/filepath" "runtime" @@ -98,11 +99,51 @@ func newUpgradeServer(t *testing.T, tag string, binContent []byte, tarChecksumOv us.tarHits++ _, _ = w.Write(us.tarBytes) }) - us.srv = httptest.NewServer(mux) + // Serve over TLS: production now scheme-pins asset URLs to https, so the + // asset URLs (and the release JSON they're embedded in) must be https. + us.srv = httptest.NewTLSServer(mux) t.Cleanup(us.srv.Close) + // The httptest server listens on 127.0.0.1, which is not a real GitHub + // release host. Inject it into the asset-host allowlist for this test so the + // production allowlist stays untouched while the integration path still runs. + if u, err := url.Parse(us.srv.URL); err == nil { + withAssetHost(t, u.Hostname()) + } + // Point both the release-lookup and asset-download clients at a transport + // that trusts the test server's self-signed cert, for this test only. + withTrustedClients(t, us.srv.Client()) return us } +// withTrustedClients swaps the transports of releaseClient + assetDownloadClient +// to trust the given test client's TLS cert, restoring them after the test. The +// asset client keeps its production CheckRedirect (the security control under +// test); only the transport (cert trust) is borrowed. +func withTrustedClients(t *testing.T, trusted *http.Client) { + t.Helper() + origRelease := releaseClient + origAssetTransport := assetDownloadClient.Transport + releaseClient = trusted + assetDownloadClient.Transport = trusted.Transport + t.Cleanup(func() { + releaseClient = origRelease + assetDownloadClient.Transport = origAssetTransport + }) +} + +// withAssetHost temporarily adds host to assetHostAllowlist for the duration of +// a test (restoring the original state afterward). This is the seam that lets +// the httptest loopback host pass validateAssetURL without weakening the +// production allowlist. +func withAssetHost(t *testing.T, host string) { + t.Helper() + if assetHostAllowlist[host] { + return + } + assetHostAllowlist[host] = true + t.Cleanup(func() { delete(assetHostAllowlist, host) }) +} + func (us *upgradeServer) releaseURL() string { return us.srv.URL + "/release" } // withExecutable pins osExecutable + evalSymlinks to a fixed path for a test. @@ -378,6 +419,131 @@ func TestUpgrade_PermissionDeniedClearError(t *testing.T) { } } +func TestValidateAssetURL(t *testing.T) { + allowed := []string{ + "https://github.com/civitai/cli/releases/download/v0.1.11/civitai_0.1.11_linux_amd64.tar.gz", + "https://objects.githubusercontent.com/foo/bar", + "https://release-assets.githubusercontent.com/foo/bar", + "https://github.com/checksums.txt", + } + for _, u := range allowed { + if err := validateAssetURL(u); err != nil { + t.Errorf("expected %q to be allowed, got: %v", u, err) + } + } + rejected := []string{ + "http://github.com/civitai/cli/releases/download/x.tar.gz", // https downgrade + "http://objects.githubusercontent.com/foo", // http CDN + "https://evil.example.com/civitai.tar.gz", // off-host + "https://github.com.evil.com/x", // host-suffix spoof + "ftp://github.com/x", // wrong scheme + "https://raw.githubusercontent.com/x", // GitHub but not a release host + "://nonsense", // unparseable + } + for _, u := range rejected { + if err := validateAssetURL(u); err == nil { + t.Errorf("expected %q to be REJECTED, but it passed", u) + } + } +} + +// TestUpgrade_RejectsHTTPAssetURL proves an http:// tarball asset URL aborts the +// upgrade BEFORE any download, leaving the binary untouched. This closes the +// self-referential-checksum foot-gun (an attacker-served http checksums+tarball +// pair). +func TestUpgrade_RejectsHTTPAssetURL(t *testing.T) { + withParseableVersion(t, "v0.1.10") + // Release JSON advertising an http:// tarball + an http:// checksums.txt. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + tarName := fmt.Sprintf("civitai_0.1.11_%s_%s.tar.gz", runtime.GOOS, runtime.GOARCH) + fmt.Fprintf(w, `{"tag_name":"v0.1.11","assets":[ + {"name":%q,"browser_download_url":"http://github.com/dl/%s"}, + {"name":"checksums.txt","browser_download_url":"http://github.com/dl/checksums.txt"}]}`, + tarName, tarName) + })) + t.Cleanup(srv.Close) + pointAtServer(t, srv.URL) + target := captureApply(t) + + exe := filepath.Join(t.TempDir(), "civitai") + original := []byte("ORIGINAL-UNTOUCHED") + if err := os.WriteFile(exe, original, 0o755); err != nil { + t.Fatal(err) + } + withExecutable(t, exe) + + var out bytes.Buffer + err := runUpgrade(&out, false, false) + if err == nil { + t.Fatal("expected an http:// asset URL to abort the upgrade") + } + if !strings.Contains(err.Error(), "not https") { + t.Errorf("error should mention the non-https scheme: %v", err) + } + if *target != "" { + t.Error("an http:// asset URL must NOT call applyUpdate") + } + got, _ := os.ReadFile(exe) + if !bytes.Equal(got, original) { + t.Errorf("binary must be untouched when the asset URL is rejected, got %q", got) + } +} + +// TestUpgrade_RejectsNonGitHubAssetHost proves an https asset URL on a host +// outside the allowlist aborts before download. +func TestUpgrade_RejectsNonGitHubAssetHost(t *testing.T) { + withParseableVersion(t, "v0.1.10") + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + tarName := fmt.Sprintf("civitai_0.1.11_%s_%s.tar.gz", runtime.GOOS, runtime.GOARCH) + fmt.Fprintf(w, `{"tag_name":"v0.1.11","assets":[ + {"name":%q,"browser_download_url":"https://evil.example.com/%s"}, + {"name":"checksums.txt","browser_download_url":"https://evil.example.com/checksums.txt"}]}`, + tarName, tarName) + })) + t.Cleanup(srv.Close) + pointAtServer(t, srv.URL) + target := captureApply(t) + withExecutable(t, filepath.Join(t.TempDir(), "civitai")) + + var out bytes.Buffer + err := runUpgrade(&out, false, false) + if err == nil || !strings.Contains(err.Error(), "not an allowed GitHub release host") { + t.Errorf("expected an off-host rejection, got: %v", err) + } + if *target != "" { + t.Error("an off-host asset URL must NOT call applyUpdate") + } +} + +// TestDownload_RejectsInsecureRedirect proves the asset download client refuses +// to follow an https->http downgrade redirect. +func TestDownload_RejectsInsecureRedirect(t *testing.T) { + // Stand up an http target the redirect would point at (must not be reached). + plain := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "ATTACKER-PAYLOAD") + })) + t.Cleanup(plain.Close) + + // A redirector that 302s to the http:// target. We allowlist its loopback + // host so the INITIAL URL passes validation and the redirect (not the first + // hop) is what gets rejected. + redir := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, plain.URL+"/payload", http.StatusFound) + })) + t.Cleanup(redir.Close) + if u, err := url.Parse(redir.URL); err == nil { + withAssetHost(t, u.Hostname()) + } + + _, err := download(t.Context(), redir.URL+"/start") + if err == nil { + t.Fatal("expected an insecure-redirect rejection") + } + if !strings.Contains(err.Error(), "insecure redirect") { + t.Errorf("error should flag the insecure redirect: %v", err) + } +} + func TestUpgrade_MissingAssetForPlatform(t *testing.T) { withParseableVersion(t, "v0.1.10") // Serve a release whose only asset is checksums.txt — no platform tarball.