WarpBox/lib/boxstore/store.go

package boxstore

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

	"warpbox/lib/helpers"
	"warpbox/lib/models"
)

const (
	UploadRoot   = "data/uploads"
	manifestFile = ".warpbox.json"
)

var manifestMu sync.Mutex

func NewBoxID() (string, error) {
	return helpers.RandomHexID(16)
}

func ValidBoxID(boxID string) bool {
	return helpers.ValidLowerHexID(boxID, 32)
}

func BoxPath(boxID string) string {
	return filepath.Join(UploadRoot, boxID)
}

func ManifestPath(boxID string) string {
	return filepath.Join(BoxPath(boxID), manifestFile)
}

func SafeBoxFilePath(boxID string, filename string) (string, bool) {
	return helpers.SafeChildPath(BoxPath(boxID), filename)
}

func ListFiles(boxID string) ([]models.BoxFile, error) {
	if manifest, err := reconcileManifest(boxID); err == nil && len(manifest.Files) > 0 {
		files := make([]models.BoxFile, 0, len(manifest.Files))
		for _, file := range manifest.Files {
			files = append(files, DecorateFile(boxID, file))
		}

		return files, nil
	}

	return listCompletedFilesFromDisk(boxID)
}

func CreateManifest(boxID string, requests []models.CreateBoxFileRequest) ([]models.BoxFile, error) {
	usedNames := make(map[string]int, len(requests))
	files := make([]models.BoxFile, 0, len(requests))

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

		filename = helpers.UniqueNameInBatch(filename, usedNames)
		fileID, err := helpers.RandomHexID(8)
		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, models.BoxFile{
			ID:       fileID,
			Name:     filename,
			Size:     request.Size,
			MimeType: mimeType,
			Status:   models.FileStatusWait,
		})
	}

	manifest := models.BoxManifest{Files: files}
	if err := WriteManifest(boxID, manifest); err != nil {
		return nil, err
	}

	decoratedFiles := make([]models.BoxFile, 0, len(files))
	for _, file := range files {
		decoratedFiles = append(decoratedFiles, DecorateFile(boxID, file))
	}

	return decoratedFiles, nil
}

func MarkFileStatus(boxID string, fileID string, status string) (models.BoxFile, error) {
	if status != models.FileStatusWait && status != models.FileStatusWork && status != models.FileStatusReady && status != models.FileStatusFailed {
		return models.BoxFile{}, fmt.Errorf("Invalid file status")
	}

	manifestMu.Lock()
	defer manifestMu.Unlock()

	manifest, err := readManifestUnlocked(boxID)
	if err != nil {
		return models.BoxFile{}, err
	}

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

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

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

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

func ReadManifest(boxID string) (models.BoxManifest, error) {
	manifestMu.Lock()
	defer manifestMu.Unlock()

	return readManifestUnlocked(boxID)
}

func WriteManifest(boxID string, manifest models.BoxManifest) error {
	manifestMu.Lock()
	defer manifestMu.Unlock()

	return writeManifestUnlocked(boxID, manifest)
}

func AddFileToZip(zipWriter *zip.Writer, boxID string, filename string) error {
	source, err := os.Open(filepath.Join(BoxPath(boxID), filename))
	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 SaveManifestUpload(boxID string, fileID string, file *multipart.FileHeader) (models.BoxFile, error) {
	manifestMu.Lock()
	defer manifestMu.Unlock()

	manifest, err := readManifestUnlocked(boxID)
	if err != nil {
		return models.BoxFile{}, err
	}

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

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

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

	destination := filepath.Join(BoxPath(boxID), filename)
	if err := saveMultipartFile(file, destination); err != nil {
		manifest.Files[fileIndex].Status = models.FileStatusFailed
		writeManifestUnlocked(boxID, manifest)
		return models.BoxFile{}, fmt.Errorf("Could not save uploaded file")
	}

	manifest.Files[fileIndex].Size = file.Size
	manifest.Files[fileIndex].MimeType = helpers.MimeTypeForFile(destination, filename)
	manifest.Files[fileIndex].Status = models.FileStatusReady
	if err := writeManifestUnlocked(boxID, manifest); err != nil {
		return models.BoxFile{}, err
	}

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

func SaveUpload(boxID string, file *multipart.FileHeader) (models.BoxFile, error) {
	filename, ok := helpers.SafeFilename(file.Filename)
	if !ok {
		return models.BoxFile{}, fmt.Errorf("Invalid filename")
	}

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

	filename = helpers.UniqueFilename(boxPath, filename)
	destination := filepath.Join(boxPath, filename)
	if err := saveMultipartFile(file, destination); err != nil {
		return models.BoxFile{}, fmt.Errorf("Could not save uploaded file")
	}

	return DecorateFile(boxID, models.BoxFile{
		ID:       filename,
		Name:     filename,
		Size:     file.Size,
		MimeType: helpers.MimeTypeForFile(destination, filename),
		Status:   models.FileStatusReady,
	}), nil
}

func DecorateFile(boxID string, file models.BoxFile) models.BoxFile {
	if file.MimeType == "" {
		file.MimeType = helpers.MimeTypeForFile(filepath.Join(BoxPath(boxID), file.Name), file.Name)
	}

	if file.SizeLabel == "" {
		file.SizeLabel = helpers.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 == models.FileStatusReady

	switch file.Status {
	case models.FileStatusReady:
		file.StatusLabel = "Ready"
		file.Title = "Download " + file.Name
	case models.FileStatusFailed:
		file.StatusLabel = "Failed"
		file.Title = "Failed to upload"
	case models.FileStatusWork:
		file.StatusLabel = "Loading"
		file.Title = "Loading"
	default:
		file.Status = models.FileStatusWait
		file.StatusLabel = "Waiting"
		file.Title = "Loading"
	}

	return file
}

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 reconcileManifest(boxID string) (models.BoxManifest, error) {
	manifestMu.Lock()
	defer manifestMu.Unlock()

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

	changed := false
	for index, file := range manifest.Files {
		path := filepath.Join(BoxPath(boxID), file.Name)
		info, err := os.Stat(path)
		if err != nil {
			continue
		}

		if file.Status == models.FileStatusReady && file.Size == info.Size() {
			continue
		}

		// The manifest is the UI source of truth, but disk wins when an upload
		// was saved and the final status write/response was interrupted.
		manifest.Files[index].Size = info.Size()
		manifest.Files[index].MimeType = helpers.MimeTypeForFile(path, file.Name)
		manifest.Files[index].Status = models.FileStatusReady
		changed = true
	}

	if changed {
		if err := writeManifestUnlocked(boxID, manifest); err != nil {
			return manifest, err
		}
	}

	return manifest, nil
}

func listCompletedFilesFromDisk(boxID string) ([]models.BoxFile, error) {
	entries, err := os.ReadDir(BoxPath(boxID))
	if err != nil {
		return nil, err
	}

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

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

		name := entry.Name()
		files = append(files, DecorateFile(boxID, models.BoxFile{
			ID:       name,
			Name:     name,
			Size:     info.Size(),
			MimeType: helpers.MimeTypeForFile(filepath.Join(BoxPath(boxID), name), name),
			Status:   models.FileStatusReady,
		}))
	}

	return files, nil
}

func readManifestUnlocked(boxID string) (models.BoxManifest, error) {
	var manifest models.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
}

// Manifest writes are serialized because the browser can upload several files
// concurrently into the same box. Without this lock, status updates can race.
func writeManifestUnlocked(boxID string, manifest models.BoxManifest) error {
	data, err := json.MarshalIndent(manifest, "", "  ")
	if err != nil {
		return err
	}

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

func saveMultipartFile(file *multipart.FileHeader, destination string) error {
	source, err := file.Open()
	if err != nil {
		return err
	}
	defer source.Close()

	target, err := os.Create(destination)
	if err != nil {
		return err
	}
	defer target.Close()

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