warpbox-dev/backend/libs/services/upload.go

package services

import (
	"archive/zip"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"mime/multipart"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"go.etcd.io/bbolt"
	"warpbox.dev/backend/libs/helpers"
)

var boxesBucket = []byte("boxes")

type UploadService struct {
	maxUploadSize int64
	baseURL       string
	dataDir       string
	filesDir      string
	db            *bbolt.DB
	logger        *slog.Logger
}

type UploadOptions struct {
	MaxDays           int
	MaxDownloads      int
	Password          string
	ObfuscateMetadata bool
	OwnerID           string
	CollectionID      string
	SkipSizeLimit     bool
}

type Box struct {
	ID              string    `json:"id"`
	OwnerID         string    `json:"ownerId,omitempty"`
	CollectionID    string    `json:"collectionId,omitempty"`
	Title           string    `json:"title,omitempty"`
	CreatedAt       time.Time `json:"createdAt"`
	ExpiresAt       time.Time `json:"expiresAt"`
	MaxDownloads    int       `json:"maxDownloads"`
	DownloadCount   int       `json:"downloadCount"`
	PasswordSalt    string    `json:"passwordSalt,omitempty"`
	PasswordHash    string    `json:"passwordHash,omitempty"`
	DeleteTokenHash string    `json:"deleteTokenHash,omitempty"`
	Obfuscate       bool      `json:"obfuscate"`
	Files           []File    `json:"files"`
}

type File struct {
	ID          string    `json:"id"`
	Name        string    `json:"name"`
	StoredName  string    `json:"storedName"`
	Size        int64     `json:"size"`
	ContentType string    `json:"contentType"`
	PreviewKind string    `json:"previewKind"`
	Thumbnail   string    `json:"thumbnail,omitempty"`
	UploadedAt  time.Time `json:"uploadedAt"`
}

type UploadResult struct {
	BoxID     string       `json:"boxId"`
	BoxURL    string       `json:"boxUrl"`
	ZipURL    string       `json:"zipUrl"`
	ManageURL string       `json:"manageUrl"`
	DeleteURL string       `json:"deleteUrl"`
	ExpiresAt string       `json:"expiresAt"`
	Files     []ResultFile `json:"files"`
}

type ResultFile struct {
	ID   string `json:"id"`
	Name string `json:"name"`
	Size string `json:"size"`
	URL  string `json:"url"`
}

type AdminStats struct {
	TotalBoxes     int
	TotalFiles     int
	TotalSize      int64
	UploadsLast24H int
	ExpiredBoxes   int
	ProtectedBoxes int
	TotalDownloads int
	TotalSizeLabel string
}

type AdminBox struct {
	ID             string
	OwnerID        string
	CreatedAt      time.Time
	ExpiresAt      time.Time
	FileCount      int
	TotalSize      int64
	TotalSizeLabel string
	DownloadCount  int
	MaxDownloads   int
	Protected      bool
	Expired        bool
}

type UserBox struct {
	Box            Box
	CollectionName string
	TotalSizeLabel string
}

func NewUploadService(maxUploadSize int64, dataDir, baseURL string, logger *slog.Logger) (*UploadService, error) {
	filesDir := filepath.Join(dataDir, "files")
	dbDir := filepath.Join(dataDir, "db")
	if err := os.MkdirAll(filesDir, 0o755); err != nil {
		return nil, err
	}
	if err := os.MkdirAll(dbDir, 0o755); err != nil {
		return nil, err
	}

	db, err := bbolt.Open(filepath.Join(dbDir, "warpbox.bbolt"), 0o600, &bbolt.Options{Timeout: time.Second})
	if err != nil {
		return nil, err
	}

	if err := db.Update(func(tx *bbolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists(boxesBucket)
		return err
	}); err != nil {
		db.Close()
		return nil, err
	}

	return &UploadService{
		maxUploadSize: maxUploadSize,
		baseURL:       strings.TrimRight(baseURL, "/"),
		dataDir:       dataDir,
		filesDir:      filesDir,
		db:            db,
		logger:        logger,
	}, nil
}

func (s *UploadService) Close() error {
	return s.db.Close()
}

func (s *UploadService) DB() *bbolt.DB {
	return s.db
}

func (s *UploadService) MaxUploadSize() int64 {
	return s.maxUploadSize
}

func (s *UploadService) MaxUploadSizeLabel() string {
	return helpers.FormatBytes(s.maxUploadSize)
}

func (s *UploadService) ValidateSize(size int64) error {
	if size > s.maxUploadSize {
		return fmt.Errorf("file exceeds max upload size of %s", s.MaxUploadSizeLabel())
	}
	return nil
}

func (s *UploadService) CreateBox(files []*multipart.FileHeader, opts UploadOptions) (UploadResult, error) {
	if len(files) == 0 {
		return UploadResult{}, fmt.Errorf("no files were uploaded")
	}
	if opts.MaxDays <= 0 {
		opts.MaxDays = 7
	}

	box := Box{
		ID:           randomID(10),
		OwnerID:      strings.TrimSpace(opts.OwnerID),
		CollectionID: strings.TrimSpace(opts.CollectionID),
		CreatedAt:    time.Now().UTC(),
		ExpiresAt:    time.Now().UTC().Add(time.Duration(opts.MaxDays) * 24 * time.Hour),
		MaxDownloads: opts.MaxDownloads,
		Obfuscate:    opts.ObfuscateMetadata && strings.TrimSpace(opts.Password) != "",
		Files:        make([]File, 0, len(files)),
	}
	deleteToken := randomID(32)
	box.DeleteTokenHash = deleteTokenHash(box.ID, deleteToken)
	if strings.TrimSpace(opts.Password) != "" {
		salt, hash := hashPassword(opts.Password)
		box.PasswordSalt = salt
		box.PasswordHash = hash
	}

	boxDir := filepath.Join(s.filesDir, box.ID)
	if err := os.MkdirAll(boxDir, 0o755); err != nil {
		return UploadResult{}, err
	}

	for _, header := range files {
		if !opts.SkipSizeLimit {
			if err := s.ValidateSize(header.Size); err != nil {
				return UploadResult{}, err
			}
		}

		maxSize := s.maxUploadSize
		if opts.SkipSizeLimit {
			maxSize = 0
		}

		file, err := header.Open()
		if err != nil {
			return UploadResult{}, err
		}

		fileID := randomID(8)
		storedName := "@each@" + fileID + strings.ToLower(filepath.Ext(header.Filename))
		storedPath := filepath.Join(boxDir, storedName)
		contentType := header.Header.Get("Content-Type")
		if contentType == "" {
			contentType = "application/octet-stream"
		}

		if err := writeUploadedFile(storedPath, file, maxSize); err != nil {
			file.Close()
			return UploadResult{}, err
		}
		file.Close()

		box.Files = append(box.Files, File{
			ID:          fileID,
			Name:        filepath.Base(header.Filename),
			StoredName:  storedName,
			Size:        header.Size,
			ContentType: contentType,
			PreviewKind: previewKind(contentType),
			UploadedAt:  time.Now().UTC(),
		})
	}

	if err := s.SaveBox(box); err != nil {
		return UploadResult{}, err
	}

	s.logger.Info("upload complete",
		"source", "user-upload",
		"severity", "user_activity",
		"code", 2001,
		"box_id", box.ID,
		"file_count", len(box.Files),
	)

	return s.resultForBox(box, deleteToken), nil
}

func (s *UploadService) GetBox(id string) (Box, error) {
	var box Box
	err := s.db.View(func(tx *bbolt.Tx) error {
		data := tx.Bucket(boxesBucket).Get([]byte(id))
		if data == nil {
			return os.ErrNotExist
		}
		return json.Unmarshal(data, &box)
	})
	if err != nil {
		return Box{}, err
	}
	return box, nil
}

func (s *UploadService) ListBoxes(limit int) ([]Box, error) {
	boxes := make([]Box, 0)
	err := s.db.View(func(tx *bbolt.Tx) error {
		cursor := tx.Bucket(boxesBucket).Cursor()
		for key, value := cursor.Last(); key != nil; key, value = cursor.Prev() {
			var box Box
			if err := json.Unmarshal(value, &box); err != nil {
				return err
			}
			boxes = append(boxes, box)
			if limit > 0 && len(boxes) >= limit {
				break
			}
		}
		return nil
	})
	return boxes, err
}

func (s *UploadService) AdminStats() (AdminStats, error) {
	boxes, err := s.ListBoxes(0)
	if err != nil {
		return AdminStats{}, err
	}

	var stats AdminStats
	cutoff := time.Now().UTC().Add(-24 * time.Hour)
	now := time.Now().UTC()
	for _, box := range boxes {
		stats.TotalBoxes++
		stats.TotalDownloads += box.DownloadCount
		if box.CreatedAt.After(cutoff) {
			stats.UploadsLast24H++
		}
		if box.ExpiresAt.Before(now) {
			stats.ExpiredBoxes++
		}
		if s.IsProtected(box) {
			stats.ProtectedBoxes++
		}
		for _, file := range box.Files {
			stats.TotalFiles++
			stats.TotalSize += file.Size
		}
	}
	stats.TotalSizeLabel = helpers.FormatBytes(stats.TotalSize)
	return stats, nil
}

func (s *UploadService) AdminBoxes(limit int) ([]AdminBox, error) {
	boxes, err := s.ListBoxes(limit)
	if err != nil {
		return nil, err
	}

	now := time.Now().UTC()
	rows := make([]AdminBox, 0, len(boxes))
	for _, box := range boxes {
		var size int64
		for _, file := range box.Files {
			size += file.Size
		}
		rows = append(rows, AdminBox{
			ID:             box.ID,
			OwnerID:        box.OwnerID,
			CreatedAt:      box.CreatedAt,
			ExpiresAt:      box.ExpiresAt,
			FileCount:      len(box.Files),
			TotalSize:      size,
			TotalSizeLabel: helpers.FormatBytes(size),
			DownloadCount:  box.DownloadCount,
			MaxDownloads:   box.MaxDownloads,
			Protected:      s.IsProtected(box),
			Expired:        box.ExpiresAt.Before(now),
		})
	}
	return rows, nil
}

func (s *UploadService) UserBoxes(userID string, collectionNames map[string]string) ([]UserBox, error) {
	boxes, err := s.ListBoxes(0)
	if err != nil {
		return nil, err
	}

	rows := make([]UserBox, 0)
	for _, box := range boxes {
		if box.OwnerID != userID {
			continue
		}
		var size int64
		for _, file := range box.Files {
			size += file.Size
		}
		rows = append(rows, UserBox{
			Box:            box,
			CollectionName: collectionNames[box.CollectionID],
			TotalSizeLabel: helpers.FormatBytes(size),
		})
	}
	sort.Slice(rows, func(i, j int) bool {
		return rows[i].Box.CreatedAt.After(rows[j].Box.CreatedAt)
	})
	return rows, nil
}

func (s *UploadService) UserStorageUsed(userID string) (int64, error) {
	boxes, err := s.ListBoxes(0)
	if err != nil {
		return 0, err
	}
	var total int64
	for _, box := range boxes {
		if box.OwnerID != userID {
			continue
		}
		for _, file := range box.Files {
			total += file.Size
		}
	}
	return total, nil
}

func (s *UploadService) RenameOwnedBox(boxID, userID, title string) error {
	box, err := s.GetBox(boxID)
	if err != nil {
		return err
	}
	if box.OwnerID != userID {
		return os.ErrPermission
	}
	box.Title = strings.TrimSpace(title)
	return s.SaveBox(box)
}

func (s *UploadService) MoveOwnedBox(boxID, userID, collectionID string) error {
	box, err := s.GetBox(boxID)
	if err != nil {
		return err
	}
	if box.OwnerID != userID {
		return os.ErrPermission
	}
	box.CollectionID = strings.TrimSpace(collectionID)
	return s.SaveBox(box)
}

func (s *UploadService) DeleteOwnedBox(boxID, userID string) error {
	box, err := s.GetBox(boxID)
	if err != nil {
		return err
	}
	if box.OwnerID != userID {
		return os.ErrPermission
	}
	return s.DeleteBoxWithSource(boxID, "user-delete")
}

func (s *UploadService) DeleteBox(boxID string) error {
	return s.DeleteBoxWithSource(boxID, "admin")
}

func (s *UploadService) DeleteBoxWithToken(boxID, token string) error {
	box, err := s.GetBox(boxID)
	if err != nil {
		return err
	}
	if !s.VerifyDeleteToken(box, token) {
		return os.ErrPermission
	}
	return s.DeleteBoxWithSource(boxID, "anonymous-delete")
}

func (s *UploadService) DeleteBoxWithSource(boxID, source string) error {
	if err := s.db.Update(func(tx *bbolt.Tx) error {
		return tx.Bucket(boxesBucket).Delete([]byte(boxID))
	}); err != nil {
		return err
	}
	if err := os.RemoveAll(filepath.Join(s.filesDir, boxID)); err != nil {
		return err
	}
	s.logger.Info("box deleted", "source", source, "severity", "user_activity", "code", 2101, "box_id", boxID)
	return nil
}

func (s *UploadService) FindFile(box Box, fileID string) (File, error) {
	for _, file := range box.Files {
		if file.ID == fileID {
			return file, nil
		}
	}
	return File{}, os.ErrNotExist
}

func (s *UploadService) FilePath(box Box, file File) string {
	return filepath.Join(s.filesDir, box.ID, file.StoredName)
}

func (s *UploadService) ThumbnailPath(box Box, file File) string {
	if file.Thumbnail == "" {
		return ""
	}
	return filepath.Join(s.filesDir, box.ID, file.Thumbnail)
}

func (s *UploadService) BoxMetadataPath(box Box) string {
	return filepath.Join(s.filesDir, box.ID, ".warpbox.box.json")
}

func (s *UploadService) IsProtected(box Box) bool {
	return box.PasswordHash != "" && box.PasswordSalt != ""
}

func (s *UploadService) VerifyPassword(box Box, password string) bool {
	if !s.IsProtected(box) {
		return true
	}
	hash := passwordHash(box.PasswordSalt, password)
	return subtle.ConstantTimeCompare([]byte(hash), []byte(box.PasswordHash)) == 1
}

func (s *UploadService) UnlockToken(box Box) string {
	sum := sha256.Sum256([]byte(box.ID + ":" + box.PasswordHash))
	return hex.EncodeToString(sum[:])
}

func (s *UploadService) VerifyDeleteToken(box Box, token string) bool {
	if box.DeleteTokenHash == "" || strings.TrimSpace(token) == "" {
		return false
	}
	hash := deleteTokenHash(box.ID, token)
	return subtle.ConstantTimeCompare([]byte(hash), []byte(box.DeleteTokenHash)) == 1
}

func (s *UploadService) CanDownload(box Box) error {
	if time.Now().UTC().After(box.ExpiresAt) {
		return fmt.Errorf("box has expired")
	}
	if box.MaxDownloads > 0 && box.DownloadCount >= box.MaxDownloads {
		return fmt.Errorf("download limit reached")
	}
	return nil
}

func (s *UploadService) RecordDownload(boxID string) error {
	return s.db.Update(func(tx *bbolt.Tx) error {
		bucket := tx.Bucket(boxesBucket)
		data := bucket.Get([]byte(boxID))
		if data == nil {
			return os.ErrNotExist
		}

		var box Box
		if err := json.Unmarshal(data, &box); err != nil {
			return err
		}
		box.DownloadCount++

		next, err := json.Marshal(box)
		if err != nil {
			return err
		}
		if err := bucket.Put([]byte(boxID), next); err != nil {
			return err
		}
		return s.writeBoxMetadata(box)
	})
}

func (s *UploadService) WriteZip(w io.Writer, box Box) error {
	archive := zip.NewWriter(w)
	defer archive.Close()

	for _, file := range box.Files {
		path := s.FilePath(box, file)
		source, err := os.Open(path)
		if err != nil {
			return err
		}

		header := &zip.FileHeader{
			Name:     file.Name,
			Method:   zip.Deflate,
			Modified: file.UploadedAt,
		}
		target, err := archive.CreateHeader(header)
		if err != nil {
			source.Close()
			return err
		}

		if _, err := io.Copy(target, source); err != nil {
			source.Close()
			return err
		}
		source.Close()
	}

	return nil
}

func (s *UploadService) SaveBox(box Box) error {
	data, err := json.Marshal(box)
	if err != nil {
		return err
	}

	return s.db.Update(func(tx *bbolt.Tx) error {
		if err := tx.Bucket(boxesBucket).Put([]byte(box.ID), data); err != nil {
			return err
		}
		return s.writeBoxMetadata(box)
	})
}

func (s *UploadService) resultForBox(box Box, deleteToken string) UploadResult {
	files := make([]ResultFile, 0, len(box.Files))
	for _, file := range box.Files {
		files = append(files, ResultFile{
			ID:   file.ID,
			Name: file.Name,
			Size: helpers.FormatBytes(file.Size),
			URL:  fmt.Sprintf("%s/d/%s/f/%s", s.baseURL, box.ID, file.ID),
		})
	}

	result := UploadResult{
		BoxID:     box.ID,
		BoxURL:    fmt.Sprintf("%s/d/%s", s.baseURL, box.ID),
		ZipURL:    fmt.Sprintf("%s/d/%s/zip", s.baseURL, box.ID),
		ExpiresAt: box.ExpiresAt.Format(time.RFC3339),
		Files:     files,
	}
	if deleteToken != "" {
		result.ManageURL = fmt.Sprintf("%s/d/%s/manage/%s", s.baseURL, box.ID, deleteToken)
		result.DeleteURL = fmt.Sprintf("%s/d/%s/manage/%s/delete", s.baseURL, box.ID, deleteToken)
	}
	return result
}

func writeUploadedFile(path string, source multipart.File, maxSize int64) error {
	target, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o644)
	if err != nil {
		return err
	}
	defer target.Close()

	var written int64
	if maxSize <= 0 {
		written, err = io.Copy(target, source)
	} else {
		written, err = io.Copy(target, io.LimitReader(source, maxSize+1))
	}
	if err != nil {
		os.Remove(path)
		return err
	}
	if maxSize > 0 && written > maxSize {
		os.Remove(path)
		return fmt.Errorf("file exceeds max upload size")
	}
	return nil
}

func randomID(byteCount int) string {
	data := make([]byte, byteCount)
	if _, err := rand.Read(data); err != nil {
		return fmt.Sprintf("%d", time.Now().UnixNano())
	}
	return base64.RawURLEncoding.EncodeToString(data)
}

func hashPassword(password string) (string, string) {
	salt := randomID(18)
	return salt, passwordHash(salt, password)
}

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

func deleteTokenHash(boxID, token string) string {
	sum := sha256.Sum256([]byte("warpbox-delete:" + boxID + ":" + token))
	return hex.EncodeToString(sum[:])
}

func previewKind(contentType string) string {
	switch {
	case strings.HasPrefix(contentType, "image/"):
		return "image"
	case strings.HasPrefix(contentType, "video/"):
		return "video"
	case strings.HasPrefix(contentType, "audio/"):
		return "audio"
	default:
		return "file"
	}
}

func (s *UploadService) writeBoxMetadata(box Box) error {
	path := s.BoxMetadataPath(box)
	data, err := json.MarshalIndent(box, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(path, data, 0o600)
}