warpbox-dev/backend/libs/handlers/resumable.go

package handlers

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"warpbox.dev/backend/libs/helpers"
	"warpbox.dev/backend/libs/jobs"
	"warpbox.dev/backend/libs/services"
)

type resumableCreateRequest struct {
	Files             []services.ResumableFileInput `json:"files"`
	MaxDays           int                           `json:"maxDays"`
	ExpiresMinutes    int                           `json:"expiresMinutes"`
	MaxDownloads      int                           `json:"maxDownloads"`
	Password          string                        `json:"password"`
	ObfuscateMetadata bool                          `json:"obfuscateMetadata"`
	CollectionID      string                        `json:"collectionId"`
}

type resumableSessionResponse struct {
	SessionID   string                   `json:"sessionId"`
	ResumeToken string                   `json:"resumeToken,omitempty"`
	ChunkSize   int64                    `json:"chunkSize"`
	Status      string                   `json:"status"`
	BoxID       string                   `json:"boxId,omitempty"`
	ExpiresAt   string                   `json:"expiresAt"`
	Files       []services.ResumableFile `json:"files"`
}

func (a *App) CreateResumableUpload(w http.ResponseWriter, r *http.Request) {
	user, loggedIn, authErr := a.currentUserWithAuthError(r)
	if authErr != nil {
		a.logger.Warn("resumable upload rejected invalid bearer token", withRequestLogAttrs(r, "source", "user-upload", "severity", "warn", "code", 4011)...)
		helpers.WriteJSONError(w, http.StatusUnauthorized, "invalid bearer token")
		return
	}
	isAdminUpload := loggedIn && user.Role == services.UserRoleAdmin
	settings, policy, ok := a.loadUploadPolicyForAPI(w, r, user, loggedIn)
	if !ok {
		return
	}
	if !settings.ResumableUploadsEnabled {
		helpers.WriteJSONError(w, http.StatusForbidden, "resumable uploads are disabled")
		return
	}
	if !loggedIn && !settings.AnonymousUploadsEnabled {
		a.logger.Warn("resumable anonymous upload rejected disabled", withRequestLogAttrs(r, "source", "user-upload", "severity", "warn", "code", 4013)...)
		helpers.WriteJSONError(w, http.StatusForbidden, "anonymous uploads are disabled")
		return
	}
	rateKey := uploadRateKey(r, user, loggedIn)
	if !isAdminUpload && policy.ShortRequests > 0 && !a.rateLimiter.Allow("upload:"+rateKey, policy.ShortRequests, policy.ShortWindow, time.Now().UTC()) {
		a.logger.Warn("resumable upload rate limited", withRequestLogAttrs(r, "source", "user-upload", "severity", "warn", "code", 4291, "user_id", user.ID)...)
		helpers.WriteJSONError(w, http.StatusTooManyRequests, "too many upload requests, please slow down")
		return
	}

	var payload resumableCreateRequest
	if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
		helpers.WriteJSONError(w, http.StatusBadRequest, "upload session request could not be read")
		return
	}
	fileSizes := make([]int64, 0, len(payload.Files))
	var totalBytes int64
	for _, file := range payload.Files {
		fileSizes = append(fileSizes, file.Size)
		totalBytes += file.Size
	}
	if !isAdminUpload {
		if status, message := a.checkUploadPolicyForSizes(r, user, loggedIn, settings, policy, fileSizes, totalBytes); message != "" {
			a.logger.Warn("resumable upload rejected by policy", withRequestLogAttrs(r, "source", "quota", "severity", "warn", "code", status, "user_id", user.ID, "message", message, "bytes", totalBytes, "files", len(payload.Files))...)
			helpers.WriteJSONError(w, status, message)
			return
		}
		if status, message := a.checkBoxCreationPolicy(r, user, loggedIn, policy); message != "" {
			a.logger.Warn("resumable upload rejected by box policy", withRequestLogAttrs(r, "source", "quota", "severity", "warn", "code", status, "user_id", user.ID, "message", message, "bytes", totalBytes, "files", len(payload.Files))...)
			helpers.WriteJSONError(w, status, message)
			return
		}
	}

	opts, err := a.resumableUploadOptions(r, payload, user, loggedIn, isAdminUpload, policy)
	if err != nil {
		helpers.WriteJSONError(w, http.StatusRequestEntityTooLarge, err.Error())
		return
	}
	chunkSize := int64(settings.ResumableChunkSizeMB * 1024 * 1024)
	retention := time.Duration(settings.ResumableRetentionHours) * time.Hour
	session, err := a.uploadService.CreateResumableSession(payload.Files, opts, chunkSize, retention, resumableChunkRoot(settings))
	if err != nil {
		a.logger.Warn("resumable session create failed", withRequestLogAttrs(r, "source", "user-upload", "severity", "warn", "code", 4002, "user_id", user.ID, "error", err.Error())...)
		helpers.WriteJSONError(w, http.StatusBadRequest, err.Error())
		return
	}
	a.logger.Info("resumable upload session created", withRequestLogAttrs(r, "source", "user-upload", "severity", "user_activity", "code", 2002, "user_id", user.ID, "session_id", session.ID, "files", len(session.Files), "bytes", totalBytes, "anonymous", !loggedIn)...)
	helpers.WriteJSON(w, http.StatusCreated, resumableResponse(session))
}

func (a *App) ResumableUploadStatus(w http.ResponseWriter, r *http.Request) {
	session, ok := a.authorizedResumableSession(w, r)
	if !ok {
		return
	}
	helpers.WriteJSON(w, http.StatusOK, resumableResponse(session))
}

func (a *App) AddResumableFiles(w http.ResponseWriter, r *http.Request) {
	session, ok := a.authorizedResumableSession(w, r)
	if !ok {
		return
	}
	user, loggedIn, _ := a.currentUserWithAuthError(r)
	isAdminUpload := loggedIn && user.Role == services.UserRoleAdmin
	settings, policy, ok := a.loadUploadPolicyForAPI(w, r, user, loggedIn)
	if !ok {
		return
	}
	var payload struct {
		Files []services.ResumableFileInput `json:"files"`
	}
	if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
		helpers.WriteJSONError(w, http.StatusBadRequest, "upload files request could not be read")
		return
	}
	fileSizes := make([]int64, 0, len(session.Files)+len(payload.Files))
	var totalBytes int64
	for _, file := range session.Files {
		fileSizes = append(fileSizes, file.Size)
		totalBytes += file.Size
	}
	for _, file := range payload.Files {
		fileSizes = append(fileSizes, file.Size)
		totalBytes += file.Size
	}
	if !isAdminUpload {
		if status, message := a.checkUploadPolicyForSizes(r, user, loggedIn, settings, policy, fileSizes, totalBytes); message != "" {
			helpers.WriteJSONError(w, status, message)
			return
		}
	}
	updated, err := a.uploadService.AddResumableFiles(session.ID, payload.Files)
	if err != nil {
		helpers.WriteJSONError(w, http.StatusBadRequest, err.Error())
		return
	}
	a.logger.Info("resumable upload files added", withRequestLogAttrs(r, "source", "user-upload", "severity", "user_activity", "code", 2006, "session_id", session.ID, "added", len(updated.Files)-len(session.Files), "files", len(updated.Files))...)
	helpers.WriteJSON(w, http.StatusOK, resumableResponse(updated))
}

func (a *App) PutResumableChunk(w http.ResponseWriter, r *http.Request) {
	session, ok := a.authorizedResumableSession(w, r)
	if !ok {
		return
	}
	fileID := r.PathValue("fileID")
	index, err := strconv.Atoi(r.PathValue("index"))
	if err != nil {
		helpers.WriteJSONError(w, http.StatusBadRequest, "chunk index is invalid")
		return
	}
	updated, err := a.uploadService.PutResumableChunk(r.Context(), session.ID, fileID, index, r.Body)
	if err != nil {
		a.logger.Warn("resumable chunk failed", withRequestLogAttrs(r, "source", "user-upload", "severity", "warn", "code", 4003, "session_id", session.ID, "file_id", fileID, "chunk", index, "error", err.Error())...)
		helpers.WriteJSONError(w, http.StatusBadRequest, err.Error())
		return
	}
	a.logger.Info("resumable chunk uploaded", withRequestLogAttrs(r, "source", "user-upload", "severity", "user_activity", "code", 2003, "session_id", session.ID, "file_id", fileID, "chunk", index)...)
	helpers.WriteJSON(w, http.StatusOK, resumableResponse(updated))
}

func (a *App) CompleteResumableUpload(w http.ResponseWriter, r *http.Request) {
	session, ok := a.authorizedResumableSession(w, r)
	if !ok {
		return
	}
	if session.Status == services.ResumableStatusCompleted {
		result, completed, err := a.uploadService.CompleteResumableSession(r.Context(), session.ID)
		if err != nil {
			a.logger.Warn("resumable upload completion replay failed", withRequestLogAttrs(r, "source", "user-upload", "severity", "warn", "code", 4004, "session_id", session.ID, "error", err.Error())...)
			helpers.WriteJSONError(w, http.StatusBadRequest, err.Error())
			return
		}
		a.logger.Info("resumable upload completion replayed", withRequestLogAttrs(r, "source", "user-upload", "severity", "user_activity", "code", 2004, "session_id", completed.ID, "box_id", result.BoxID, "files", len(result.Files))...)
		helpers.WriteJSON(w, http.StatusOK, result)
		return
	}
	if session.Status == services.ResumableStatusProcessing {
		result, err := a.uploadService.FinalizeProcessingResumableSession(r.Context(), session.ID)
		if err != nil {
			a.logger.Warn("resumable upload completion replay failed", withRequestLogAttrs(r, "source", "user-upload", "severity", "warn", "code", 4004, "session_id", session.ID, "error", err.Error())...)
			helpers.WriteJSONError(w, http.StatusBadRequest, err.Error())
			return
		}
		a.logger.Info("resumable upload completion replayed", withRequestLogAttrs(r, "source", "user-upload", "severity", "user_activity", "code", 2004, "session_id", session.ID, "box_id", result.BoxID, "files", len(result.Files))...)
		helpers.WriteJSON(w, http.StatusOK, result)
		return
	}
	user, loggedIn, _ := a.currentUserWithAuthError(r)
	isAdminUpload := loggedIn && user.Role == services.UserRoleAdmin
	settings, policy, ok := a.loadUploadPolicyForAPI(w, r, user, loggedIn)
	if !ok {
		return
	}
	fileSizes := make([]int64, 0, len(session.Files))
	var totalBytes int64
	for _, file := range session.Files {
		fileSizes = append(fileSizes, file.Size)
		totalBytes += file.Size
	}
	if !isAdminUpload {
		if status, message := a.checkUploadPolicyForSizes(r, user, loggedIn, settings, policy, fileSizes, totalBytes); message != "" {
			helpers.WriteJSONError(w, status, message)
			return
		}
		if status, message := a.checkBoxCreationPolicy(r, user, loggedIn, policy); message != "" {
			helpers.WriteJSONError(w, status, message)
			return
		}
		if status, message := a.checkStorageBackendCapacity(session.Options.StorageBackendID, settings, totalBytes); message != "" {
			helpers.WriteJSONError(w, status, message)
			return
		}
	}
	result, completed, err := a.uploadService.CreateProcessingBoxFromResumable(session.ID)
	if err != nil {
		a.logger.Warn("resumable upload complete failed", withRequestLogAttrs(r, "source", "user-upload", "severity", "warn", "code", 4004, "session_id", session.ID, "error", err.Error())...)
		helpers.WriteJSONError(w, http.StatusBadRequest, err.Error())
		return
	}
	if !isAdminUpload {
		if err := a.recordUploadUsage(r, user, loggedIn, totalBytes, 1); err != nil {
			a.logger.Warn("failed to record resumable upload usage", "source", "quota", "severity", "warn", "code", 4404, "error", err.Error())
		}
		if err := a.settingsService.CleanupUsage(time.Now().UTC(), settings.UsageRetentionDays); err != nil {
			a.logger.Warn("failed to cleanup upload usage", "source", "quota", "severity", "warn", "code", 4405, "error", err.Error())
		}
	}
	a.finalizeResumableUploadAsync(completed.ID, result.BoxID)
	a.logger.Info("resumable upload queued for processing", withRequestLogAttrs(r, "source", "user-upload", "severity", "user_activity", "code", 2004, "user_id", user.ID, "session_id", completed.ID, "box_id", result.BoxID, "files", len(result.Files), "bytes", totalBytes, "admin", isAdminUpload, "anonymous", !loggedIn)...)
	helpers.WriteJSON(w, http.StatusCreated, result)
}

func (a *App) CompleteUploadedResumableUpload(w http.ResponseWriter, r *http.Request) {
	session, ok := a.authorizedResumableSession(w, r)
	if !ok {
		return
	}
	user, loggedIn, _ := a.currentUserWithAuthError(r)
	isAdminUpload := loggedIn && user.Role == services.UserRoleAdmin
	settings, policy, ok := a.loadUploadPolicyForAPI(w, r, user, loggedIn)
	if !ok {
		return
	}
	fileSizes := make([]int64, 0, len(session.Files))
	var totalBytes int64
	var completeCount int
	for _, file := range session.Files {
		if len(file.UploadedChunks) != file.ChunkCount {
			continue
		}
		fileSizes = append(fileSizes, file.Size)
		totalBytes += file.Size
		completeCount++
	}
	if completeCount == 0 {
		helpers.WriteJSONError(w, http.StatusBadRequest, "no fully uploaded files to finish")
		return
	}
	if !isAdminUpload {
		if status, message := a.checkUploadPolicyForSizes(r, user, loggedIn, settings, policy, fileSizes, totalBytes); message != "" {
			helpers.WriteJSONError(w, status, message)
			return
		}
		if status, message := a.checkBoxCreationPolicy(r, user, loggedIn, policy); message != "" {
			helpers.WriteJSONError(w, status, message)
			return
		}
		if status, message := a.checkStorageBackendCapacity(session.Options.StorageBackendID, settings, totalBytes); message != "" {
			helpers.WriteJSONError(w, status, message)
			return
		}
	}
	result, completed, err := a.uploadService.CompleteUploadedResumableSession(r.Context(), session.ID)
	if err != nil {
		a.logger.Warn("resumable partial complete failed", withRequestLogAttrs(r, "source", "user-upload", "severity", "warn", "code", 4005, "session_id", session.ID, "error", err.Error())...)
		helpers.WriteJSONError(w, http.StatusBadRequest, err.Error())
		return
	}
	if !isAdminUpload {
		if err := a.recordUploadUsage(r, user, loggedIn, totalBytes, 1); err != nil {
			a.logger.Warn("failed to record partial resumable upload usage", "source", "quota", "severity", "warn", "code", 4406, "error", err.Error())
		}
		if err := a.settingsService.CleanupUsage(time.Now().UTC(), settings.UsageRetentionDays); err != nil {
			a.logger.Warn("failed to cleanup upload usage", "source", "quota", "severity", "warn", "code", 4405, "error", err.Error())
		}
	}
	jobs.GenerateThumbnailsForBoxAsync(a.uploadService, a.logger, result.BoxID)
	a.logger.Info("resumable uploaded files completed", withRequestLogAttrs(r, "source", "user-upload", "severity", "user_activity", "code", 2007, "user_id", user.ID, "session_id", completed.ID, "box_id", result.BoxID, "files", len(result.Files), "bytes", totalBytes, "admin", isAdminUpload, "anonymous", !loggedIn)...)
	helpers.WriteJSON(w, http.StatusCreated, result)
}

func (a *App) finalizeResumableUploadAsync(sessionID, boxID string) {
	go func() {
		a.logger.Info("resumable upload processing started", "source", "user-upload", "severity", "user_activity", "code", 2009, "session_id", sessionID, "box_id", boxID)
		result, err := a.uploadService.FinalizeProcessingResumableSession(context.Background(), sessionID)
		if err != nil {
			a.logger.Warn("resumable upload processing failed", "source", "user-upload", "severity", "warn", "code", 4010, "session_id", sessionID, "box_id", boxID, "error", err.Error())
			return
		}
		jobs.GenerateThumbnailsForBoxAsync(a.uploadService, a.logger, result.BoxID)
		a.logger.Info("resumable upload processing completed", "source", "user-upload", "severity", "user_activity", "code", 2010, "session_id", sessionID, "box_id", result.BoxID, "files", len(result.Files))
	}()
}

func resumableChunkRoot(settings services.UploadPolicySettings) string {
	if settings.ResumableChunkMode != "custom" {
		return ""
	}
	return strings.TrimSpace(settings.ResumableChunkPath)
}

func (a *App) CancelResumableUpload(w http.ResponseWriter, r *http.Request) {
	session, ok := a.authorizedResumableSession(w, r)
	if !ok {
		return
	}
	if err := a.uploadService.CancelResumableSession(session.ID); err != nil {
		helpers.WriteJSONError(w, http.StatusBadRequest, err.Error())
		return
	}
	a.logger.Info("resumable upload cancelled", withRequestLogAttrs(r, "source", "user-upload", "severity", "user_activity", "code", 2005, "session_id", session.ID)...)
	w.WriteHeader(http.StatusNoContent)
}

func (a *App) authorizedResumableSession(w http.ResponseWriter, r *http.Request) (services.ResumableSession, bool) {
	user, loggedIn, authErr := a.currentUserWithAuthError(r)
	if authErr != nil {
		helpers.WriteJSONError(w, http.StatusUnauthorized, "invalid bearer token")
		return services.ResumableSession{}, false
	}
	session, err := a.uploadService.GetResumableSession(r.PathValue("sessionID"))
	if err != nil {
		helpers.WriteJSONError(w, http.StatusNotFound, "upload session not found")
		return services.ResumableSession{}, false
	}
	if !a.uploadService.VerifyResumableToken(session, r.Header.Get("X-Warpbox-Resume-Token")) {
		helpers.WriteJSONError(w, http.StatusUnauthorized, "upload session not found")
		return services.ResumableSession{}, false
	}
	if loggedIn {
		if session.Options.OwnerID != user.ID {
			helpers.WriteJSONError(w, http.StatusForbidden, "upload session not found")
			return services.ResumableSession{}, false
		}
		return session, true
	}
	if session.Options.OwnerID != "" || session.Options.CreatorIP != uploadClientIP(r) {
		helpers.WriteJSONError(w, http.StatusForbidden, "upload session not found")
		return services.ResumableSession{}, false
	}
	return session, true
}

func (a *App) loadUploadPolicyForAPI(w http.ResponseWriter, r *http.Request, user services.User, loggedIn bool) (services.UploadPolicySettings, services.EffectiveUploadPolicy, bool) {
	settings, err := a.settingsService.UploadPolicy()
	if err != nil {
		a.logger.Error("failed to load upload policy", "source", "settings", "severity", "error", "code", 5006, "error", err.Error())
		helpers.WriteJSONError(w, http.StatusInternalServerError, "upload policy could not be loaded")
		return services.UploadPolicySettings{}, services.EffectiveUploadPolicy{}, false
	}
	return settings, a.effectiveUploadPolicy(settings, user, loggedIn), true
}

func (a *App) resumableUploadOptions(r *http.Request, payload resumableCreateRequest, user services.User, loggedIn, isAdminUpload bool, policy services.EffectiveUploadPolicy) (services.UploadOptions, error) {
	var ownerID string
	var collectionID string
	if loggedIn {
		ownerID = user.ID
		collectionID = strings.TrimSpace(payload.CollectionID)
		if !a.authService.CollectionOwnedBy(collectionID, user.ID) {
			return services.UploadOptions{}, fmt.Errorf("collection not found")
		}
	}
	unlimitedExpiry := isAdminUpload || policy.MaxDays < 0
	rawMaxDays := payload.MaxDays
	maxDays := rawMaxDays
	if maxDays <= 0 {
		maxDays = 7
		if policy.MaxDays > 0 && policy.MaxDays < maxDays {
			maxDays = policy.MaxDays
		}
	}
	expiresMinutes := payload.ExpiresMinutes
	if expiresMinutes < 0 || rawMaxDays < 0 {
		if !unlimitedExpiry {
			return services.UploadOptions{}, fmt.Errorf("expiration cannot exceed %d days", policy.MaxDays)
		}
		expiresMinutes = -1
	} else if !unlimitedExpiry {
		if maxDays > policy.MaxDays {
			return services.UploadOptions{}, fmt.Errorf("expiration cannot exceed %d days", policy.MaxDays)
		}
		if expiresMinutes > 0 && expiresMinutes > policy.MaxDays*24*60 {
			return services.UploadOptions{}, fmt.Errorf("expiration cannot exceed %d days", policy.MaxDays)
		}
	}
	return services.UploadOptions{
		MaxDays:           maxDays,
		ExpiresInMinutes:  expiresMinutes,
		MaxDownloads:      payload.MaxDownloads,
		Password:          payload.Password,
		ObfuscateMetadata: payload.ObfuscateMetadata,
		OwnerID:           ownerID,
		CollectionID:      collectionID,
		SkipSizeLimit:     isAdminUpload || policy.MaxUploadMB < 0,
		CreatorIP:         uploadClientIP(r),
		StorageBackendID:  policy.StorageBackendID,
	}, nil
}

func resumableResponse(session services.ResumableSession) resumableSessionResponse {
	return resumableSessionResponse{
		SessionID:   session.ID,
		ResumeToken: session.ResumeToken,
		ChunkSize:   session.ChunkSize,
		Status:      session.Status,
		BoxID:       session.BoxID,
		ExpiresAt:   session.ExpiresAt.Format(time.RFC3339),
		Files:       session.Files,
	}
}