BatchuKVM/server/service/application/manual_update.go

package application

import (
	"NanoKVM-Server/proto"
	"NanoKVM-Server/utils"
	"os"
	"os/exec"
	"path/filepath"
	"time"

	"github.com/gin-gonic/gin"
	log "github.com/sirupsen/logrus"
)

func (s *Service) UpdateServer(c *gin.Context) {
	var rsp proto.Response

	file, err := c.FormFile("file")
	if err != nil {
		rsp.ErrRsp(c, -1, "invalid file")
		return
	}

	// Get executable path
	execPath, err := os.Executable()
	if err != nil {
		rsp.ErrRsp(c, -2, "get executable path failed")
		return
	}
	execDir := filepath.Dir(execPath)

	// Save to temp file in SAME directory to ensure atomic rename
	tmpPath := filepath.Join(execDir, "NanoKVM-Server.new")
	if err := c.SaveUploadedFile(file, tmpPath); err != nil {
		rsp.ErrRsp(c, -3, "save file failed")
		return
	}

	// Verify file size
	info, err := os.Stat(tmpPath)
	if err != nil || info.Size() == 0 {
		_ = os.Remove(tmpPath)
		rsp.ErrRsp(c, -4, "invalid file size")
		return
	}

	// Verify it's executable
	if err := os.Chmod(tmpPath, 0755); err != nil {
		_ = os.Remove(tmpPath)
		rsp.ErrRsp(c, -5, "chmod failed")
		return
	}

	// Backup current binary
	backupPath := execPath + ".bak"
	// Try to remove old backup first
	_ = os.Remove(backupPath)
	
	if err := os.Rename(execPath, backupPath); err != nil {
		// If rename fails (e.g. running binary locked?), try copy
		if err := utils.CopyFile(execPath, backupPath); err != nil {
			log.Warnf("backup failed: %v", err)
			// Proceed with caution or fail? 
			// If we can't backup, maybe we shouldn't proceed.
			// But for embedded systems, sometimes we just overwrite.
		}
	}

	// Replace binary
	if err := os.Rename(tmpPath, execPath); err != nil {
		// Attempt rollback
		_ = os.Rename(backupPath, execPath)
		rsp.ErrRsp(c, -6, "replace binary failed")
		return
	}

	rsp.OkRsp(c)
	
	// Restart service in background
	go func() {
		time.Sleep(1 * time.Second)
		_ = exec.Command("sh", "-c", "/etc/init.d/S95nanokvm restart").Run()
	}()
}

func (s *Service) UpdateWeb(c *gin.Context) {
	var rsp proto.Response

	file, err := c.FormFile("file")
	if err != nil {
		rsp.ErrRsp(c, -1, "invalid file")
		return
	}

	// Get executable path
	execPath, err := os.Executable()
	if err != nil {
		rsp.ErrRsp(c, -2, "get executable path failed")
		return
	}
	execDir := filepath.Dir(execPath)

	// Save to temp file in SAME directory
	tmpPath := filepath.Join(execDir, "web.tar.gz.new")
	if err := c.SaveUploadedFile(file, tmpPath); err != nil {
		rsp.ErrRsp(c, -3, "save file failed")
		return
	}

	// Verify file size
	info, err := os.Stat(tmpPath)
	if err != nil || info.Size() == 0 {
		_ = os.Remove(tmpPath)
		rsp.ErrRsp(c, -4, "invalid file size")
		return
	}

	// Define web directory
	webDir := filepath.Join(execDir, "web")
	
	// Extract to temp dir first
	extractDir := filepath.Join(execDir, "web_extract_tmp")
	os.RemoveAll(extractDir)
	os.MkdirAll(extractDir, 0755)

	if _, err := utils.UnTarGz(tmpPath, extractDir); err != nil {
		os.RemoveAll(extractDir)
		_ = os.Remove(tmpPath)
		rsp.ErrRsp(c, -5, "extract failed")
		return
	}

	// Check if there is a single directory inside extractDir (e.g. 'dist')
	entries, _ := os.ReadDir(extractDir)
	sourceDir := extractDir
	if len(entries) == 1 && entries[0].IsDir() {
		sourceDir = filepath.Join(extractDir, entries[0].Name())
	}

	// Backup old web dir
	backupWebDir := filepath.Join(execDir, "web.bak")
	os.RemoveAll(backupWebDir)
	_ = os.Rename(webDir, backupWebDir)

	// Move new web dir
	if err := os.Rename(sourceDir, webDir); err != nil {
		// Rollback
		_ = os.Rename(backupWebDir, webDir)
		os.RemoveAll(extractDir)
		_ = os.Remove(tmpPath)
		rsp.ErrRsp(c, -6, "replace web dir failed")
		return
	}

	// Cleanup
	os.RemoveAll(extractDir)
	os.RemoveAll(backupWebDir)
	_ = os.Remove(tmpPath)

	rsp.OkRsp(c)
}