From 0a1d6fa4db3b770cc59841649d92d1de7549436a Mon Sep 17 00:00:00 2001 From: Cuong Manh Le Date: Mon, 28 Jul 2025 17:51:11 +0700 Subject: [PATCH] feat: create commands_upgrade.go and add UpgradeCommand with complete logic Create separate file for upgrade command handling to improve code organization. Add UpgradeCommand struct with Upgrade method that includes all original logic: channel management, service restart, rollback handling, and version verification. Includes InitUpgradeCmd function with proper argument validation and privilege checks. --- cmd/cli/commands_upgrade.go | 209 ++++++++++++++++++++++++++++++++++++ 1 file changed, 209 insertions(+) create mode 100644 cmd/cli/commands_upgrade.go diff --git a/cmd/cli/commands_upgrade.go b/cmd/cli/commands_upgrade.go new file mode 100644 index 0000000..b6fc472 --- /dev/null +++ b/cmd/cli/commands_upgrade.go @@ -0,0 +1,209 @@ +package cli + +import ( + "context" + "errors" + "net/http" + "os" + "os/exec" + "strings" + "time" + + "github.com/kardianos/service" + "github.com/minio/selfupdate" + "github.com/spf13/cobra" +) + +const ( + upgradeChannelDev = "dev" + upgradeChannelProd = "prod" + upgradeChannelDefault = "default" +) + +// UpgradeCommand handles upgrade-related operations +type UpgradeCommand struct { + serviceManager *ServiceManager +} + +// NewUpgradeCommand creates a new upgrade command handler +func NewUpgradeCommand() (*UpgradeCommand, error) { + sm, err := NewServiceManager() + if err != nil { + return nil, err + } + + return &UpgradeCommand{ + serviceManager: sm, + }, nil +} + +// Upgrade performs the upgrade operation +func (uc *UpgradeCommand) Upgrade(cmd *cobra.Command, args []string) error { + upgradeChannel := map[string]string{ + upgradeChannelDefault: "https://dl.controld.dev", + upgradeChannelDev: "https://dl.controld.dev", + upgradeChannelProd: "https://dl.controld.com", + } + if isStableVersion(curVersion()) { + upgradeChannel[upgradeChannelDefault] = upgradeChannel[upgradeChannelProd] + } + + bin, err := os.Executable() + if err != nil { + mainLog.Load().Fatal().Err(err).Msg("failed to get current ctrld binary path") + } + + // Create service config with executable path + sc := &service.Config{ + Name: ctrldServiceName, + DisplayName: "Control-D Helper Service", + Description: "A highly configurable, multi-protocol DNS forwarding proxy", + Option: service.KeyValue{}, + Executable: bin, + } + + readConfig(false) + v.Unmarshal(&cfg) + p := &prog{} + s, err := newService(p, sc) + if err != nil { + mainLog.Load().Error().Msg(err.Error()) + return nil + } + + if iface == "" { + iface = "auto" + } + p.preRun() + if ir := runningIface(s); ir != nil { + p.runningIface = ir.Name + p.requiredMultiNICsConfig = ir.All + } + + svcInstalled := true + if _, err := s.Status(); errors.Is(err, service.ErrNotInstalled) { + svcInstalled = false + } + + oldBin := bin + oldBinSuffix + baseUrl := upgradeChannel[upgradeChannelDefault] + if len(args) > 0 { + channel := args[0] + switch channel { + case upgradeChannelProd, upgradeChannelDev: // ok + default: + mainLog.Load().Fatal().Msgf("uprade argument must be either %q or %q", upgradeChannelProd, upgradeChannelDev) + } + baseUrl = upgradeChannel[channel] + } + + dlUrl := upgradeUrl(baseUrl) + mainLog.Load().Debug().Msgf("Downloading binary: %s", dlUrl) + + resp, err := getWithRetry(dlUrl, downloadServerIp) + if err != nil { + mainLog.Load().Fatal().Err(err).Msg("failed to download binary") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + mainLog.Load().Fatal().Msgf("could not download binary: %s", http.StatusText(resp.StatusCode)) + } + + mainLog.Load().Debug().Msg("Updating current binary") + if err := selfupdate.Apply(resp.Body, selfupdate.Options{OldSavePath: oldBin}); err != nil { + if rerr := selfupdate.RollbackError(err); rerr != nil { + mainLog.Load().Error().Err(rerr).Msg("could not rollback old binary") + } + mainLog.Load().Fatal().Err(err).Msg("failed to update current binary") + } + + doRestart := func() bool { + if !svcInstalled { + return true + } + tasks := []task{ + {s.Stop, true, "Stop"}, + {func() error { + // restore static DNS settings or DHCP + p.resetDNS(false, true) + return nil + }, false, "Cleanup"}, + {func() error { + time.Sleep(time.Second * 1) + return nil + }, false, "Waiting for service to stop"}, + } + doTasks(tasks) + + tasks = []task{ + {s.Start, true, "Start"}, + } + if doTasks(tasks) { + if dir, err := socketDir(); err == nil { + if cc := newSocketControlClient(context.TODO(), s, dir); cc != nil { + _, _ = cc.post(ifacePath, nil) + return true + } + } + } + return false + } + + if svcInstalled { + mainLog.Load().Debug().Msg("Restarting ctrld service using new binary") + } + + if doRestart() { + _ = os.Remove(oldBin) + _ = os.Chmod(bin, 0755) + ver := "unknown version" + out, err := exec.Command(bin, "--version").CombinedOutput() + if err != nil { + mainLog.Load().Warn().Err(err).Msg("Failed to get new binary version") + } + if after, found := strings.CutPrefix(string(out), "ctrld version "); found { + ver = after + } + mainLog.Load().Notice().Msgf("Upgrade successful - %s", ver) + return nil + } + + mainLog.Load().Warn().Msgf("Upgrade failed, restoring previous binary: %s", oldBin) + if err := os.Remove(bin); err != nil { + mainLog.Load().Fatal().Err(err).Msg("failed to remove new binary") + } + if err := os.Rename(oldBin, bin); err != nil { + mainLog.Load().Fatal().Err(err).Msg("failed to restore old binary") + } + if doRestart() { + mainLog.Load().Notice().Msg("Restored previous binary successfully") + return nil + } + + return nil +} + +// InitUpgradeCmd creates the upgrade command with proper logic +func InitUpgradeCmd() *cobra.Command { + upgradeCmd := &cobra.Command{ + Use: "upgrade", + Short: "Upgrading ctrld to latest version", + ValidArgs: []string{upgradeChannelDev, upgradeChannelProd}, + Args: cobra.MaximumNArgs(1), + PreRun: func(cmd *cobra.Command, args []string) { + checkHasElevatedPrivilege() + }, + RunE: func(cmd *cobra.Command, args []string) error { + uc, err := NewUpgradeCommand() + if err != nil { + return err + } + return uc.Upgrade(cmd, args) + }, + } + + rootCmd.AddCommand(upgradeCmd) + + return upgradeCmd +}