rill/internal/file/file.go

// Package file 提供文件的上传、删除、查看接口与本地存储服务。
package file

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"gorm.io/gorm"

	"rill/internal/config"
	"rill/internal/model"
)

const (
	storageLocal = "local"
	tempPrefix   = ".upload-"
	// viewCacheControl 文件名由内容哈希决定，内容不会变化，可长缓存。
	viewCacheControl = "public, max-age=31536000, immutable"
)

// ErrTooLarge 上传内容超过大小限制。
var ErrTooLarge = errors.New("file too large")

// ErrEmpty 上传内容为空。
var ErrEmpty = errors.New("empty file")

// ErrFileNotFound 文件记录不存在或已禁用。
var ErrFileNotFound = errors.New("file not found")

// ErrFileInUse 文件仍被业务引用，不允许删除。
var ErrFileInUse = errors.New("file is in use")

// Operator 操作人快照，用于写文件操作日志。
type Operator struct {
	ID   *uint
	Name string
	IP   string
}

// Save 保存上传内容并按 sha256 去重：命中已有记录时直接复用，不写日志。
// 返回记录不增加引用计数，业务引用请调用 Acquire。
func Save(ctx context.Context, db *gorm.DB, cfg *config.Config, operator Operator, filename string, src io.Reader) (*model.File, error) {
	root := cfg.Storage.Dir
	if err := os.MkdirAll(root, 0o755); err != nil {
		return nil, fmt.Errorf("创建存储目录失败: %w", err)
	}

	tmp, err := os.CreateTemp(root, tempPrefix+"*")
	if err != nil {
		return nil, fmt.Errorf("创建临时文件失败: %w", err)
	}
	tmpName := tmp.Name()
	defer func() {
		_ = tmp.Close()
		_ = os.Remove(tmpName)
	}()

	hasher := sha256.New()
	limit := cfg.MaxUploadBytes()
	size, err := io.Copy(io.MultiWriter(tmp, hasher), io.LimitReader(src, limit+1))
	if err != nil {
		return nil, fmt.Errorf("写入上传内容失败: %w", err)
	}
	if size > limit {
		return nil, ErrTooLarge
	}
	if size == 0 {
		return nil, ErrEmpty
	}

	hash := hex.EncodeToString(hasher.Sum(nil))
	var existing model.File
	if err := db.WithContext(ctx).Where("hash = ?", hash).First(&existing).Error; err == nil {
		return &existing, nil
	} else if !errors.Is(err, gorm.ErrRecordNotFound) {
		return nil, err
	}

	mimeType, err := detectMimeType(tmpName)
	if err != nil {
		return nil, err
	}
	extension := extensionFor(mimeType)

	relPath := path.Join(hash[:2], hash+extension)
	finalPath := filepath.Join(root, filepath.FromSlash(relPath))
	if err := os.MkdirAll(filepath.Dir(finalPath), 0o755); err != nil {
		return nil, fmt.Errorf("创建存储子目录失败: %w", err)
	}
	if err := tmp.Close(); err != nil {
		return nil, fmt.Errorf("关闭临时文件失败: %w", err)
	}
	if err := os.Rename(tmpName, finalPath); err != nil {
		return nil, fmt.Errorf("保存文件失败: %w", err)
	}

	record := model.File{
		Name:       displayName(filename),
		Path:       relPath,
		Extension:  extension,
		MimeType:   mimeType,
		Size:       size,
		Hash:       hash,
		RefCount:   0,
		UploaderID: operator.ID,
		Storage:    storageLocal,
		Status:     model.FileStatusEnabled,
	}
	err = db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if err := tx.Create(&record).Error; err != nil {
			return err
		}
		return writeOperation(tx, model.FileOperationCreate, record, operator, "", "", record.Path, record.Name)
	})
	if err != nil {
		if errors.Is(err, gorm.ErrDuplicatedKey) {
			if lookupErr := db.WithContext(ctx).Where("hash = ?", hash).First(&existing).Error; lookupErr == nil {
				return &existing, nil
			}
		}
		return nil, err
	}
	return &record, nil
}

// DeleteFile 删除文件：仍被引用时返回 ErrFileInUse；物理删除后保留记录（status=0）并写日志。
func DeleteFile(ctx context.Context, db *gorm.DB, cfg *config.Config, f model.File, operator Operator) error {
	if f.RefCount > 0 {
		return ErrFileInUse
	}
	if err := removeLocal(cfg.Storage.Dir, f.Path); err != nil {
		return err
	}
	return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if err := tx.Model(&model.File{}).Where("id = ?", f.ID).
			Update("status", model.FileStatusDisabled).Error; err != nil {
			return err
		}
		return writeOperation(tx, model.FileOperationDelete, f, operator, f.Path, f.Name, "", "")
	})
}

// Acquire 增加文件引用计数并刷新最后引用时间。
func Acquire(ctx context.Context, tx *gorm.DB, id uint) error {
	now := time.Now()
	result := tx.WithContext(ctx).Model(&model.File{}).
		Where("id = ? AND status = ?", id, model.FileStatusEnabled).
		Updates(map[string]any{"ref_count": gorm.Expr("ref_count + 1"), "last_referenced_at": now})
	if result.Error != nil {
		return result.Error
	}
	if result.RowsAffected == 0 {
		return ErrFileNotFound
	}
	return nil
}

// Release 减少文件引用计数，最低减到 0；不会删除物理文件。
func Release(ctx context.Context, tx *gorm.DB, id uint) error {
	now := time.Now()
	return tx.WithContext(ctx).Model(&model.File{}).
		Where("id = ? AND ref_count > 0", id).
		Updates(map[string]any{"ref_count": gorm.Expr("ref_count - 1"), "last_referenced_at": now}).Error
}

// Open 打开文件记录对应的本地文件。
func Open(root string, f model.File) (*os.File, error) {
	full, err := localPath(root, f.Path)
	if err != nil {
		return nil, err
	}
	handle, err := os.Open(full)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil, ErrFileNotFound
		}
		return nil, err
	}
	return handle, nil
}

// URL 返回文件的公开访问地址。
func URL(prefix string, id uint) string {
	return strings.TrimSuffix(prefix, "/") + "/files/" + strconv.FormatUint(uint64(id), 10)
}

// ParseLocalURL 从本站文件地址中解析文件 ID，非本站地址返回 false。
func ParseLocalURL(prefix, value string) (uint, bool) {
	base := strings.TrimSuffix(prefix, "/") + "/files/"
	if !strings.HasPrefix(value, base) {
		return 0, false
	}
	id, err := strconv.ParseUint(strings.TrimPrefix(value, base), 10, 64)
	if err != nil || id == 0 {
		return 0, false
	}
	return uint(id), true
}

// cspHeader 用户上传内容统一附加的 CSP：内联展示 SVG 时禁用脚本与外部资源。
const cspHeader = "default-src 'none'; style-src 'unsafe-inline'; sandbox"

// CanInline 是否可内联展示。SVG 允许内联，但查看接口会附加 CSP 禁止脚本执行。
func CanInline(mimeType string) bool {
	switch strings.ToLower(strings.TrimSpace(strings.Split(mimeType, ";")[0])) {
	case "image/jpeg", "image/png", "image/gif", "image/webp", "image/bmp", "image/avif", "image/svg+xml",
		"video/mp4", "video/webm", "audio/mpeg", "audio/ogg", "audio/wav", "application/pdf", "text/plain":
		return true
	default:
		return false
	}
}

// IsImageUpload 通过文件头探测上传内容是否为栅格图片，避免仅信任客户端声明的类型。
func IsImageUpload(header *multipart.FileHeader) (bool, error) {
	head, err := readHead(header)
	if err != nil {
		return false, err
	}
	return strings.HasPrefix(http.DetectContentType(head), "image/"), nil
}

// IsSVGUpload 通过文件头探测上传内容是否为 SVG（可内联但受 CSP 限制）。
func IsSVGUpload(header *multipart.FileHeader) (bool, error) {
	head, err := readHead(header)
	if err != nil {
		return false, err
	}
	return isSVGContent(head), nil
}

// readHead 读取上传文件头部，最多 512 字节。
func readHead(header *multipart.FileHeader) ([]byte, error) {
	src, err := header.Open()
	if err != nil {
		return nil, err
	}
	defer src.Close()

	head := make([]byte, 512)
	n, err := src.Read(head)
	if err != nil && !errors.Is(err, io.EOF) {
		return nil, err
	}
	return head[:n], nil
}

// isSVGContent 判断内容是否为 SVG 文档。
func isSVGContent(head []byte) bool {
	lower := strings.ToLower(string(head))
	return strings.Contains(lower, "<svg")
}

// detectMimeType 读取文件头部探测 MIME；探测失败时回退扩展名与通用类型。
func detectMimeType(name string) (string, error) {
	handle, err := os.Open(name)
	if err != nil {
		return "", fmt.Errorf("读取上传内容失败: %w", err)
	}
	defer handle.Close()

	head := make([]byte, 512)
	n, err := handle.Read(head)
	if err != nil && !errors.Is(err, io.EOF) {
		return "", fmt.Errorf("读取上传内容失败: %w", err)
	}
	detected := http.DetectContentType(head[:n])
	if strings.HasPrefix(detected, "image/") {
		return detected, nil
	}
	// DetectContentType 对 XML 文本返回 text/xml，需单独识别 SVG。
	if isSVGContent(head[:n]) {
		return "image/svg+xml", nil
	}
	if detected != "" && detected != "application/octet-stream" {
		return detected, nil
	}
	if byExt := mime.TypeByExtension(strings.ToLower(filepath.Ext(name))); byExt != "" {
		return byExt, nil
	}
	return "application/octet-stream", nil
}

// extensionFor 返回 MIME 对应的存储扩展名。
func extensionFor(mimeType string) string {
	switch strings.ToLower(strings.TrimSpace(strings.Split(mimeType, ";")[0])) {
	case "image/jpeg":
		return ".jpg"
	case "image/png":
		return ".png"
	case "image/gif":
		return ".gif"
	case "image/webp":
		return ".webp"
	case "image/bmp":
		return ".bmp"
	case "image/avif":
		return ".avif"
	case "image/svg+xml":
		return ".svg"
	case "video/mp4":
		return ".mp4"
	case "video/webm":
		return ".webm"
	case "audio/mpeg":
		return ".mp3"
	case "audio/ogg":
		return ".ogg"
	case "audio/wav":
		return ".wav"
	case "application/pdf":
		return ".pdf"
	default:
		return ""
	}
}

// displayName 清洗上传文件名，仅用于展示与下载名。
func displayName(name string) string {
	name = strings.TrimSpace(path.Base(strings.ReplaceAll(name, "\\", "/")))
	if name == "" || name == "." || name == ".." {
		return "file"
	}
	if len(name) > 255 {
		name = name[len(name)-255:]
	}
	return name
}

// localPath 拼接并校验本地存储路径，防止目录穿越。
func localPath(root, rel string) (string, error) {
	if rel == "" {
		return "", ErrFileNotFound
	}
	cleanRoot, err := filepath.Abs(root)
	if err != nil {
		return "", err
	}
	full := filepath.Join(cleanRoot, filepath.FromSlash(rel))
	if full != cleanRoot && !strings.HasPrefix(full, cleanRoot+string(os.PathSeparator)) {
		return "", ErrFileNotFound
	}
	return full, nil
}

func removeLocal(root, rel string) error {
	full, err := localPath(root, rel)
	if err != nil {
		return err
	}
	if err := os.Remove(full); err != nil && !errors.Is(err, os.ErrNotExist) {
		return fmt.Errorf("删除文件失败: %w", err)
	}
	return nil
}

// writeOperation 追加文件操作日志。
func writeOperation(tx *gorm.DB, operation string, f model.File, operator Operator, pathBefore, nameBefore, pathAfter, nameAfter string) error {
	record := model.FileOperation{
		FileID:     f.ID,
		FileName:   f.Name,
		FileHash:   f.Hash,
		Operation:  operation,
		OperatorID: operator.ID,
		Operator:   operator.Name,
		PathBefore: pathBefore,
		NameBefore: nameBefore,
		PathAfter:  pathAfter,
		NameAfter:  nameAfter,
		IP:         operator.IP,
	}
	return tx.Create(&record).Error
}