WarpBox/lib/server/server.go

package server

import (
	"archive/zip"
	"crypto/rand"
	"encoding/json"
	"encoding/hex"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/gin-contrib/gzip"
	"github.com/gin-gonic/gin"
)

const (
	uploadRoot       = "data/uploads"
	boxManifestFile  = ".warpbox.json"
	fileStatusFailed = "failed"
	fileStatusReady  = "complete"
	fileStatusWait   = "pending"
	fileStatusWork   = "uploading"
)

var boxManifestMu sync.Mutex

type boxFile struct {
	ID           string `json:"id"`
	Name         string `json:"name"`
	Size         int64  `json:"size"`
	SizeLabel    string `json:"size_label"`
	MimeType     string `json:"mime_type"`
	Status       string `json:"status"`
	StatusLabel  string `json:"status_label"`
	Title        string `json:"title"`
	IconPath     string `json:"icon_path"`
	DownloadPath string `json:"download_path"`
	UploadPath   string `json:"upload_path"`
	IsComplete   bool   `json:"is_complete"`
}

type boxManifest struct {
	Files []boxFile `json:"files"`
}

type createBoxRequest struct {
	Files []createBoxFileRequest `json:"files"`
}

type createBoxFileRequest struct {
	Name string `json:"name"`
	Size int64  `json:"size"`
}

type updateFileStatusRequest struct {
	Status string `json:"status"`
}

func Run(addr string) error {
	router := gin.Default()
	router.LoadHTMLGlob("templates/*.html")

	router.GET("/", func(ctx *gin.Context) {
		ctx.HTML(http.StatusOK, "index.html", gin.H{})
	})

	router.GET("/box/:id", func(ctx *gin.Context) {
		boxID := ctx.Param("id")
		if !validBoxID(boxID) {
			ctx.String(http.StatusBadRequest, "Invalid box id")
			return
		}

		files, err := listBoxFiles(boxID)
		if err != nil {
			ctx.String(http.StatusNotFound, "Box not found")
			return
		}

		ctx.HTML(http.StatusOK, "box.html", gin.H{
			"BoxID":       boxID,
			"Files":       files,
			"FileCount":   len(files),
			"DownloadAll": "/box/" + boxID + "/download",
		})
	})

	router.GET("/box/:id/status", func(ctx *gin.Context) {
		boxID := ctx.Param("id")
		if !validBoxID(boxID) {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "Invalid box id"})
			return
		}

		files, err := listBoxFiles(boxID)
		if err != nil {
			ctx.JSON(http.StatusNotFound, gin.H{"error": "Box not found"})
			return
		}

		ctx.JSON(http.StatusOK, gin.H{
			"box_id": boxID,
			"files":  files,
		})
	})

	router.GET("/box/:id/download", func(ctx *gin.Context) {
		boxID := ctx.Param("id")
		if !validBoxID(boxID) {
			ctx.String(http.StatusBadRequest, "Invalid box id")
			return
		}

		files, err := listBoxFiles(boxID)
		if err != nil {
			ctx.String(http.StatusNotFound, "Box not found")
			return
		}

		ctx.Header("Content-Type", "application/zip")
		ctx.Header("Content-Disposition", fmt.Sprintf(`attachment; filename="warpbox-%s.zip"`, boxID))

		zipWriter := zip.NewWriter(ctx.Writer)
		defer zipWriter.Close()

		for _, file := range files {
			if !file.IsComplete {
				continue
			}

			if err := addFileToZip(zipWriter, boxID, file.Name); err != nil {
				ctx.Status(http.StatusInternalServerError)
				return
			}
		}
	})

	router.GET("/box/:id/files/:filename", func(ctx *gin.Context) {
		boxID := ctx.Param("id")
		filename, ok := safeFilename(ctx.Param("filename"))
		if !validBoxID(boxID) || !ok {
			ctx.String(http.StatusBadRequest, "Invalid file")
			return
		}

		path := filepath.Join(boxPath(boxID), filename)
		if !strings.HasPrefix(path, boxPath(boxID)+string(filepath.Separator)) {
			ctx.String(http.StatusBadRequest, "Invalid file")
			return
		}

		if _, err := os.Stat(path); err != nil {
			ctx.String(http.StatusNotFound, "File not found")
			return
		}

		ctx.FileAttachment(path, filename)
	})

	router.POST("/box", func(ctx *gin.Context) {
		boxID, err := newBoxID()
		if err != nil {
			ctx.JSON(http.StatusInternalServerError, gin.H{"error": "Could not create upload box"})
			return
		}

		if err := os.MkdirAll(boxPath(boxID), 0755); err != nil {
			ctx.JSON(http.StatusInternalServerError, gin.H{"error": "Could not prepare upload box"})
			return
		}

		var request createBoxRequest
		if err := ctx.ShouldBindJSON(&request); err != nil && err != io.EOF {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "Invalid box payload"})
			return
		}

		files, err := createBoxManifest(boxID, request.Files)
		if err != nil {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
			return
		}

		ctx.JSON(http.StatusOK, gin.H{
			"box_id":  boxID,
			"box_url": "/box/" + boxID,
			"files":   files,
		})
	})

	router.POST("/box/:id/files/:file_id/upload", func(ctx *gin.Context) {
		boxID := ctx.Param("id")
		fileID := ctx.Param("file_id")
		if !validBoxID(boxID) {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "Invalid box id"})
			return
		}

		file, err := ctx.FormFile("file")
		if err != nil {
			markManifestFileStatus(boxID, fileID, fileStatusFailed)
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "No file received"})
			return
		}

		savedFile, err := saveManifestUploadedFile(ctx, boxID, fileID, file)
		if err != nil {
			markManifestFileStatus(boxID, fileID, fileStatusFailed)
			ctx.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
			return
		}

		ctx.JSON(http.StatusOK, gin.H{
			"box_id":  boxID,
			"box_url": "/box/" + boxID,
			"file":    savedFile,
		})
	})

	router.POST("/box/:id/files/:file_id/status", func(ctx *gin.Context) {
		boxID := ctx.Param("id")
		fileID := ctx.Param("file_id")
		if !validBoxID(boxID) {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "Invalid box id"})
			return
		}

		var request updateFileStatusRequest
		if err := ctx.ShouldBindJSON(&request); err != nil {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "Invalid status payload"})
			return
		}

		file, err := markManifestFileStatus(boxID, fileID, request.Status)
		if err != nil {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
			return
		}

		ctx.JSON(http.StatusOK, gin.H{"file": file})
	})

	router.POST("/box/:id/upload", func(ctx *gin.Context) {
		boxID := ctx.Param("id")
		if !validBoxID(boxID) {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "Invalid box id"})
			return
		}

		file, err := ctx.FormFile("file")
		if err != nil {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "No file received"})
			return
		}

		savedFile, err := saveUploadedFile(ctx, boxID, file)
		if err != nil {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
			return
		}

		ctx.JSON(http.StatusOK, gin.H{
			"box_id":  boxID,
			"box_url": "/box/" + boxID,
			"file":    savedFile,
		})
	})

	router.POST("/upload", func(ctx *gin.Context) {
		form, err := ctx.MultipartForm()
		if err != nil {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "No files received"})
			return
		}

		files := form.File["files"]
		if len(files) == 0 {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "No files received"})
			return
		}

		boxID, err := newBoxID()
		if err != nil {
			ctx.JSON(http.StatusInternalServerError, gin.H{"error": "Could not create upload box"})
			return
		}

		if err := os.MkdirAll(boxPath(boxID), 0755); err != nil {
			ctx.JSON(http.StatusInternalServerError, gin.H{"error": "Could not prepare upload box"})
			return
		}

		savedFiles := make([]gin.H, 0, len(files))

		for _, file := range files {
			savedFile, err := saveUploadedFile(ctx, boxID, file)
			if err != nil {
				ctx.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
				return
			}

			savedFiles = append(savedFiles, savedFile)
		}

		ctx.JSON(http.StatusOK, gin.H{
			"box_id":  boxID,
			"box_url": "/box/" + boxID,
			"files":   savedFiles,
		})
	})

	compressed := router.Group("/", gzip.Gzip(gzip.DefaultCompression))
	compressed.Static("/static", "./static")

	return router.Run(addr)
}

func newBoxID() (string, error) {
	bytes := make([]byte, 16)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}

	return hex.EncodeToString(bytes), nil
}

func newFileID() (string, error) {
	bytes := make([]byte, 8)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}

	return hex.EncodeToString(bytes), nil
}

func validBoxID(boxID string) bool {
	if len(boxID) != 32 {
		return false
	}

	for _, character := range boxID {
		if !strings.ContainsRune("0123456789abcdef", character) {
			return false
		}
	}

	return true
}

func boxPath(boxID string) string {
	return filepath.Join(uploadRoot, boxID)
}

func manifestPath(boxID string) string {
	return filepath.Join(boxPath(boxID), boxManifestFile)
}

func listBoxFiles(boxID string) ([]boxFile, error) {
	if manifest, err := readBoxManifest(boxID); err == nil && len(manifest.Files) > 0 {
		files := make([]boxFile, 0, len(manifest.Files))
		for _, file := range manifest.Files {
			files = append(files, decorateBoxFile(boxID, file))
		}

		return files, nil
	}

	entries, err := os.ReadDir(boxPath(boxID))
	if err != nil {
		return nil, err
	}

	files := make([]boxFile, 0, len(entries))
	for _, entry := range entries {
		if entry.IsDir() || entry.Name() == boxManifestFile {
			continue
		}

		info, err := entry.Info()
		if err != nil {
			return nil, err
		}

		name := entry.Name()
		mimeType := mimeTypeForFile(filepath.Join(boxPath(boxID), name), name)
		files = append(files, decorateBoxFile(boxID, boxFile{
			ID:       name,
			Name:     name,
			Size:     info.Size(),
			MimeType: mimeType,
			Status:   fileStatusReady,
		}))
	}

	return files, nil
}

func createBoxManifest(boxID string, requests []createBoxFileRequest) ([]boxFile, error) {
	usedNames := make(map[string]int, len(requests))
	files := make([]boxFile, 0, len(requests))

	for _, request := range requests {
		filename, ok := safeFilename(request.Name)
		if !ok {
			return nil, fmt.Errorf("Invalid filename")
		}

		filename = uniqueManifestFilename(filename, usedNames)
		fileID, err := newFileID()
		if err != nil {
			return nil, fmt.Errorf("Could not create file id")
		}

		mimeType := mime.TypeByExtension(strings.ToLower(filepath.Ext(filename)))
		if mimeType == "" {
			mimeType = "application/octet-stream"
		}

		files = append(files, boxFile{
			ID:       fileID,
			Name:     filename,
			Size:     request.Size,
			MimeType: mimeType,
			Status:   fileStatusWait,
		})
	}

	manifest := boxManifest{Files: files}
	if err := writeBoxManifest(boxID, manifest); err != nil {
		return nil, err
	}

	decoratedFiles := make([]boxFile, 0, len(files))
	for _, file := range files {
		decoratedFiles = append(decoratedFiles, decorateBoxFile(boxID, file))
	}

	return decoratedFiles, nil
}

func readBoxManifest(boxID string) (boxManifest, error) {
	boxManifestMu.Lock()
	defer boxManifestMu.Unlock()

	return readBoxManifestUnlocked(boxID)
}

func readBoxManifestUnlocked(boxID string) (boxManifest, error) {
	var manifest boxManifest
	data, err := os.ReadFile(manifestPath(boxID))
	if err != nil {
		return manifest, err
	}

	if err := json.Unmarshal(data, &manifest); err != nil {
		return manifest, err
	}

	return manifest, nil
}

func writeBoxManifest(boxID string, manifest boxManifest) error {
	boxManifestMu.Lock()
	defer boxManifestMu.Unlock()

	return writeBoxManifestUnlocked(boxID, manifest)
}

func writeBoxManifestUnlocked(boxID string, manifest boxManifest) error {
	data, err := json.MarshalIndent(manifest, "", "  ")
	if err != nil {
		return err
	}

	return os.WriteFile(manifestPath(boxID), data, 0644)
}

func markManifestFileStatus(boxID string, fileID string, status string) (boxFile, error) {
	if status != fileStatusWait && status != fileStatusWork && status != fileStatusReady && status != fileStatusFailed {
		return boxFile{}, fmt.Errorf("Invalid file status")
	}

	boxManifestMu.Lock()
	defer boxManifestMu.Unlock()

	manifest, err := readBoxManifestUnlocked(boxID)
	if err != nil {
		return boxFile{}, err
	}

	for index, file := range manifest.Files {
		if file.ID != fileID {
			continue
		}

		manifest.Files[index].Status = status
		if err := writeBoxManifestUnlocked(boxID, manifest); err != nil {
			return boxFile{}, err
		}

		return decorateBoxFile(boxID, manifest.Files[index]), nil
	}

	return boxFile{}, fmt.Errorf("File not found")
}

func decorateBoxFile(boxID string, file boxFile) boxFile {
	if file.MimeType == "" {
		file.MimeType = mimeTypeForFile(filepath.Join(boxPath(boxID), file.Name), file.Name)
	}

	if file.SizeLabel == "" {
		file.SizeLabel = formatBytes(file.Size)
	}

	file.IconPath = iconForMimeType(file.MimeType, file.Name)
	file.DownloadPath = "/box/" + boxID + "/files/" + url.PathEscape(file.Name)
	file.UploadPath = "/box/" + boxID + "/files/" + url.PathEscape(file.ID) + "/upload"
	file.IsComplete = file.Status == fileStatusReady

	switch file.Status {
	case fileStatusReady:
		file.StatusLabel = "Ready"
		file.Title = "Download " + file.Name
	case fileStatusFailed:
		file.StatusLabel = "Failed"
		file.Title = "Failed to upload"
	case fileStatusWork:
		file.StatusLabel = "Loading"
		file.Title = "Loading"
	default:
		file.Status = fileStatusWait
		file.StatusLabel = "Waiting"
		file.Title = "Loading"
	}

	return file
}

func addFileToZip(zipWriter *zip.Writer, boxID string, filename string) error {
	path := filepath.Join(boxPath(boxID), filename)
	source, err := os.Open(path)
	if err != nil {
		return err
	}
	defer source.Close()

	destination, err := zipWriter.Create(filename)
	if err != nil {
		return err
	}

	_, err = io.Copy(destination, source)
	return err
}

func saveManifestUploadedFile(ctx *gin.Context, boxID string, fileID string, file *multipart.FileHeader) (boxFile, error) {
	boxManifestMu.Lock()
	defer boxManifestMu.Unlock()

	manifest, err := readBoxManifestUnlocked(boxID)
	if err != nil {
		return boxFile{}, err
	}

	fileIndex := -1
	for index, manifestFile := range manifest.Files {
		if manifestFile.ID == fileID {
			fileIndex = index
			break
		}
	}

	if fileIndex < 0 {
		return boxFile{}, fmt.Errorf("File not found")
	}

	filename := manifest.Files[fileIndex].Name
	if err := os.MkdirAll(boxPath(boxID), 0755); err != nil {
		return boxFile{}, fmt.Errorf("Could not prepare upload box")
	}

	destination := filepath.Join(boxPath(boxID), filename)
	if err := ctx.SaveUploadedFile(file, destination); err != nil {
		manifest.Files[fileIndex].Status = fileStatusFailed
		writeBoxManifestUnlocked(boxID, manifest)
		return boxFile{}, fmt.Errorf("Could not save uploaded file")
	}

	manifest.Files[fileIndex].Size = file.Size
	manifest.Files[fileIndex].MimeType = mimeTypeForFile(destination, filename)
	manifest.Files[fileIndex].Status = fileStatusReady
	if err := writeBoxManifestUnlocked(boxID, manifest); err != nil {
		return boxFile{}, err
	}

	return decorateBoxFile(boxID, manifest.Files[fileIndex]), nil
}

func saveUploadedFile(ctx *gin.Context, boxID string, file *multipart.FileHeader) (gin.H, error) {
	filename, ok := safeFilename(file.Filename)
	if !ok {
		return nil, fmt.Errorf("Invalid filename")
	}

	boxPath := boxPath(boxID)
	if err := os.MkdirAll(boxPath, 0755); err != nil {
		return nil, fmt.Errorf("Could not prepare upload box")
	}

	filename = uniqueFilename(boxPath, filename)
	destination := filepath.Join(boxPath, filename)
	if err := ctx.SaveUploadedFile(file, destination); err != nil {
		return nil, fmt.Errorf("Could not save uploaded file")
	}

	return gin.H{
		"name": filename,
		"size": file.Size,
	}, nil
}

func safeFilename(name string) (string, bool) {
	filename := filepath.Base(name)
	filename = strings.TrimSpace(filename)
	return filename, filename != "" && filename != "." && filename != string(filepath.Separator)
}

func uniqueFilename(directory string, filename string) string {
	if _, err := os.Stat(filepath.Join(directory, filename)); os.IsNotExist(err) {
		return filename
	}

	extension := filepath.Ext(filename)
	base := strings.TrimSuffix(filename, extension)
	for count := 2; ; count++ {
		candidate := fmt.Sprintf("%s-%d%s", base, count, extension)
		if _, err := os.Stat(filepath.Join(directory, candidate)); os.IsNotExist(err) {
			return candidate
		}
	}
}

func uniqueManifestFilename(filename string, usedNames map[string]int) string {
	count := usedNames[filename]
	usedNames[filename] = count + 1

	if count == 0 {
		return filename
	}

	extension := filepath.Ext(filename)
	base := strings.TrimSuffix(filename, extension)
	return fmt.Sprintf("%s-%d%s", base, count+1, extension)
}

func mimeTypeForFile(path string, filename string) string {
	if mimeType := mime.TypeByExtension(strings.ToLower(filepath.Ext(filename))); mimeType != "" {
		return mimeType
	}

	file, err := os.Open(path)
	if err != nil {
		return "application/octet-stream"
	}
	defer file.Close()

	buffer := make([]byte, 512)
	bytesRead, err := file.Read(buffer)
	if err != nil && err != io.EOF {
		return "application/octet-stream"
	}

	return http.DetectContentType(buffer[:bytesRead])
}

func iconForMimeType(mimeType string, filename string) string {
	extension := strings.ToLower(filepath.Ext(filename))

	switch {
	case extension == ".exe":
		return "/static/img/icons/Program Files Icons - PNG/MSONSEXT.DLL_14_6-0.png"
	case strings.HasPrefix(mimeType, "image/"):
		return "/static/img/sprites/bitmap.png"
	case strings.HasPrefix(mimeType, "video/"):
		return "/static/img/icons/netshow_notransm-1.png"
	case strings.HasPrefix(mimeType, "audio/"):
		return "/static/img/icons/netshow_notransm-1.png"
	case strings.HasPrefix(mimeType, "text/") || extension == ".md":
		return "/static/img/sprites/notepad_file-1.png"
	case strings.Contains(mimeType, "zip") || strings.Contains(mimeType, "compressed") || extension == ".rar" || extension == ".7z" || extension == ".tar" || extension == ".gz":
		return "/static/img/icons/Windows Icons - PNG/zipfldr.dll_14_101-0.png"
	case extension == ".ttf" || extension == ".otf" || extension == ".woff" || extension == ".woff2":
		return "/static/img/sprites/font.png"
	case extension == ".pdf":
		return "/static/img/sprites/journal.png"
	case extension == ".html" || extension == ".css" || extension == ".js":
		return "/static/img/sprites/frame_web-0.png"
	default:
		return "/static/img/icons/Windows Icons - PNG/ole2.dll_14_DEFICON.png"
	}
}

func formatBytes(bytes int64) string {
	units := []string{"B", "KB", "MB", "GB"}
	size := float64(bytes)
	unitIndex := 0

	for size >= 1024 && unitIndex < len(units)-1 {
		size /= 1024
		unitIndex++
	}

	if unitIndex == 0 {
		return fmt.Sprintf("%d %s", bytes, units[unitIndex])
	}

	return fmt.Sprintf("%.1f %s", size, units[unitIndex])
}