refactor: refactor offline update function

server/service/application/install.go

@@ -0,0 +1,76 @@
+package application
+
+import (
+	"fmt"
+	"os"
+	"sync"
+
+	"NanoKVM-Server/utils"
+	log "github.com/sirupsen/logrus"
+)
+
+var (
+	mutex      sync.Mutex
+	isUpdating bool
+)
+
+func acquireUpdateLock() bool {
+	mutex.Lock()
+	defer mutex.Unlock()
+
+	if isUpdating {
+		return false
+	}
+	isUpdating = true
+	return true
+}
+
+func releaseUpdateLock() {
+	mutex.Lock()
+	defer mutex.Unlock()
+	isUpdating = false
+}
+
+func installPackage(source string) error {
+	dir, err := utils.UnTarGz(source, CacheDir)
+	if err != nil {
+		return fmt.Errorf("failed to decompress app: %w", err)
+	}
+
+	if err := backupCurrentApp(); err != nil {
+		return err
+	}
+
+	if err := applyUpdate(dir); err != nil {
+		return err
+	}
+
+	if err := utils.ChmodRecursively(AppDir, 0o755); err != nil {
+		return fmt.Errorf("failed to chmod: %w", err)
+	}
+
+	return nil
+}
+
+func backupCurrentApp() error {
+	if err := os.RemoveAll(BackupDir); err != nil {
+		return fmt.Errorf("failed to remove backup: %w", err)
+	}
+
+	if err := utils.MoveFilesRecursively(AppDir, BackupDir); err != nil {
+		return fmt.Errorf("failed to backup app: %w", err)
+	}
+
+	return nil
+}
+
+func applyUpdate(sourceDir string) error {
+	if err := utils.MoveFilesRecursively(sourceDir, AppDir); err != nil {
+		// Try to restore backup on failure
+		if restoreErr := utils.MoveFilesRecursively(BackupDir, AppDir); restoreErr != nil {
+			log.Errorf("Failed to restore backup after update failure: %v", restoreErr)
+		}
+		return fmt.Errorf("failed to move update in place: %w", err)
+	}
+	return nil
+}

server/service/application/progress_writer.go

@@ -0,0 +1,84 @@
+package application
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"strings"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	sentinelPath       = "/tmp/.download_in_progress"
+	tickerInterval     = 2500 * time.Millisecond
+	sentinelPermission = 0644
+)
+
+// progressWriter wraps an io.Writer to track and report upload progress.
+type progressWriter struct {
+	writer    io.Writer
+	written   int64
+	totalSize int64
+	ticker    *time.Ticker
+	done      chan struct{}
+}
+
+func newProgressWriter(w io.Writer, totalSize int64) *progressWriter {
+	pw := &progressWriter{
+		writer:    w,
+		totalSize: totalSize,
+		ticker:    time.NewTicker(tickerInterval),
+		done:      make(chan struct{}),
+	}
+
+	go pw.reportProgress()
+	return pw
+}
+
+func (pw *progressWriter) reportProgress() {
+	for {
+		select {
+		case <-pw.done:
+			return
+		case <-pw.ticker.C:
+			pw.updateSentinel()
+		}
+	}
+}
+
+func (pw *progressWriter) updateSentinel() {
+	if pw.totalSize == 0 {
+		return
+	}
+
+	content, err := os.ReadFile(sentinelPath)
+	if err != nil {
+		log.Error("Failed to read sentinel file")
+		return
+	}
+
+	parts := strings.SplitN(string(content), ";", 2)
+	if len(parts) == 0 {
+		return
+	}
+
+	percentage := float64(pw.written) / float64(pw.totalSize) * 100
+	newContent := fmt.Sprintf("%s;%.2f%%", parts[0], percentage)
+
+	if err := os.WriteFile(sentinelPath, []byte(newContent), sentinelPermission); err != nil {
+		log.Error("Failed to update sentinel file")
+	}
+}
+
+func (pw *progressWriter) Write(p []byte) (int, error) {
+	n, err := pw.writer.Write(p)
+	pw.written += int64(n)
+	return n, err
+}
+
+func (pw *progressWriter) Stop() {
+	pw.ticker.Stop()
+	close(pw.done)
+}

server/service/application/update.go

@@ -8,11 +8,7 @@ import (
 	"net/http"
 	"os"
 	"os/exec"
-	"sync"
 	"time"
-	"path/filepath"
-	"strings"
-	"regexp"
 
 	"github.com/gin-gonic/gin"
 	log "github.com/sirupsen/logrus"
@@ -21,235 +17,18 @@ import (
 	"NanoKVM-Server/utils"
 )
 
-var sentinelPath = "/tmp/.download_in_progress"
-
 const (
 	maxTries = 3
 )
 
-var (
-	updateMutex sync.Mutex
-	isUpdating  bool
-)
-
-func (s *Service) UploadUpdate(c *gin.Context) {
-	var rsp proto.Response
-
-	updateMutex.Lock()
-	if isUpdating {
-		updateMutex.Unlock()
-		rsp.ErrRsp(c, -1, "update already in progress")
-		return
-	}
-	isUpdating = true
-	updateMutex.Unlock()
-
-	defer func() {
-		updateMutex.Lock()
-		isUpdating = false
-		updateMutex.Unlock()
-	}()
-
-	if err := uploadupdate(rsp, c); err != nil {
-		rsp.ErrRsp(c, -1, fmt.Sprintf("update failed: %s", err))
-		return
-	}
-
-	rsp.OkRsp(c)
-	log.Debugf("update application success")
-
-	// Sleep for a second before restarting the device
-	time.Sleep(1 * time.Second)
-
-	_ = exec.Command("sh", "-c", "/etc/init.d/S95nanokvm restart").Run()
-}
-
-func uploadupdate(rsp proto.Response, c *gin.Context) error {
-	_ = os.RemoveAll(CacheDir)
-	_ = os.MkdirAll(CacheDir, 0o755)
-	defer func() {
-		_ = os.RemoveAll(CacheDir)
-	}()
-
-	// Set a sentinel file to mark that there is a download in progress
-	// This is to prevent multiple downloads at the same time
-	if _, err := os.Stat(sentinelPath); err == nil {
-		log.Debug("Download in progress")
-		rsp.ErrRsp(c, -1, "download in progress")
-		return err
-	}
-
-	// Create the sentinel file
-	err := os.WriteFile(sentinelPath, []byte("start"), 0644)
-	if err != nil {
-		log.Error("Failed to create sentinel file")
-		rsp.ErrRsp(c, -1, "failed to create sentinel file")
-		return err
-	}
-
-    // Multipart Reader direkt nutzen (keine FormFile!)
-    reader, err := c.Request.MultipartReader()
-    if err != nil {
-		log.Error("invalid multipart data")
-		rsp.ErrRsp(c, -1, "invalid multipart data")
-		defer os.Remove(sentinelPath)
-        return err
-    }
-
-	
-	var lw *loggingWriter
-	var outPath = ""
-
-    for {
-        part, err := reader.NextPart()
-        if err == io.EOF {
-            break
-        }
-        if err != nil {
-			lw.stopTicker()
-			defer os.Remove(sentinelPath)
-            return err
-        }
-
-        if part.FormName() != "file" {
-            continue
-        }
-
-        filename := part.FileName()
-        if filename == "" {
-			lw.stopTicker()
-			defer os.Remove(sentinelPath)
-            return fmt.Errorf("no filename")
-        }
-
-		filename = filepath.Base(filename)
-
-		if filename != part.FileName() {
-			defer os.Remove(sentinelPath)
-			return fmt.Errorf("path detected in filename")
-		}
-
-		if strings.Contains(filename, "..") {
-			log.Warn("path traversal attempt")
-			rsp.ErrRsp(c, -1, "invalid filename")
-			defer os.Remove(sentinelPath)
-			return fmt.Errorf("path traversal attempt")
-		}
-
-		valid := regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)
-		if !valid.MatchString(filename) {
-			rsp.ErrRsp(c, -1, "invalid filename")
-			defer os.Remove(sentinelPath)
-			return fmt.Errorf("err4")
-		}
-
-		data, err := os.ReadFile(sentinelPath)
-		if err != nil {
-			lw.stopTicker()
-			defer os.Remove(sentinelPath)
-			return err
-		}
-
-        outPath = "/data/" + filename
-        out, err := os.Create(outPath)
-        if err != nil {
-			lw.stopTicker()
-			defer os.Remove(sentinelPath)
-            return err
-        }
-        defer out.Close()
-
-		if strings.Contains(string(data), "start") {
-			err = os.WriteFile(sentinelPath, []byte(filename), 0644)
-			if err != nil {
-				lw.stopTicker()
-				defer os.Remove(outPath)
-				defer os.Remove(sentinelPath)
-				return err
-			}
-			
-			lw = &loggingWriter{writer: out, totalSize: c.Request.ContentLength}
-			lw.startTicker()
-		} else {
-			if !strings.Contains(string(data), filename) {
-				lw.stopTicker()
-				defer os.Remove(outPath)
-				defer os.Remove(sentinelPath)
-				return fmt.Errorf("failed")
-			}
-		}
-
-        // Direkt streamen → kein RAM-Bedarf außer kleinem Buffer
-        _, err = io.Copy(lw, part)
-        if err != nil {
-			lw.stopTicker()
-			defer os.Remove(outPath)
-			defer os.Remove(sentinelPath)
-            return err
-        }
-    }
-	lw.stopTicker()
-
-	rsp.OkRspWithData(c, &proto.StatusImageRsp{
-		Status:     "idle",
-		File:       "",
-		Percentage: "",
-	})
-	
-	defer os.Remove(sentinelPath)
-
-	// decompress
-	dir, err := utils.UnTarGz(outPath, CacheDir)
-	if err != nil {
-		fmt.Errorf("decompress app failed: %s", err)
-		return err
-	}
-
-	// backup old version
-	if err := os.RemoveAll(BackupDir); err != nil {
-		fmt.Errorf("remove backup failed: %s", err)
-		return err
-	}
-
-	if err := utils.MoveFilesRecursively(AppDir, BackupDir); err != nil {
-		fmt.Errorf("backup app failed: %s", err)
-		return err
-	}
-
-	// update
-	if err := utils.MoveFilesRecursively(dir, AppDir); err != nil {
-		fmt.Errorf("failed to move update back in place: %s", err)
-		return err
-	}
-
-	// modify permissions
-	if err := utils.ChmodRecursively(AppDir, 0o755); err != nil {
-		fmt.Errorf("chmod failed: %s", err)
-		return err
-	}
-
-	defer os.Remove(outPath)
-	
-	return nil
-}
-
 func (s *Service) Update(c *gin.Context) {
 	var rsp proto.Response
 
-	updateMutex.Lock()
-	if isUpdating {
-		updateMutex.Unlock()
+	if !acquireUpdateLock() {
 		rsp.ErrRsp(c, -1, "update already in progress")
 		return
 	}
-	isUpdating = true
-	updateMutex.Unlock()
-
-	defer func() {
-		updateMutex.Lock()
-		isUpdating = false
-		updateMutex.Unlock()
-	}()
+	defer releaseUpdateLock()
 
 	if err := update(); err != nil {
 		rsp.ErrRsp(c, -1, fmt.Sprintf("update failed: %s", err))
@@ -280,7 +59,6 @@ func update() error {
 
 	// download
 	target := fmt.Sprintf("%s/%s", CacheDir, latest.Name)
-
 	if err := download(latest.Url, target); err != nil {
 		log.Errorf("download app failed: %s", err)
 		return err
@@ -292,34 +70,9 @@ func update() error {
 		return err
 	}
 
-	// decompress
-	dir, err := utils.UnTarGz(target, CacheDir)
-	log.Debugf("untar: %s", dir)
-	if err != nil {
-		log.Errorf("decompress app failed: %s", err)
-		return err
-	}
-
-	// backup old version
-	if err := os.RemoveAll(BackupDir); err != nil {
-		log.Errorf("remove backup failed: %s", err)
-		return err
-	}
-
-	if err := utils.MoveFilesRecursively(AppDir, BackupDir); err != nil {
-		log.Errorf("backup app failed: %s", err)
-		return err
-	}
-
-	// update
-	if err := utils.MoveFilesRecursively(dir, AppDir); err != nil {
-		log.Errorf("failed to move update back in place: %s", err)
-		return err
-	}
-
-	// modify permissions
-	if err := utils.ChmodRecursively(AppDir, 0o755); err != nil {
-		log.Errorf("chmod failed: %s", err)
+	// install
+	if err := installPackage(target); err != nil {
+		log.Errorf("failed to install package: %v", err)
 		return err
 	}
 
@@ -330,7 +83,7 @@ func download(url string, target string) (err error) {
 	for i := range maxTries {
 		log.Debugf("attempt #%d/%d", i+1, maxTries)
 		if i > 0 {
-			time.Sleep(time.Second * 3) // wait for 3 seconds before retrying the download attempt
+			time.Sleep(time.Second * 3)
 		}
 
 		var req *http.Request
@@ -378,54 +131,3 @@ func checksum(filePath string, expectedHash string) error {
 
 	return nil
 }
-
-type loggingWriter struct {
-	writer    io.Writer
-	total     int64
-	totalSize int64
-	ticker    *time.Ticker
-	done      chan bool
-}
-
-func (lw *loggingWriter) startTicker() {
-	lw.ticker = time.NewTicker(2500 * time.Millisecond)
-	lw.done = make(chan bool)
-	go func() {
-		for {
-			select {
-			case <-lw.done:
-				return
-			case <-lw.ticker.C:
-				lw.updateSentinel()
-			}
-		}
-	}()
-}
-
-func (lw *loggingWriter) stopTicker() {
-	lw.ticker.Stop()
-	lw.done <- true
-}
-
-func (lw *loggingWriter) updateSentinel() {
-	percentage := float64(lw.total) / float64(lw.totalSize) * 100
-	content, err := os.ReadFile(sentinelPath)
-	if err != nil {
-		log.Error("Failed to read sentinel file")
-		return
-	}
-	splitted := strings.Split(string(content), ";")
-	if len(splitted) == 0 {
-		return
-	}
-	err = os.WriteFile(sentinelPath, []byte(fmt.Sprintf("%s;%.2f%%", splitted[0], percentage)), 0644)
-	if err != nil {
-		log.Error("Failed to update sentinel file")
-	}
-}
-
-func (lw *loggingWriter) Write(p []byte) (int, error) {
-	n, err := lw.writer.Write(p)
-	lw.total += int64(n)
-	return n, err
-}

server/service/application/update_offline.go

@@ -0,0 +1,176 @@
+package application
+
+import (
+	"fmt"
+	"io"
+	"mime/multipart"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"regexp"
+	"strings"
+	"time"
+
+	"NanoKVM-Server/proto"
+	"github.com/gin-gonic/gin"
+	log "github.com/sirupsen/logrus"
+)
+
+var validFilenameRegex = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)
+
+func (s *Service) OfflineUpdate(c *gin.Context) {
+	var rsp proto.Response
+
+	if !acquireUpdateLock() {
+		rsp.ErrRsp(c, -1, "update already in progress")
+		return
+	}
+	defer releaseUpdateLock()
+
+	if err := offlineUpdate(c); err != nil {
+		rsp.ErrRsp(c, -1, fmt.Sprintf("update failed: %s", err))
+		return
+	}
+
+	rsp.OkRsp(c)
+	log.Debugf("offline update application success")
+
+	time.Sleep(1 * time.Second)
+	_ = exec.Command("sh", "-c", "/etc/init.d/S95nanokvm restart").Run()
+}
+
+func offlineUpdate(c *gin.Context) error {
+	_ = os.RemoveAll(CacheDir)
+	_ = os.MkdirAll(CacheDir, 0o755)
+	defer func() {
+		_ = os.RemoveAll(CacheDir)
+	}()
+
+	if err := checkDownloadInProgress(); err != nil {
+		return err
+	}
+
+	if err := createSentinelFile(); err != nil {
+		return err
+	}
+	defer removeSentinelFile()
+
+	reader, err := c.Request.MultipartReader()
+	if err != nil {
+		log.Errorf("Invalid multipart data: %v", err)
+		return fmt.Errorf("invalid multipart data: %w", err)
+	}
+
+	target, err := processUpload(reader, c.Request.ContentLength)
+	if err != nil {
+		log.Errorf("failed to upload install package: %v", err)
+		return err
+	}
+
+	if err := installPackage(target); err != nil {
+		log.Errorf("failed to install package: %v", err)
+		return err
+	}
+
+	return nil
+}
+
+func checkDownloadInProgress() error {
+	if _, err := os.Stat(sentinelPath); err == nil {
+		log.Debug("Download in progress")
+		return fmt.Errorf("download already in progress")
+	}
+	return nil
+}
+
+func createSentinelFile() error {
+	if err := os.WriteFile(sentinelPath, []byte("downloading"), sentinelPermission); err != nil {
+		log.Errorf("Failed to create sentinel file: %v", err)
+		return fmt.Errorf("failed to create sentinel file: %w", err)
+	}
+	return nil
+}
+
+func processUpload(reader *multipart.Reader, contentLength int64) (string, error) {
+	var outPath string
+
+	for {
+		part, err := reader.NextPart()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return "", fmt.Errorf("failed to read multipart: %w", err)
+		}
+
+		if part.FormName() != "file" {
+			continue
+		}
+
+		outPath, err = saveUploadedFile(part, contentLength)
+		if err != nil {
+			return "", err
+		}
+	}
+
+	if outPath == "" {
+		return "", fmt.Errorf("no file uploaded")
+	}
+
+	return outPath, nil
+}
+
+func saveUploadedFile(part *multipart.Part, contentLength int64) (string, error) {
+	filename := part.FileName()
+	if filename == "" {
+		return "", fmt.Errorf("no filename provided")
+	}
+
+	if err := validateFilename(filename); err != nil {
+		return "", err
+	}
+
+	outPath := filepath.Join(CacheDir, filename)
+	out, err := os.Create(outPath)
+	if err != nil {
+		return "", fmt.Errorf("failed to create output file: %w", err)
+	}
+	defer out.Close()
+
+	pw := newProgressWriter(out, contentLength)
+	defer pw.Stop()
+
+	if _, err := io.Copy(pw, part); err != nil {
+		return "", fmt.Errorf("failed to write file: %w", err)
+	}
+
+	return outPath, nil
+}
+
+func validateFilename(filename string) error {
+	baseName := filepath.Base(filename)
+
+	// Check if the path contains directory components
+	if baseName != filename {
+		log.Warnf("Path detected in filename: %s", filename)
+		return fmt.Errorf("path detected in filename")
+	}
+
+	// Check for path traversal attempts
+	if strings.Contains(filename, "..") {
+		log.Warnf("Path traversal attempt: %s", filename)
+		return fmt.Errorf("invalid filename: path traversal detected")
+	}
+
+	// Validate filename characters
+	if !validFilenameRegex.MatchString(filename) {
+		log.Warnf("Invalid filename characters: %s", filename)
+		return fmt.Errorf("invalid filename: contains invalid characters")
+	}
+
+	return nil
+}
+
+func removeSentinelFile() {
+	_ = os.Remove(sentinelPath)
+}

server/service/application/version.go

@@ -1,7 +1,6 @@
 package application
 
 import (
-	"NanoKVM-Server/proto"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -10,6 +9,8 @@ import (
 	"strings"
 	"time"
 
+	"NanoKVM-Server/proto"
+
 	"github.com/gin-gonic/gin"
 	log "github.com/sirupsen/logrus"
 )
@@ -36,33 +37,15 @@ func (s *Service) GetVersion(c *gin.Context) {
 	log.Debugf("current version: %s", currentVersion)
 
 	// latest version
+	latestVersion := ""
 	latest, err := getLatest()
-	if err != nil {
-		rsp.ErrRsp(c, -1, "get latest version failed")
-		return
+	if err == nil && latest != nil {
+		latestVersion = latest.Version
 	}
 
 	rsp.OkRspWithData(c, &proto.GetVersionRsp{
 		Current: currentVersion,
-		Latest:  latest.Version,
-	})
-}
-
-func (s *Service) GetCurrentVersion(c *gin.Context) {
-	var rsp proto.Response
-
-	// current version
-	currentVersion := "1.0.0"
-
-	versionFile := fmt.Sprintf("%s/version", AppDir)
-	if version, err := os.ReadFile(versionFile); err == nil {
-		currentVersion = strings.ReplaceAll(string(version), "\n", "")
-	}
-
-	log.Debugf("current version: %s", currentVersion)
-
-	rsp.OkRspWithData(c, &proto.GetVersionRsp{
-		Current: currentVersion,
+		Latest:  latestVersion,
 	})
 }