drive/backend/handlers/tus.go

package handlers

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io"
	"mime"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"git.elijahkuntz.com/Elijah/drive/config"
	"git.elijahkuntz.com/Elijah/drive/database"
	"git.elijahkuntz.com/Elijah/drive/workers"
	"github.com/gofiber/fiber/v2"
	"github.com/cespare/xxhash/v2"
)

// TusHandler implements a simplified tus v1.0.0 protocol for resumable uploads.
// This handles the core Creation and Termination extensions needed for reliable
// large file uploads over unstable connections.
type TusHandler struct {
	DB     *database.DB
	Config *config.Config
	mu     sync.RWMutex
	uploads map[string]*tusUpload
}

type tusUpload struct {
	ID       string
	FilePath string // final destination path (relative)
	FileName string
	FileSize int64
	Offset   int64
	TempPath string // path to the incomplete upload on disk
	MimeType string
}

func NewTusHandler(db *database.DB, cfg *config.Config) *TusHandler {
	h := &TusHandler{
		DB:      db,
		Config:  cfg,
		uploads: make(map[string]*tusUpload),
	}
	go h.cleanupStaleUploads()
	return h
}

func (h *TusHandler) cleanupStaleUploads() {
	ticker := time.NewTicker(1 * time.Hour)
	for range ticker.C {
		tempDir := filepath.Join(h.Config.StorageDir, ".uploads")
		entries, err := os.ReadDir(tempDir)
		if err != nil {
			continue
		}

		cutoff := time.Now().Add(-24 * time.Hour)

		for _, entry := range entries {
			info, err := entry.Info()
			if err != nil {
				continue
			}

			if info.ModTime().Before(cutoff) {
				path := filepath.Join(tempDir, entry.Name())
				os.Remove(path)

				h.mu.Lock()
				delete(h.uploads, entry.Name())
				h.mu.Unlock()
			}
		}
	}
}

// Options returns tus protocol capability headers.
// OPTIONS /api/tus
func (h *TusHandler) Options(c *fiber.Ctx) error {
	c.Set("Tus-Resumable", "1.0.0")
	c.Set("Tus-Version", "1.0.0")
	c.Set("Tus-Extension", "creation,termination")
	c.Set("Tus-Max-Size", "107374182400") // 100 GB
	return c.SendStatus(fiber.StatusNoContent)
}

// Create initializes a new resumable upload.
// POST /api/tus
func (h *TusHandler) Create(c *fiber.Ctx) error {
	uploadLength := c.Get("Upload-Length")
	if uploadLength == "" {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "Upload-Length header is required"})
	}

	fileSize, err := strconv.ParseInt(uploadLength, 10, 64)
	if err != nil {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid Upload-Length"})
	}

	// Parse metadata (contains filename and destination path)
	metadata := parseTusMetadata(c.Get("Upload-Metadata"))
	fileName := metadata["filename"]
	destPath := metadata["path"]

	if fileName == "" {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "filename metadata is required"})
	}
	if destPath == "" {
		destPath = "."
	}

	// Resolve destination directory and ensure unique filename
	destFull, err := resolveSafe(h.Config.StorageDir, destPath)
	if err != nil {
		return c.Status(fiber.StatusForbidden).JSON(fiber.Map{"error": "invalid destination path"})
	}
	
	uniqueFileName := GetUniquePath(destFull, fileName)

	// Generate unique upload ID
	uploadID, err := generateUploadID()
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to generate upload ID"})
	}

	// Create temp file for the incomplete upload
	tempDir := filepath.Join(h.Config.StorageDir, ".uploads")
	os.MkdirAll(tempDir, 0755)
	tempPath := filepath.Join(tempDir, uploadID)

	// Create the empty temp file
	f, err := os.Create(tempPath)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to create upload"})
	}
	f.Close()

	upload := &tusUpload{
		ID:       uploadID,
		FilePath: filepath.ToSlash(filepath.Join(destPath, uniqueFileName)),
		FileName: uniqueFileName,
		FileSize: fileSize,
		Offset:   0,
		TempPath: tempPath,
		MimeType: mime.TypeByExtension(filepath.Ext(uniqueFileName)),
	}

	h.mu.Lock()
	h.uploads[uploadID] = upload
	h.mu.Unlock()

	c.Set("Tus-Resumable", "1.0.0")
	c.Set("Location", fmt.Sprintf("/api/tus/%s", uploadID))
	c.Set("Upload-Offset", "0")

	return c.SendStatus(fiber.StatusCreated)
}

// Head returns the current offset of an upload (for resuming).
// HEAD /api/tus/:id
func (h *TusHandler) Head(c *fiber.Ctx) error {
	uploadID := c.Params("id")

	h.mu.RLock()
	upload, exists := h.uploads[uploadID]
	h.mu.RUnlock()

	if !exists {
		return c.SendStatus(fiber.StatusNotFound)
	}

	c.Set("Tus-Resumable", "1.0.0")
	c.Set("Upload-Offset", strconv.FormatInt(upload.Offset, 10))
	c.Set("Upload-Length", strconv.FormatInt(upload.FileSize, 10))

	return c.SendStatus(fiber.StatusOK)
}

// Patch appends a chunk of data to an existing upload.
// PATCH /api/tus/:id
func (h *TusHandler) Patch(c *fiber.Ctx) error {
	uploadID := c.Params("id")

	h.mu.RLock()
	upload, exists := h.uploads[uploadID]
	h.mu.RUnlock()

	if !exists {
		return c.SendStatus(fiber.StatusNotFound)
	}

	// Validate offset matches
	clientOffset := c.Get("Upload-Offset")
	if clientOffset == "" {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "Upload-Offset header required"})
	}

	offset, err := strconv.ParseInt(clientOffset, 10, 64)
	if err != nil || offset != upload.Offset {
		return c.Status(fiber.StatusConflict).JSON(fiber.Map{
			"error":           "offset mismatch",
			"expected_offset": upload.Offset,
		})
	}

	// Open temp file at the correct offset and write the chunk
	f, err := os.OpenFile(upload.TempPath, os.O_WRONLY, 0644)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to open upload file"})
	}
	defer f.Close()

	if _, err := f.Seek(upload.Offset, io.SeekStart); err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "seek failed"})
	}

	reader := c.Context().RequestBodyStream()
	if reader == nil {
		// Fallback for smaller payloads or if streaming isn't fully active
		body := c.Body()
		n, err := f.Write(body)
		if err != nil {
			return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "write failed"})
		}
		
		h.mu.Lock()
		upload.Offset += int64(n)
		h.mu.Unlock()
	} else {
		written, err := io.Copy(f, reader)
		if err != nil {
			return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "stream read/write failed"})
		}
		h.mu.Lock()
		upload.Offset += written
		h.mu.Unlock()
	}

	c.Set("Tus-Resumable", "1.0.0")
	c.Set("Upload-Offset", strconv.FormatInt(upload.Offset, 10))

	// Check if upload is complete
	if upload.Offset >= upload.FileSize {
		h.finalizeUpload(upload, c.IP())
	}

	return c.SendStatus(fiber.StatusNoContent)
}

// Delete cancels an in-progress upload (Termination extension).
// DELETE /api/tus/:id
func (h *TusHandler) Delete(c *fiber.Ctx) error {
	uploadID := c.Params("id")

	h.mu.Lock()
	upload, exists := h.uploads[uploadID]
	if exists {
		delete(h.uploads, uploadID)
	}
	h.mu.Unlock()

	if !exists {
		return c.SendStatus(fiber.StatusNotFound)
	}

	os.Remove(upload.TempPath)

	return c.SendStatus(fiber.StatusNoContent)
}

// finalizeUpload moves the completed temp file to its destination and computes checksum.
func (h *TusHandler) finalizeUpload(upload *tusUpload, ip string) {
	// Resolve final destination
	destFull := filepath.Join(h.Config.StorageDir, filepath.FromSlash(upload.FilePath))

	// Ensure parent dir exists
	os.MkdirAll(filepath.Dir(destFull), 0755)

	// Move temp file to final location
	if err := os.Rename(upload.TempPath, destFull); err != nil {
		fmt.Printf("Error finalizing upload %s: %v\n", upload.ID, err)
		return
	}

	// Generate XXHash checksum
	checksum := ""
	if f, err := os.Open(destFull); err == nil {
		hasher := xxhash.New()
		io.Copy(hasher, f)
		f.Close()
		checksum = fmt.Sprintf("%016x", hasher.Sum64())
	}

	// Upsert file record
	h.DB.Exec(`
		INSERT INTO files (path, name, size, mime_type, checksum)
		VALUES (?, ?, ?, ?, ?)
		ON CONFLICT(path) DO UPDATE SET
			size = excluded.size,
			mime_type = excluded.mime_type,
			checksum = excluded.checksum,
			updated_at = CURRENT_TIMESTAMP
	`, upload.FilePath, upload.FileName, upload.FileSize, upload.MimeType, checksum)

	h.DB.AddAuditLog("upload_complete", fmt.Sprintf("%s (%d bytes, tus)", upload.FilePath, upload.FileSize), ip)

	// Queue thumbnail generation
	if workers.Instance != nil {
		workers.Instance.Enqueue(workers.ThumbnailJob{
			FilePath: upload.FilePath,
			Checksum: checksum,
			MimeType: upload.MimeType,
		})
	}

	// Remove from active uploads
	h.mu.Lock()
	delete(h.uploads, upload.ID)
	h.mu.Unlock()

	fmt.Printf("✅ Upload finalized: %s (checksum: %s)\n", upload.FilePath, checksum)
}

// --- Helpers ---

func parseTusMetadata(header string) map[string]string {
	result := make(map[string]string)
	if header == "" {
		return result
	}

	pairs := strings.Split(header, ",")
	for _, pair := range pairs {
		pair = strings.TrimSpace(pair)
		parts := strings.SplitN(pair, " ", 2)
		if len(parts) == 2 {
			// tus metadata values are base64 encoded
			decoded, err := decodeBase64(parts[1])
			if err == nil {
				result[parts[0]] = decoded
			}
		} else if len(parts) == 1 {
			result[parts[0]] = ""
		}
	}
	return result
}

func decodeBase64(s string) (string, error) {
	// Try standard base64 first, then URL-safe
	decoded, err := base64.StdEncoding.DecodeString(s)
	if err != nil {
		decoded, err = base64.RawStdEncoding.DecodeString(s)
	}
	return string(decoded), err
}

func generateUploadID() (string, error) {
	bytes := make([]byte, 16)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}
	return hex.EncodeToString(bytes), nil
}