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

package services

import (
	"archive/zip"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"mime/multipart"
	"os"
	"path/filepath"
	"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
}

type Box struct {
	ID            string    `json:"id"`
	CreatedAt     time.Time `json:"createdAt"`
	ExpiresAt     time.Time `json:"expiresAt"`
	MaxDownloads  int       `json:"maxDownloads"`
	DownloadCount int       `json:"downloadCount"`
	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"`
	UploadedAt  time.Time `json:"uploadedAt"`
}

type UploadResult struct {
	BoxID     string       `json:"boxId"`
	BoxURL    string       `json:"boxUrl"`
	ZipURL    string       `json:"zipUrl"`
	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"`
}

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) 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),
		CreatedAt:    time.Now().UTC(),
		ExpiresAt:    time.Now().UTC().Add(time.Duration(opts.MaxDays) * 24 * time.Hour),
		MaxDownloads: opts.MaxDownloads,
		Files:        make([]File, 0, len(files)),
	}

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

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

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

		fileID := randomID(8)
		storedName := 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, s.maxUploadSize); 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,
			UploadedAt:  time.Now().UTC(),
		})
	}

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

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

	return s.resultForBox(box), 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) 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) 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
		}
		return bucket.Put([]byte(boxID), next)
	})
}

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 {
		return tx.Bucket(boxesBucket).Put([]byte(box.ID), data)
	})
}

func (s *UploadService) resultForBox(box Box) 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),
		})
	}

	return 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,
	}
}

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()

	written, err := io.Copy(target, io.LimitReader(source, maxSize+1))
	if err != nil {
		os.Remove(path)
		return err
	}
	if 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)
}