WarpBox/lib/boxstore/store.go

package boxstore

import (
	"archive/zip"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/bcrypt"

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

const (
	manifestFile = ".warpbox.json"

	OneTimeDownloadRetentionKey = "one-time"
)

var (
	uploadRoot = filepath.Join("data", "uploads")
	manifestMu sync.Mutex
)

var retentionOptions = []models.RetentionOption{
	{Key: "10s", Label: "10 seconds", Seconds: 10},
	{Key: "10m", Label: "10 minutes", Seconds: 10 * 60},
	{Key: "1h", Label: "1 hour", Seconds: 60 * 60},
	{Key: "12h", Label: "12 hours", Seconds: 12 * 60 * 60},
	{Key: "24h", Label: "24 hours", Seconds: 24 * 60 * 60},
	{Key: "48h", Label: "48 hours", Seconds: 48 * 60 * 60},
	{Key: OneTimeDownloadRetentionKey, Label: "One time download", Seconds: 0},
}

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

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

func RetentionOptions() []models.RetentionOption {
	options := make([]models.RetentionOption, len(retentionOptions))
	copy(options, retentionOptions)
	return options
}

func DefaultRetentionOption() models.RetentionOption {
	return retentionOptions[0]
}

func SetUploadRoot(path string) {
	if path == "" {
		return
	}
	uploadRoot = filepath.Clean(path)
}

func UploadRoot() string {
	return uploadRoot
}

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

func safeBoxPath(boxID string) (string, bool) {
	if !ValidBoxID(boxID) {
		return "", false
	}
	return helpers.SafeChildPath(uploadRoot, boxID)
}

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

func SafeBoxFilePath(boxID string, filename string) (string, bool) {
	boxPath, ok := safeBoxPath(boxID)
	if !ok {
		return "", false
	}
	return helpers.SafeChildPath(boxPath, filename)
}

func IsSafeRegularBoxFile(boxID string, filename string) bool {
	path, ok := SafeBoxFilePath(boxID, filename)
	if !ok {
		return false
	}
	return ensureRegularFile(path) == nil
}

func DeleteBox(boxID string) error {
	boxPath, ok := safeBoxPath(boxID)
	if !ok {
		return fmt.Errorf("Invalid box id")
	}
	return os.RemoveAll(boxPath)
}

func ListBoxSummaries() ([]models.BoxSummary, error) {
	entries, err := os.ReadDir(uploadRoot)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}

	summaries := make([]models.BoxSummary, 0, len(entries))
	for _, entry := range entries {
		if !entry.IsDir() || !ValidBoxID(entry.Name()) {
			continue
		}

		summary, err := BoxSummary(entry.Name())
		if err != nil {
			continue
		}
		summaries = append(summaries, summary)
	}

	sort.Slice(summaries, func(i int, j int) bool {
		return summaries[i].CreatedAt.After(summaries[j].CreatedAt)
	})
	return summaries, nil
}

func BoxSummary(boxID string) (models.BoxSummary, error) {
	files, err := ListFiles(boxID)
	if err != nil {
		return models.BoxSummary{}, err
	}

	var manifest models.BoxManifest
	hasManifest := false
	if readManifest, err := ReadManifest(boxID); err == nil {
		manifest = readManifest
		hasManifest = true
	}

	totalSize := int64(0)
	for _, file := range files {
		totalSize += file.Size
	}

	summary := models.BoxSummary{
		ID:             boxID,
		FileCount:      len(files),
		TotalSize:      totalSize,
		TotalSizeLabel: helpers.FormatBytes(totalSize),
	}
	if hasManifest {
		summary.CreatedAt = manifest.CreatedAt
		summary.ExpiresAt = manifest.ExpiresAt
		summary.Expired = IsExpired(manifest)
		summary.OneTimeDownload = manifest.OneTimeDownload
		summary.PasswordProtected = IsPasswordProtected(manifest)
	} else if info, err := os.Stat(BoxPath(boxID)); err == nil {
		summary.CreatedAt = info.ModTime().UTC()
	}

	return summary, nil
}

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, request models.CreateBoxRequest) ([]models.BoxFile, error) {
	retention := normalizeRetentionOption(request.RetentionKey)
	usedNames := make(map[string]int, len(request.Files))
	files := make([]models.BoxFile, 0, len(request.Files))

	for _, fileRequest := range request.Files {
		filename, ok := helpers.SafeFilename(fileRequest.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:     fileRequest.Size,
			MimeType: mimeType,
			Status:   models.FileStatusWait,
		})
	}

	now := time.Now().UTC()
	disableZip := false
	if request.AllowZip != nil {
		disableZip = !*request.AllowZip
	}
	oneTimeDownload := retention.Key == OneTimeDownloadRetentionKey
	if oneTimeDownload {
		disableZip = false
	}

	manifest := models.BoxManifest{
		Files:           files,
		CreatedAt:       now,
		RetentionKey:    retention.Key,
		RetentionLabel:  retention.Label,
		RetentionSecs:   retention.Seconds,
		DisableZip:      disableZip,
		OneTimeDownload: oneTimeDownload,
	}

	if password := strings.TrimSpace(request.Password); password != "" {
		authToken, err := helpers.RandomHexID(16)
		if err != nil {
			return nil, fmt.Errorf("Could not secure upload box")
		}
		passwordHash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
		if err != nil {
			return nil, fmt.Errorf("Could not secure upload box")
		}

		manifest.PasswordHash = string(passwordHash)
		manifest.PasswordHashAlg = "bcrypt"
		manifest.AuthToken = authToken
	}

	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 IsExpired(manifest models.BoxManifest) bool {
	return !manifest.ExpiresAt.IsZero() && time.Now().UTC().After(manifest.ExpiresAt)
}

func IsPasswordProtected(manifest models.BoxManifest) bool {
	return manifest.PasswordHash != "" && manifest.AuthToken != ""
}

func VerifyPassword(manifest models.BoxManifest, password string) bool {
	if !IsPasswordProtected(manifest) {
		return true
	}

	expected := manifest.PasswordHash
	if manifest.PasswordHashAlg == "bcrypt" || strings.HasPrefix(expected, "$2") {
		return bcrypt.CompareHashAndPassword([]byte(expected), []byte(password)) == nil
	}

	actual := legacyPasswordHash(manifest.PasswordSalt, password)
	return subtle.ConstantTimeCompare([]byte(expected), []byte(actual)) == 1
}

func VerifyAuthToken(manifest models.BoxManifest, token string) bool {
	if !IsPasswordProtected(manifest) {
		return true
	}

	if token == "" {
		return false
	}

	return subtle.ConstantTimeCompare([]byte(manifest.AuthToken), []byte(token)) == 1
}

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
		startRetentionIfTerminalUnlocked(&manifest)
		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 RenewManifest(boxID string, seconds int64) (models.BoxManifest, error) {
	manifestMu.Lock()
	defer manifestMu.Unlock()

	manifest, err := readManifestUnlocked(boxID)
	if err != nil {
		return manifest, err
	}
	if seconds <= 0 || manifest.OneTimeDownload || manifest.ExpiresAt.IsZero() {
		return manifest, nil
	}
	manifest.ExpiresAt = time.Now().UTC().Add(time.Duration(seconds) * time.Second)
	return manifest, writeManifestUnlocked(boxID, manifest)
}

func AddFileToZip(zipWriter *zip.Writer, boxID string, filename string) error {
	path, ok := SafeBoxFilePath(boxID, filename)
	if !ok {
		return fmt.Errorf("Invalid file")
	}
	if err := ensureRegularFile(path); err != nil {
		return err
	}
	zipName, ok := safeZipEntryName(filename)
	if !ok {
		return fmt.Errorf("Invalid zip entry")
	}

	source, err := os.Open(path)
	if err != nil {
		return err
	}
	defer source.Close()

	destination, err := zipWriter.Create(zipName)
	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
	}
	if IsExpired(manifest) {
		return models.BoxFile{}, fmt.Errorf("Box expired")
	}

	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, ok := SafeBoxFilePath(boxID, filename)
	if !ok {
		return models.BoxFile{}, fmt.Errorf("Invalid filename")
	}
	if err := saveMultipartFile(file, destination); err != nil {
		manifest.Files[fileIndex].Status = models.FileStatusFailed
		startRetentionIfTerminalUnlocked(&manifest)
		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
	startRetentionIfTerminalUnlocked(&manifest)
	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, ok := SafeBoxFilePath(boxID, filename)
	if !ok {
		return models.BoxFile{}, fmt.Errorf("Invalid 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 == "" {
		if path, ok := SafeBoxFilePath(boxID, file.Name); ok {
			file.MimeType = helpers.MimeTypeForFile(path, file.Name)
		}
	}

	if file.SizeLabel == "" {
		file.SizeLabel = helpers.FormatBytes(file.Size)
	}

	file.IconPath = IconForMimeType(file.MimeType, file.Name)
	if file.ThumbnailPath != nil {
		file.ThumbnailURL = *file.ThumbnailPath
	}
	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, ok := SafeBoxFilePath(boxID, file.Name)
		if !ok || ensureRegularFile(path) != nil {
			continue
		}
		info, err := os.Stat(path)
		if err != nil || !info.Mode().IsRegular() {
			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 {
		startRetentionIfTerminalUnlocked(&manifest)
		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 || entry.Type()&os.ModeSymlink != 0 {
			continue
		}

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

		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
}

func normalizeRetentionOption(key string) models.RetentionOption {
	for _, option := range retentionOptions {
		if option.Key == key {
			return option
		}
	}

	return DefaultRetentionOption()
}

func startRetentionIfTerminalUnlocked(manifest *models.BoxManifest) {
	if !manifest.ExpiresAt.IsZero() || len(manifest.Files) == 0 {
		return
	}
	if manifest.OneTimeDownload {
		return
	}

	for _, file := range manifest.Files {
		if file.Status != models.FileStatusReady && file.Status != models.FileStatusFailed {
			return
		}
	}

	seconds := manifest.RetentionSecs
	if seconds <= 0 {
		seconds = normalizeRetentionOption(manifest.RetentionKey).Seconds
	}

	// Retention starts after uploads settle so slow or very large uploads do
	// not expire before users get a real chance to open the box.
	manifest.ExpiresAt = time.Now().UTC().Add(time.Duration(seconds) * time.Second)
}

func legacyPasswordHash(salt string, password string) string {
	sum := sha256.Sum256([]byte(salt + ":" + password))
	return hex.EncodeToString(sum[:])
}

// 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, tempPath, err := createTempSibling(destination)
	if err != nil {
		return err
	}
	committed := false
	defer func() {
		target.Close()
		if !committed {
			os.Remove(tempPath)
		}
	}()

	if _, err := io.Copy(target, source); err != nil {
		return err
	}
	if err := target.Close(); err != nil {
		return err
	}
	if err := os.Rename(tempPath, destination); err != nil {
		return err
	}
	committed = true
	return nil
}

func createTempSibling(destination string) (*os.File, string, error) {
	directory := filepath.Dir(destination)
	if err := os.MkdirAll(directory, 0755); err != nil {
		return nil, "", err
	}

	target, err := os.CreateTemp(directory, ".warpbox-upload-*")
	if err != nil {
		return nil, "", err
	}
	return target, target.Name(), nil
}

func safeZipEntryName(filename string) (string, bool) {
	filename = strings.TrimSpace(filename)
	if filename == "" || filepath.IsAbs(filename) {
		return "", false
	}

	cleaned := filepath.ToSlash(filepath.Clean(filename))
	if cleaned == "." || cleaned == ".." || strings.HasPrefix(cleaned, "../") || strings.HasPrefix(cleaned, "/") {
		return "", false
	}
	return cleaned, true
}

func ensureRegularFile(path string) error {
	info, err := os.Lstat(path)
	if err != nil {
		return err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return fmt.Errorf("Invalid file")
	}
	return nil
}