From 8d2bef03a10ee64645f5511aab1e6d76d880c06e Mon Sep 17 00:00:00 2001 From: "qcqcqc@wsl" <1220204124@zust.edu.cn> Date: Wed, 23 Apr 2025 09:20:29 +0800 Subject: [PATCH] =?UTF-8?q?=E6=8B=86=E5=88=86=E5=87=BD=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- version/cmd/main.go | 2 +- version/handler/handler.go | 145 ------------------------ version/pkg/handler/handler.go | 194 +++++++++++++++++++++++++++++++++ 3 files changed, 195 insertions(+), 146 deletions(-) delete mode 100644 version/handler/handler.go create mode 100644 version/pkg/handler/handler.go diff --git a/version/cmd/main.go b/version/cmd/main.go index 1e00697..c0c91e3 100644 --- a/version/cmd/main.go +++ b/version/cmd/main.go @@ -4,7 +4,7 @@ import ( "bash_go_service/shared/pkg/constants" "bash_go_service/shared/pkg/logger" "fmt" - "version/handler" + "version/pkg/handler" "github.com/spf13/viper" ) diff --git a/version/handler/handler.go b/version/handler/handler.go deleted file mode 100644 index fd6122f..0000000 --- a/version/handler/handler.go +++ /dev/null @@ -1,145 +0,0 @@ -package handler - -import ( - "bash_go_service/shared/pkg/client" - "bash_go_service/shared/pkg/constants" - "bash_go_service/shared/pkg/logger" - "crypto/md5" - "fmt" - "io" - "os" - "os/exec" - "path/filepath" - "time" - "version/api" -) - -const ( - updaterExeName = "updater-main" -) - -func DoUpdate() { - logger.Debug("start check updating") - res, err := api.GetVersion() - if err != nil { - logger.Error("Failed to get version: %v", err) - return - } - logger.Debug("res: %v", res) - - if !res.NeedUpdate { - logger.Info("No update needed.") - return - } - - exePath, err := os.Executable() - if err != nil { - logger.Error("Failed to get executable path: %v", err) - return - } - exeDir := filepath.Dir(exePath) - exeBase := filepath.Base(exePath) - - // 备份旧程序 - backupName := fmt.Sprintf("%s.bak.%s", exeBase, time.Now().Format("20060102_150405")) - backupPath := filepath.Join(exeDir, backupName) - logger.Debug("backupPath: %v", backupPath) - err = copyFile(exePath, backupPath) - if err != nil { - logger.Error("Failed to backup current executable: %v", err) - return - } - logger.Info("Backup created: %s", backupPath) - - // 下载新程序到临时路径 - newExePath := filepath.Join(exeDir, exeBase+".new") - cli := client.NewClient() - query := map[string]string{ - "version": res.Version, - } - err = cli.Download(constants.DownloadNewApi, query, newExePath) - if err != nil { - logger.Error("Download failed: %v", err) - _ = os.Rename(backupPath, exePath) // rollback - return - } - logger.Info("New executable downloaded to: %s", newExePath) - - // 🔥 校验 MD5 - if res.MD5 == "" { - logger.Error("MD5 is empty") - return - } - match, Merr := checkFileMD5(newExePath, res.MD5) - if Merr != nil { - logger.Error("Failed to check MD5: %v", Merr) - _ = os.Rename(backupPath, exePath) - return - } - if !match { - logger.Error("MD5 mismatch: expected %s", res.MD5) - _ = os.Rename(backupPath, exePath) - return - } - logger.Info("MD5 checksum verified.") - - // 设置新文件执行权限 - err = os.Chmod(newExePath, 0755) - if err != nil { - logger.Error("Failed to set executable permission: %v", err) - _ = os.Remove(newExePath) - _ = copyFile(backupPath, exePath) - return - } - logger.Info("Executable permission set for new file.") - - // 启动 updater 进行替换 - updaterPath := filepath.Join(exeDir, updaterExeName) - cmd := exec.Command(updaterPath, backupPath, newExePath, exePath, "true") - cmd.Stdout = os.Stdout - cmd.Stderr = os.Stderr - - err = cmd.Start() - if err != nil { - logger.Error("Failed to start updater: %v", err) - _ = os.Rename(backupPath, exePath) - return - } - logger.Info("Updater started, exiting current process.") - os.Exit(0) -} - -// MD5 校验函数 -func checkFileMD5(filePath, expectedMD5 string) (bool, error) { - f, err := os.Open(filePath) - if err != nil { - return false, err - } - defer f.Close() - - hash := md5.New() - if _, err := io.Copy(hash, f); err != nil { - return false, err - } - actualMD5 := fmt.Sprintf("%x", hash.Sum(nil)) - logger.Info("actual MD5: %v", actualMD5) - return actualMD5 == expectedMD5, nil -} - -// 复制文件函数 -func copyFile(src, dst string) error { - sourceFile, err := os.Open(src) - if err != nil { - return err - } - defer sourceFile.Close() - - destFile, err := os.Create(dst) - if err != nil { - return err - } - defer destFile.Close() - - _, err = io.Copy(destFile, sourceFile) - return err -} diff --git a/version/pkg/handler/handler.go b/version/pkg/handler/handler.go new file mode 100644 index 0000000..1168125 --- /dev/null +++ b/version/pkg/handler/handler.go @@ -0,0 +1,194 @@ +package handler + +import ( + "bash_go_service/shared/pkg/client" + "bash_go_service/shared/pkg/constants" + "bash_go_service/shared/pkg/logger" + "crypto/md5" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "time" + "version/api" +) + +const ( + updaterExeName = "updater-main" +) + +// DoUpdate checks for updates, downloads if needed, and executes the update. +func DoUpdate() { + logger.Debug("start check updating") + + needUpdate, newVersion, md5sum, err := checkVersionUpdate() + if err != nil { + logger.Error("Failed to check for updates: %v", err) + return + } + + if !needUpdate { + logger.Info("No update needed.") + return + } + + currentExePath, err := os.Executable() + if err != nil { + logger.Error("Failed to get executable path: %v", err) + return + } + currentExeDir := filepath.Dir(currentExePath) + currentExeBase := filepath.Base(currentExePath) + + backupPath, err := backupCurrentExecutable(currentExePath, currentExeBase, currentExeDir) + if err != nil { + logger.Error("Failed to backup current executable: %v", err) + return + } + + // 这只下载exe文件进行更新，如果需要更新其他，可以自定义逻辑 + newExePath, err := downloadNewExecutable(currentExeDir, currentExeBase, newVersion, "exe") + if err != nil { + logger.Error("Failed to download new executable: %v", err) + // Attempt rollback + _ = os.Rename(backupPath, currentExePath) + return + } + + if err := verifyMD5(newExePath, md5sum); err != nil { + logger.Error("MD5 verification failed: %v", err) + // Attempt rollback + _ = os.Rename(backupPath, currentExePath) + _ = os.Remove(newExePath) + return + } + + if err := executeUpdate(currentExeDir, backupPath, newExePath, currentExePath); err != nil { + logger.Error("Failed to execute update: %v", err) + // Attempt rollback + _ = os.Rename(backupPath, currentExePath) + _ = os.Remove(newExePath) + return + } + + logger.Info("Update process initiated, current process exiting.") + os.Exit(0) +} + +// checkVersionUpdate checks if a new version is available. +func checkVersionUpdate() (needUpdate bool, newVersion string, md5sum string, err error) { + res, err := api.GetVersion() + if err != nil { + logger.Error("Failed to get version information: %v", err) + return false, "", "", err + } + logger.Debug("version response: %v", res) + return res.NeedUpdate, res.Version, res.MD5, nil +} + +// backupCurrentExecutable creates a backup of the current executable. +func backupCurrentExecutable(exePath, exeBase, exeDir string) (backupPath string, err error) { + backupName := fmt.Sprintf("%s.bak.%s", exeBase, time.Now().Format("20060102_150405")) + backupPath = filepath.Join(exeDir, backupName) + logger.Debug("backup path: %v", backupPath) + err = copyFile(exePath, backupPath) + if err != nil { + logger.Error("Failed to backup current executable: %v", err) + return "", err + } + logger.Info("Backup created: %s", backupPath) + return backupPath, nil +} + +// downloadNewExecutable downloads the new executable file. +func downloadNewExecutable(exeDir, exeBase, newVersion string, filetype string) (newExePath string, err error) { + newExePath = filepath.Join(exeDir, exeBase+".new") + cli := client.NewClient() + query := map[string]string{ + "version": newVersion, + "filetype": filetype, + } + err = cli.Download(constants.DownloadNewApi, query, newExePath) + if err != nil { + logger.Error("Download failed: %v", err) + return "", err + } + logger.Info("New executable downloaded to: %s", newExePath) + return newExePath, nil +} + +// verifyMD5 checks if the downloaded file's MD5 matches the expected value. +func verifyMD5(filePath, expectedMD5 string) error { + if expectedMD5 == "" { + logger.Error("MD5 is empty, skipping verification.") + return nil + } + match, err := checkFileMD5(filePath, expectedMD5) + if err != nil { + logger.Error("Failed to check MD5: %v", err) + return err + } + if !match { + logger.Error("MD5 mismatch: expected %s", expectedMD5) + return fmt.Errorf("MD5 mismatch: expected %s", expectedMD5) + } + logger.Info("MD5 checksum verified.") + err = os.Chmod(filePath, 0755) + if err != nil { + logger.Error("Failed to set executable permission for new file: %v", err) + return err + } + logger.Info("Executable permission set for new file.") + return nil +} + +// executeUpdate starts the updater process to replace the current executable. +func executeUpdate(exeDir, backupPath, newExePath, currentExePath string) error { + updaterPath := filepath.Join(exeDir, updaterExeName) + cmd := exec.Command(updaterPath, backupPath, newExePath, currentExePath, "true") + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + + err := cmd.Start() + if err != nil { + logger.Error("Failed to start updater: %v", err) + return err + } + return nil +} + +// checkFileMD5 calculates the MD5 hash of a file and compares it with the expected MD5. +func checkFileMD5(filePath, expectedMD5 string) (bool, error) { + f, err := os.Open(filePath) + if err != nil { + return false, err + } + defer f.Close() + + hash := md5.New() + if _, err := io.Copy(hash, f); err != nil { + return false, err + } + actualMD5 := fmt.Sprintf("%x", hash.Sum(nil)) + logger.Debug("actual MD5: %v", actualMD5) + return actualMD5 == expectedMD5, nil +} + +// copyFile copies a file from source to destination. +func copyFile(src, dst string) error { + sourceFile, err := os.Open(src) + if err != nil { + return err + } + defer sourceFile.Close() + + destFile, err := os.Create(dst) + if err != nil { + return err + } + defer destFile.Close() + + _, err = io.Copy(destFile, sourceFile) + return err +}