warpbox/lib/server/downloads.go

package server

import (
	"archive/zip"
	"fmt"
	"io"
	"net/http"
	"os"
	"sync"

	"github.com/gin-gonic/gin"

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

var oneTimeDownloadLocks sync.Map

func (app *App) handleDownloadBox(ctx *gin.Context) {
	boxID := ctx.Param("id")
	if !boxstore.ValidBoxID(boxID) {
		ctx.String(http.StatusBadRequest, "Invalid box id")
		return
	}

	if !app.config.ZipDownloadsEnabled {
		ctx.String(http.StatusForbidden, "Zip downloads are disabled")
		return
	}

	manifest, hasManifest, ok := app.authorizeBoxRequest(ctx, boxID, true)
	if !ok {
		return
	}
	if hasManifest && manifest.OneTimeDownload {
		app.handleOneTimeDownloadBox(ctx, boxID)
		return
	}

	if hasManifest && manifest.DisableZip {
		ctx.String(http.StatusForbidden, "Zip download disabled for this box")
		return
	}

	files, err := boxstore.ListFiles(boxID)
	if err != nil {
		ctx.String(http.StatusNotFound, "Box not found")
		return
	}
	if !app.writeBoxZip(ctx, boxID, files) {
		return
	}
	if hasManifest && app.config.RenewOnDownloadEnabled {
		boxstore.RenewManifest(boxID, manifest.RetentionSecs)
	}
}

func (app *App) handleOneTimeDownloadBox(ctx *gin.Context, boxID string) {
	lock := oneTimeDownloadLock(boxID)
	lock.Lock()
	defer lock.Unlock()
	defer oneTimeDownloadLocks.Delete(boxID)

	manifest, hasManifest, ok := app.authorizeBoxRequest(ctx, boxID, true)
	if !ok {
		return
	}
	if !hasManifest || !manifest.OneTimeDownload || manifest.Consumed {
		ctx.String(http.StatusGone, "Box already consumed")
		return
	}

	files, err := boxstore.ListFiles(boxID)
	if err != nil {
		ctx.String(http.StatusNotFound, "Box not found")
		return
	}
	if !allFilesComplete(files) {
		ctx.String(http.StatusConflict, "Box is not ready yet")
		return
	}

	if app.config.OneTimeDownloadRetryOnFailure {
		app.handleRetryableOneTimeZip(ctx, boxID, manifest, files)
		return
	}

	manifest.Consumed = true
	if err := boxstore.WriteManifest(boxID, manifest); err != nil {
		ctx.String(http.StatusInternalServerError, "Could not mark box as consumed")
		return
	}
	if !app.writeBoxZip(ctx, boxID, files) {
		boxstore.DeleteBox(boxID)
		return
	}
	boxstore.DeleteBox(boxID)
}

func (app *App) writeBoxZip(ctx *gin.Context, boxID string, files []models.BoxFile) bool {
	writeBoxZipHeaders(ctx, boxID)
	if err := writeBoxZipTo(ctx.Writer, boxID, files); err != nil {
		ctx.Status(http.StatusInternalServerError)
		return false
	}
	return true
}

func (app *App) handleRetryableOneTimeZip(ctx *gin.Context, boxID string, manifest models.BoxManifest, files []models.BoxFile) {
	tempZip, err := os.CreateTemp("", "warpbox-"+boxID+"-*.zip")
	if err != nil {
		ctx.String(http.StatusInternalServerError, "Could not prepare ZIP download")
		return
	}
	tempPath := tempZip.Name()
	defer os.Remove(tempPath)

	if err := writeBoxZipTo(tempZip, boxID, files); err != nil {
		tempZip.Close()
		ctx.String(http.StatusInternalServerError, "Could not build ZIP download")
		return
	}
	if _, err := tempZip.Seek(0, 0); err != nil {
		tempZip.Close()
		ctx.String(http.StatusInternalServerError, "Could not read ZIP download")
		return
	}

	writeBoxZipHeaders(ctx, boxID)
	if _, err := io.Copy(ctx.Writer, tempZip); err != nil {
		tempZip.Close()
		return
	}
	if err := tempZip.Close(); err != nil {
		return
	}

	manifest.Consumed = true
	if err := boxstore.WriteManifest(boxID, manifest); err != nil {
		return
	}
	boxstore.DeleteBox(boxID)
}

func writeBoxZipHeaders(ctx *gin.Context, boxID string) {
	ctx.Header("Content-Type", "application/zip")
	ctx.Header("Content-Disposition", fmt.Sprintf(`attachment; filename="warpbox-%s.zip"`, boxID))
}

func writeBoxZipTo(destination io.Writer, boxID string, files []models.BoxFile) error {
	zipWriter := zip.NewWriter(destination)

	for _, file := range files {
		if !file.IsComplete {
			continue
		}
		if err := boxstore.AddFileToZip(zipWriter, boxID, file.Name); err != nil {
			return err
		}
	}

	if err := zipWriter.Close(); err != nil {
		return err
	}
	return nil
}

func oneTimeDownloadLock(boxID string) *sync.Mutex {
	lock, _ := oneTimeDownloadLocks.LoadOrStore(boxID, &sync.Mutex{})
	return lock.(*sync.Mutex)
}

func allFilesComplete(files []models.BoxFile) bool {
	if len(files) == 0 {
		return false
	}

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

	return true
}

func manifestFilesReady(files []models.BoxFile) bool {
	if len(files) == 0 {
		return false
	}
	for _, file := range files {
		if file.Status != models.FileStatusReady {
			return false
		}
	}
	return true
}

func stripOneTimeThumbnailState(files []models.BoxFile) []models.BoxFile {
	stripped := make([]models.BoxFile, 0, len(files))
	for _, file := range files {
		file.ThumbnailPath = nil
		file.ThumbnailURL = ""
		if file.ThumbnailStatus == "" {
			file.ThumbnailStatus = models.ThumbnailStatusUnsupported
		}
		stripped = append(stripped, file)
	}
	return stripped
}

func (app *App) handleDownloadFile(ctx *gin.Context) {
	boxID := ctx.Param("id")
	filename, ok := helpers.SafeFilename(ctx.Param("filename"))
	if !boxstore.ValidBoxID(boxID) || !ok {
		ctx.String(http.StatusBadRequest, "Invalid file")
		return
	}

	manifest, hasManifest, authorized := app.authorizeBoxRequest(ctx, boxID, true)
	if !authorized {
		return
	}
	if hasManifest && manifest.OneTimeDownload {
		ctx.String(http.StatusForbidden, "Individual downloads disabled for this box")
		return
	}

	path, ok := boxstore.SafeBoxFilePath(boxID, filename)
	if !ok {
		ctx.String(http.StatusBadRequest, "Invalid file")
		return
	}

	if _, err := os.Stat(path); err != nil {
		ctx.String(http.StatusNotFound, "File not found")
		return
	}
	if !boxstore.IsSafeRegularBoxFile(boxID, filename) {
		ctx.String(http.StatusBadRequest, "Invalid file")
		return
	}

	ctx.FileAttachment(path, filename)
	if hasManifest && app.config.RenewOnDownloadEnabled {
		boxstore.RenewManifest(boxID, manifest.RetentionSecs)
	}
}

func (app *App) handleDownloadThumbnail(ctx *gin.Context) {
	boxID := ctx.Param("id")
	fileID := ctx.Param("file_id")
	if !boxstore.ValidBoxID(boxID) {
		ctx.String(http.StatusBadRequest, "Invalid box id")
		return
	}

	manifest, hasManifest, authorized := app.authorizeBoxRequest(ctx, boxID, true)
	if !authorized {
		return
	}
	if hasManifest && manifest.OneTimeDownload {
		ctx.String(http.StatusForbidden, "Thumbnails disabled for one-time boxes")
		return
	}

	path, ok := boxstore.ThumbnailFilePath(boxID, fileID)
	if !ok {
		ctx.String(http.StatusBadRequest, "Invalid thumbnail")
		return
	}

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

	ctx.Header("Content-Type", "image/jpeg")
	ctx.File(path)
}