278 lines
8.4 KiB
Go
278 lines
8.4 KiB
Go
package db
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/pkg/errors"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/clause"
|
|
|
|
"github.com/tamakyi/TamaBox/internal/conf"
|
|
"github.com/tamakyi/TamaBox/internal/dbutil"
|
|
"github.com/tamakyi/TamaBox/internal/storage"
|
|
)
|
|
|
|
var R2DeleteTasks R2DeleteTaskStore
|
|
|
|
var _ R2DeleteTaskStore = (*r2DeleteTasks)(nil)
|
|
|
|
// R2DeleteTask 是 R2 对象删除的持久化 outbox 任务。
|
|
// 数据库事务只负责登记任务，实际对象删除由 worker 在事务外执行。
|
|
type R2DeleteTask struct {
|
|
ID uint `gorm:"primaryKey"`
|
|
Bucket string `gorm:"size:255;not null;uniqueIndex:uk_r2_delete_tasks_bucket_key_hash,priority:1"`
|
|
Key string `gorm:"type:text;not null"`
|
|
KeyHash string `gorm:"size:64;not null;uniqueIndex:uk_r2_delete_tasks_bucket_key_hash,priority:2"`
|
|
Attempts int `gorm:"not null;default:0"`
|
|
NextRetryAt time.Time `gorm:"not null;index:idx_r2_delete_tasks_next_retry_at"`
|
|
LastError string `gorm:"type:text"`
|
|
CreatedAt time.Time
|
|
UpdatedAt time.Time
|
|
}
|
|
|
|
type R2DeleteTaskStore interface {
|
|
Enqueue(ctx context.Context, bucket string, keys []string) error
|
|
EnqueueValues(ctx context.Context, bucket string, values []string) error
|
|
ListDue(ctx context.Context, limit int) ([]*R2DeleteTask, error)
|
|
Delete(id uint) error
|
|
MarkFailed(id uint, nextRetry time.Time, lastError string) error
|
|
// HasReferences 仅供 worker 对当前配置的 bucket 调用。
|
|
// 它只查询数据库，不执行网络请求；查询或旧 URL 解析失败时必须返回错误。
|
|
HasReferences(ctx context.Context, bucket, key string) (bool, error)
|
|
}
|
|
|
|
func NewR2DeleteTaskStore(db *gorm.DB) R2DeleteTaskStore {
|
|
return &r2DeleteTasks{DB: db}
|
|
}
|
|
|
|
type r2DeleteTasks struct {
|
|
*gorm.DB
|
|
}
|
|
|
|
// EnqueueR2DeleteTasks 在调用方事务中登记对象删除任务。
|
|
// 相同 (bucket,key) 重复入队会重置重试状态，使新事务对应的删除请求立即可见。
|
|
func EnqueueR2DeleteTasks(tx *gorm.DB, bucket string, keys []string) error {
|
|
if bucket == "" || len(keys) == 0 {
|
|
return nil
|
|
}
|
|
|
|
now := time.Now()
|
|
tasks := make([]*R2DeleteTask, 0, len(keys))
|
|
seen := make(map[string]struct{}, len(keys))
|
|
for _, key := range keys {
|
|
if key == "" {
|
|
continue
|
|
}
|
|
if _, ok := seen[key]; ok {
|
|
continue
|
|
}
|
|
seen[key] = struct{}{}
|
|
sum := sha256.Sum256([]byte(key))
|
|
tasks = append(tasks, &R2DeleteTask{
|
|
Bucket: bucket,
|
|
Key: key,
|
|
KeyHash: hex.EncodeToString(sum[:]),
|
|
NextRetryAt: now,
|
|
})
|
|
}
|
|
if len(tasks) == 0 {
|
|
return nil
|
|
}
|
|
|
|
if err := tx.Clauses(clause.OnConflict{
|
|
Columns: []clause.Column{{Name: "bucket"}, {Name: "key_hash"}},
|
|
DoUpdates: clause.Assignments(map[string]interface{}{
|
|
"attempts": 0,
|
|
"next_retry_at": now,
|
|
"last_error": "",
|
|
"updated_at": now,
|
|
}),
|
|
}).Create(&tasks).Error; err != nil {
|
|
return errors.Wrap(err, "enqueue R2 delete tasks")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (db *r2DeleteTasks) Enqueue(ctx context.Context, bucket string, keys []string) error {
|
|
return EnqueueR2DeleteTasks(db.WithContext(ctx), bucket, keys)
|
|
}
|
|
|
|
func (db *r2DeleteTasks) EnqueueValues(ctx context.Context, bucket string, values []string) error {
|
|
return enqueueR2DeleteTasksForValues(db.WithContext(ctx), bucket, values)
|
|
}
|
|
|
|
func (db *r2DeleteTasks) ListDue(ctx context.Context, limit int) ([]*R2DeleteTask, error) {
|
|
tasks := make([]*R2DeleteTask, 0)
|
|
if limit <= 0 {
|
|
return tasks, nil
|
|
}
|
|
if err := db.WithContext(ctx).
|
|
Where("next_retry_at <= ?", time.Now()).
|
|
Order("next_retry_at ASC, id ASC").
|
|
Limit(limit).
|
|
Find(&tasks).Error; err != nil {
|
|
return nil, errors.Wrap(err, "list due R2 delete tasks")
|
|
}
|
|
return tasks, nil
|
|
}
|
|
|
|
func (db *r2DeleteTasks) Delete(id uint) error {
|
|
if err := db.DB.Delete(&R2DeleteTask{}, id).Error; err != nil {
|
|
return errors.Wrap(err, "delete R2 delete task")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (db *r2DeleteTasks) MarkFailed(id uint, nextRetry time.Time, lastError string) error {
|
|
if err := db.Model(&R2DeleteTask{}).Where("id = ?", id).Updates(map[string]interface{}{
|
|
"attempts": gorm.Expr("attempts + 1"),
|
|
"next_retry_at": nextRetry,
|
|
"last_error": lastError,
|
|
"updated_at": time.Now(),
|
|
}).Error; err != nil {
|
|
return errors.Wrap(err, "mark R2 delete task failed")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (db *r2DeleteTasks) HasReferences(ctx context.Context, bucket, key string) (bool, error) {
|
|
if bucket == "" || key == "" {
|
|
return false, nil
|
|
}
|
|
defaultKey, owned, err := r2ObjectKeyFromValue(key)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if owned && isDefaultProfileObject(defaultKey) {
|
|
return true, nil
|
|
}
|
|
|
|
quotedKey := dbutil.QuoteIdentifier(db.Dialector.Name(), "key")
|
|
var count int64
|
|
if err := db.WithContext(ctx).Unscoped().Model(&UploadImage{}).
|
|
Where(quotedKey+" = ?", key).Count(&count).Error; err != nil {
|
|
return false, errors.Wrap(err, "check upload image references")
|
|
}
|
|
if count > 0 {
|
|
return true, nil
|
|
}
|
|
|
|
count = 0
|
|
if err := db.WithContext(ctx).Unscoped().Model(&MessageImage{}).
|
|
Where("object_key = ? OR url = ?", key, key).Count(&count).Error; err != nil {
|
|
return false, errors.Wrap(err, "check message image references")
|
|
}
|
|
if count > 0 {
|
|
return true, nil
|
|
}
|
|
|
|
var uploadURLs []string
|
|
if err := db.WithContext(ctx).Unscoped().Model(&UploadImage{}).
|
|
Where(quotedKey+" LIKE ? OR "+quotedKey+" LIKE ?", "%://%", "//%").
|
|
Pluck(quotedKey, &uploadURLs).Error; err != nil {
|
|
return false, errors.Wrap(err, "list legacy upload image URLs")
|
|
}
|
|
referenced, err := hasLegacyCDNReference(uploadURLs, key)
|
|
if err != nil || referenced {
|
|
return referenced, err
|
|
}
|
|
|
|
var messageURLs []string
|
|
if err := db.WithContext(ctx).Unscoped().Model(&MessageImage{}).
|
|
Where("url LIKE ? OR url LIKE ?", "%://%", "//%").
|
|
Pluck("url", &messageURLs).Error; err != nil {
|
|
return false, errors.Wrap(err, "list legacy message image URLs")
|
|
}
|
|
referenced, err = hasLegacyCDNReference(messageURLs, key)
|
|
if err != nil || referenced {
|
|
return referenced, err
|
|
}
|
|
|
|
var profileUsers []*User
|
|
if err := db.WithContext(ctx).Unscoped().Model(&User{}).
|
|
Select("avatar", "background", "background_image").
|
|
Where("avatar = ? OR background = ? OR background_image = ? OR avatar LIKE ? OR background LIKE ? OR background_image LIKE ?", key, key, key, "%://%", "%://%", "%://%").
|
|
Find(&profileUsers).Error; err != nil {
|
|
return false, errors.Wrap(err, "list profile image references")
|
|
}
|
|
for _, user := range profileUsers {
|
|
referenced, err = hasLegacyCDNReference([]string{user.Avatar, user.Background, user.BackgroundImage}, key)
|
|
if err != nil || referenced {
|
|
return referenced, err
|
|
}
|
|
}
|
|
return false, nil
|
|
}
|
|
|
|
func isDefaultProfileObject(key string) bool {
|
|
for _, value := range []string{conf.Upload.DefaultAvatarURL, conf.Upload.DefaultBackground, conf.Upload.DefaultBackgroundImageURL} {
|
|
defaultKey, owned, err := r2ObjectKeyFromValue(value)
|
|
if err == nil && owned && defaultKey == key {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func hasLegacyCDNReference(values []string, targetKey string) (bool, error) {
|
|
for _, value := range values {
|
|
key, owned, err := r2ObjectKeyFromValue(value)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if owned && key == targetKey {
|
|
return true, nil
|
|
}
|
|
}
|
|
return false, nil
|
|
}
|
|
|
|
// r2ObjectKeyFromValue 接受当前对象 key 或旧 CDN URL。
|
|
// 外部 URL 不属于当前 CDN，安全返回 owned=false；畸形 URL 返回错误以避免误判无引用。
|
|
func r2ObjectKeyFromValue(value string) (key string, owned bool, err error) {
|
|
if value == "" {
|
|
return "", false, nil
|
|
}
|
|
if !strings.Contains(value, "://") && !strings.HasPrefix(value, "//") {
|
|
return value, true, nil
|
|
}
|
|
key, owned, err = storage.ObjectKeyFromCDNURL(value, conf.Upload.ImageBucketCDNHost)
|
|
if err != nil {
|
|
return "", false, errors.Wrap(err, "parse legacy CDN URL")
|
|
}
|
|
if !owned || key == "" {
|
|
return "", false, nil
|
|
}
|
|
return key, true, nil
|
|
}
|
|
|
|
func enqueueR2DeleteTasksForValues(tx *gorm.DB, bucket string, values []string, excludedKeys ...string) error {
|
|
if bucket == "" || len(values) == 0 {
|
|
return nil
|
|
}
|
|
excluded := make(map[string]struct{}, len(excludedKeys))
|
|
for _, key := range excludedKeys {
|
|
excluded[key] = struct{}{}
|
|
}
|
|
keys := make([]string, 0, len(values))
|
|
for _, value := range values {
|
|
key, owned, err := r2ObjectKeyFromValue(value)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !owned || isDefaultProfileObject(key) {
|
|
continue
|
|
}
|
|
if _, skip := excluded[key]; skip {
|
|
continue
|
|
}
|
|
keys = append(keys, key)
|
|
}
|
|
return EnqueueR2DeleteTasks(tx, bucket, keys)
|
|
}
|