BatchuKVM/server/service/download/service.go

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
}