Refactor: Rename NanoKVM to BatchuKVM and update server URL

server/service/download/service.go

@@ -0,0 +1,230 @@
+package download
+
+import (
+	"NanoKVM-Server/proto"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	log "github.com/sirupsen/logrus"
+	"io"
+	"net/http"
+	"net/url"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+)
+
+type Service struct{}
+
+var sentinelPath = "/tmp/.download_in_progress"
+
+func NewService() *Service {
+	// Clear sentinel
+	// If we are starting from scratch, we need to remove the sentinel file as any downloads at this point are done or broken
+	_ = os.Remove(sentinelPath)
+	return &Service{}
+}
+
+func (s *Service) ImageEnabled(c *gin.Context) {
+	var rsp proto.Response
+
+	// Check if /data mount is RO/RW
+	testFile := "/data/.testfile"
+	file, err := os.Create(testFile)
+	defer file.Close()
+	defer os.Remove(testFile)
+	if err != nil {
+		if os.IsPermission(err) {
+			rsp.OkRspWithData(c, &proto.ImageEnabledRsp{
+				Enabled: false,
+			})
+			return
+		}
+		rsp.OkRspWithData(c, &proto.ImageEnabledRsp{
+			Enabled: false,
+		})
+		return
+	}
+
+	rsp.OkRspWithData(c, &proto.ImageEnabledRsp{
+		Enabled: true,
+	})
+}
+
+func (s *Service) StatusImage(c *gin.Context) {
+	var rsp proto.Response
+
+	// Check if the sentinel file exists
+	log.Debug("StatusImage")
+	if _, err := os.Stat(sentinelPath); err == nil {
+		content, err := os.ReadFile(sentinelPath)
+		if err != nil {
+			log.Error("Failed to read sentinel file")
+			rsp.OkRspWithData(c, &proto.StatusImageRsp{
+				Status:     "in_progress",
+				File:       "",
+				Percentage: "",
+			})
+			return
+		}
+		splitted := strings.Split(string(content), ";")
+		if len(splitted) == 1 {
+			// No percentage, just the URL
+			rsp.OkRspWithData(c, &proto.StatusImageRsp{
+				Status:     "in_progress",
+				File:       splitted[0],
+				Percentage: "",
+			})
+		} else {
+			// Percentage is available
+			rsp.OkRspWithData(c, &proto.StatusImageRsp{
+				Status:     "in_progress",
+				File:       splitted[0],
+				Percentage: splitted[1],
+			})
+		}
+
+		return
+	}
+	rsp.OkRspWithData(c, &proto.StatusImageRsp{
+		Status:     "idle",
+		File:       "",
+		Percentage: "",
+	})
+}
+
+func (s *Service) DownloadImage(c *gin.Context) {
+	var req proto.MountImageReq
+	var rsp proto.Response
+
+	log.Debug("DownloadImage")
+
+	if err := proto.ParseFormRequest(c, &req); err != nil {
+		rsp.ErrRsp(c, -1, "invalid arguments")
+		return
+	}
+
+	if req.File == "" {
+		rsp.ErrRsp(c, -1, "invalid arguments")
+		return
+	}
+	// Parse the URI to see if its valid http/s
+	u, err := url.Parse(req.File)
+	if err != nil || u.Scheme == "" || u.Host == "" {
+		rsp.ErrRsp(c, -1, "invalid url")
+		return
+	}
+
+	// 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
+	}
+	// Create the sentinel file
+	err = os.WriteFile(sentinelPath, []byte(req.File), 0644)
+	if err != nil {
+		log.Error("Failed to create sentinel file")
+		rsp.ErrRsp(c, -1, "failed to create sentinel file")
+		return
+	}
+
+	// Check if it actually exists and fail if it doesn't
+	resp, err := http.Head(req.File)
+	if resp.StatusCode != http.StatusOK || err != nil {
+		rsp.ErrRsp(c, resp.StatusCode, "failed when checking the url")
+		log.Error("Failed to check the URL")
+		defer os.Remove(sentinelPath)
+		return
+	}
+	defer resp.Body.Close()
+
+	// Download the image in a goroutine to not block the request
+	go func() {
+		defer os.Remove(sentinelPath)
+		resp, err = http.Get(req.File)
+		if err != nil {
+			log.Error("Failed to download the file")
+			rsp.ErrRsp(c, -1, "failed to download the file")
+			return
+		}
+		defer resp.Body.Close()
+		// Create the destination file
+		destPath := filepath.Join("/data", filepath.Base(u.Path))
+		out, err := os.Create(destPath)
+		if err != nil {
+			log.Error("Failed to create destination file")
+			rsp.ErrRsp(c, -1, "failed to create destination file")
+			return
+		}
+		defer out.Close()
+
+		lw := &loggingWriter{writer: out, totalSize: resp.ContentLength}
+		lw.startTicker()
+		_, err = io.Copy(lw, resp.Body)
+		if err != nil {
+			log.Error("Failed to save the file")
+			rsp.ErrRsp(c, -1, "failed to save the file")
+			lw.stopTicker()
+			return
+		}
+		lw.stopTicker()
+	}()
+	rsp.OkRspWithData(c, &proto.StatusImageRsp{
+		Status:     "in_progress",
+		File:       req.File,
+		Percentage: "",
+	})
+}
+
+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
+}