可靠性: 完善 R2 备份恢复与对象生命周期
This commit is contained in:
tamakyi
+412
-143
@@ -6,40 +6,46 @@ package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
"github.com/aws/aws-sdk-go-v2/config"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/pkg/errors"
|
||||
"github.com/sirupsen/logrus"
|
||||
"github.com/urfave/cli/v2"
|
||||
|
||||
"github.com/tamakyi/TamaBox/internal/conf"
|
||||
"github.com/tamakyi/TamaBox/internal/storage"
|
||||
)
|
||||
|
||||
const (
|
||||
imageObjectAttemptTimeout = 60 * time.Second
|
||||
imageObjectMaxAttempts = 3
|
||||
imageObjectRetryBackoff = 100 * time.Millisecond
|
||||
)
|
||||
|
||||
// Backup 全量备份数据库与对象存储（Cloudflare R2）图片到本地目录。
|
||||
//
|
||||
// 为什么必须同时备份图片：站点的所有图片（头像 / 背景 / 二维码 logo / 提问图 / 回答图）
|
||||
// 都以 picture/YYYY/MM/DD/<hex>（新）或 tamabox/YYYY/MM/<hex>.<ext>（历史）的形式存在
|
||||
// R2 上，数据库里只存 key（upload_images.key）或 URL（users.avatar/background 里的完整
|
||||
// CDN 地址）。只备份数据库的话，恢复后这些引用指向的图片若已不在 R2（bucket 误删 /
|
||||
// 账号丢失 / 图片被清理），页面就全是裂图。因此备份必须把 R2 上的图片一并落盘。
|
||||
// 都以根日期 key、picture/ 或 tamabox/ 等形式存在 R2 上，数据库里只存 key 或 CDN URL。
|
||||
// 只备份数据库的话，恢复后这些引用指向的图片若已不在 R2，页面就全是裂图。
|
||||
//
|
||||
// 输出目录结构：
|
||||
//
|
||||
// <output>/
|
||||
// db-<20060102-150405>.sql # pg_dump 全量导出（schema + data，可用 psql 恢复）
|
||||
// images/<key> # R2 对象，保持原 key 路径
|
||||
// images-manifest.json # key → md5 / size 清单，用于恢复后完整性校验
|
||||
// images-manifest.json # key → 实际大小 / SHA-256 / GET 元数据
|
||||
//
|
||||
// 用法：TamaBox backup --output /var/backups/tamabox
|
||||
//
|
||||
@@ -47,7 +53,7 @@ import (
|
||||
// TamaBox backup --output . --images-only # 只备份图片
|
||||
var Backup = &cli.Command{
|
||||
Name: "backup",
|
||||
Usage: "Back up the database and R2 images to a local directory",
|
||||
Usage: "Back up the database and all R2 objects to a local directory",
|
||||
Action: runBackup,
|
||||
Flags: []cli.Flag{
|
||||
&cli.StringFlag{
|
||||
@@ -67,6 +73,9 @@ var Backup = &cli.Command{
|
||||
}
|
||||
|
||||
func runBackup(ctx *cli.Context) error {
|
||||
if ctx.Bool("db-only") && ctx.Bool("images-only") {
|
||||
return errors.New("--db-only and --images-only cannot be used together")
|
||||
}
|
||||
if err := conf.Init(); err != nil {
|
||||
return errors.Wrap(err, "load configuration")
|
||||
}
|
||||
@@ -127,190 +136,450 @@ func backupDatabase(ctx context.Context, outDir string) error {
|
||||
)
|
||||
cmd.Env = append(os.Environ(), "PGPASSWORD="+conf.Database.Password)
|
||||
|
||||
f, err := os.OpenFile(outFile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "create backup file %q", outFile)
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
var dumpErr strings.Builder
|
||||
cmd.Stdout = f
|
||||
cmd.Stderr = &dumpErr
|
||||
if err := cmd.Run(); err != nil {
|
||||
return errors.Wrapf(err, "pg_dump failed: %s", strings.TrimSpace(dumpErr.String()))
|
||||
if err := writeFileAtomically(outFile, func(w io.Writer) error {
|
||||
cmd.Stdout = w
|
||||
cmd.Stderr = &dumpErr
|
||||
if err := cmd.Run(); err != nil {
|
||||
return errors.Wrapf(err, "pg_dump failed: %s", strings.TrimSpace(dumpErr.String()))
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
logrus.WithField("file", outFile).Info("数据库已备份")
|
||||
return nil
|
||||
}
|
||||
|
||||
// imageManifestEntry 图片备份清单里的一条：key → 大小 → md5（用于恢复后校验完整性）。
|
||||
type imageManifestEntry struct {
|
||||
Key string `json:"key"`
|
||||
Size int64 `json:"size"`
|
||||
MD5 string `json:"md5,omitempty"`
|
||||
// writeFileAtomically 保留数据库、manifest 与对象下载共用的原子写修复：
|
||||
// 只有完整写入、Sync 并关闭成功后才替换目标，失败时清理临时 .part 文件。
|
||||
func writeFileAtomically(outFile string, write func(io.Writer) error) error {
|
||||
f, err := os.CreateTemp(filepath.Dir(outFile), "."+filepath.Base(outFile)+"-*.part")
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "create temporary file for %q", outFile)
|
||||
}
|
||||
tmpFile := f.Name()
|
||||
defer os.Remove(tmpFile)
|
||||
defer f.Close()
|
||||
|
||||
if err := write(f); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := f.Sync(); err != nil {
|
||||
return errors.Wrapf(err, "sync temporary file for %q", outFile)
|
||||
}
|
||||
if err := f.Close(); err != nil {
|
||||
return errors.Wrapf(err, "close temporary file for %q", outFile)
|
||||
}
|
||||
if err := os.Rename(tmpFile, outFile); err != nil {
|
||||
return errors.Wrapf(err, "finalize %q", outFile)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// backupImages 从 R2 下载 tamabox 站的全部图片对象到本地，并生成 manifest。
|
||||
//
|
||||
// R2 里 tamabox 站图片的 key 都带 `tamabox/` 前缀（如 tamabox/2025/12/<ts>.jpg），
|
||||
// 这是「bucket 用 path-style 访问时 List 返回的 key 带 bucket 名作前缀」这一 R2 行为。
|
||||
// 因此只需列出 `tamabox/` 前缀即可覆盖全部图片。blog/ 是同一 R2 account 下另一个站点
|
||||
// 的内容，不在 tamabox 的备份范围。
|
||||
// imageManifestEntry 记录实际下载内容的大小和 SHA-256，以及 GET 返回的原始元数据。
|
||||
// ETag 只是服务商元数据，不能当作内容 MD5 或完整性校验的替代品。
|
||||
type imageManifestEntry struct {
|
||||
Key string `json:"key"`
|
||||
Size int64 `json:"size"`
|
||||
SHA256 string `json:"sha256"`
|
||||
ETag string `json:"etag,omitempty"`
|
||||
ContentType string `json:"content_type,omitempty"`
|
||||
CacheControl string `json:"cache_control,omitempty"`
|
||||
}
|
||||
|
||||
type backupS3API interface {
|
||||
ListObjectsV2(context.Context, *s3.ListObjectsV2Input, ...func(*s3.Options)) (*s3.ListObjectsV2Output, error)
|
||||
GetObject(context.Context, *s3.GetObjectInput, ...func(*s3.Options)) (*s3.GetObjectOutput, error)
|
||||
}
|
||||
|
||||
type backupDownloadPlan struct {
|
||||
key string
|
||||
destination string
|
||||
}
|
||||
|
||||
type downloadedImageObject struct {
|
||||
size int64
|
||||
sha256 string
|
||||
etag string
|
||||
contentType string
|
||||
cacheControl string
|
||||
}
|
||||
|
||||
// backupImages 默认使用空前缀列举配置 bucket 的全部对象，而不是假定图片都在 tamabox/ 下。
|
||||
func backupImages(ctx context.Context, outDir string) error {
|
||||
client, err := backupS3Client()
|
||||
client, err := backupS3Client(ctx)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "get R2 client for backup")
|
||||
}
|
||||
return backupImagesWithClient(ctx, client, conf.Upload.ImageBucket, outDir, "")
|
||||
}
|
||||
|
||||
imageDir := filepath.Join(outDir, "images")
|
||||
if err := os.MkdirAll(imageDir, 0o700); err != nil {
|
||||
return errors.Wrap(err, "create images directory")
|
||||
}
|
||||
|
||||
var manifest []imageManifestEntry
|
||||
var count int64
|
||||
var totalBytes int64
|
||||
|
||||
// 见函数注释：tamabox/ 前缀即覆盖全部图片（含新旧 key 格式）。
|
||||
if err := backupPrefix(ctx, client, imageDir, "tamabox/", &manifest, &count, &totalBytes); err != nil {
|
||||
// backupImagesWithClient 允许测试和隔离演练显式指定 prefix。
|
||||
// 所有分页 List、key 校验和本地路径预检都在第一次 GetObject 之前完成。
|
||||
func backupImagesWithClient(ctx context.Context, client backupS3API, bucket, outDir, prefix string) error {
|
||||
keys, err := listAllBackupKeys(ctx, client, bucket, prefix)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 写 manifest：记录 key → 大小 → md5，恢复后可用于校验完整性。
|
||||
mf, err := os.Create(filepath.Join(outDir, "images-manifest.json"))
|
||||
imageDir := filepath.Join(outDir, "images")
|
||||
plans, err := planBackupDownloads(imageDir, keys)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "create manifest file")
|
||||
return err
|
||||
}
|
||||
defer mf.Close()
|
||||
enc := json.NewEncoder(mf)
|
||||
enc.SetIndent("", " ")
|
||||
if err := enc.Encode(manifest); err != nil {
|
||||
return errors.Wrap(err, "write manifest")
|
||||
|
||||
manifest := make([]imageManifestEntry, 0, len(plans))
|
||||
var totalBytes int64
|
||||
for _, plan := range plans {
|
||||
if err := ensureBackupDestinationDirectory(imageDir, plan.destination); err != nil {
|
||||
return err
|
||||
}
|
||||
downloaded, err := downloadObject(ctx, client, bucket, plan.key, plan.destination)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
manifest = append(manifest, imageManifestEntry{
|
||||
Key: plan.key,
|
||||
Size: downloaded.size,
|
||||
SHA256: downloaded.sha256,
|
||||
ETag: downloaded.etag,
|
||||
ContentType: downloaded.contentType,
|
||||
CacheControl: downloaded.cacheControl,
|
||||
})
|
||||
totalBytes += downloaded.size
|
||||
}
|
||||
|
||||
manifestFile := filepath.Join(outDir, "images-manifest.json")
|
||||
if err := writeFileAtomically(manifestFile, func(w io.Writer) error {
|
||||
enc := json.NewEncoder(w)
|
||||
enc.SetIndent("", " ")
|
||||
if err := enc.Encode(manifest); err != nil {
|
||||
return errors.Wrap(err, "write manifest")
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
logrus.WithFields(logrus.Fields{
|
||||
"objects": count,
|
||||
"objects": len(manifest),
|
||||
"bytes": totalBytes,
|
||||
"dir": imageDir,
|
||||
"prefix": prefix,
|
||||
}).Info("图片已备份")
|
||||
return nil
|
||||
}
|
||||
|
||||
// backupPrefix 分页列出并下载指定前缀下的全部对象。
|
||||
func backupPrefix(ctx context.Context, client *s3.Client, imageDir, prefix string, manifest *[]imageManifestEntry, count, totalBytes *int64) error {
|
||||
paginator := s3.NewListObjectsV2Paginator(client, &s3.ListObjectsV2Input{
|
||||
Bucket: aws.String(conf.Upload.ImageBucket),
|
||||
Prefix: aws.String(prefix),
|
||||
})
|
||||
for paginator.HasMorePages() {
|
||||
page, err := paginator.NextPage(ctx)
|
||||
func listAllBackupKeys(ctx context.Context, client backupS3API, bucket, prefix string) ([]string, error) {
|
||||
keys := make([]string, 0)
|
||||
var continuationToken *string
|
||||
seenTokens := make(map[string]struct{})
|
||||
|
||||
for {
|
||||
page, err := client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
|
||||
Bucket: aws.String(bucket),
|
||||
Prefix: aws.String(prefix),
|
||||
ContinuationToken: continuationToken,
|
||||
})
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "list R2 objects prefix %q", prefix)
|
||||
return nil, errors.Wrapf(err, "list R2 objects prefix %q", prefix)
|
||||
}
|
||||
if page == nil {
|
||||
return nil, errors.New("list R2 objects returned a nil page")
|
||||
}
|
||||
for _, object := range page.Contents {
|
||||
// nil key 与空 key 一样是无效输入；统一转成空串，由完整 List 后的预检拒绝。
|
||||
keys = append(keys, aws.ToString(object.Key))
|
||||
}
|
||||
for _, obj := range page.Contents {
|
||||
if obj.Key == nil {
|
||||
continue
|
||||
}
|
||||
key := *obj.Key
|
||||
|
||||
// 落盘路径 = imageDir + "/" + key，保持对象 key 结构（含 tamabox/ 前缀），
|
||||
// 恢复时可直接用同一 key 回传 R2。
|
||||
dst := filepath.Join(imageDir, filepath.FromSlash(key))
|
||||
if err := os.MkdirAll(filepath.Dir(dst), 0o700); err != nil {
|
||||
return errors.Wrapf(err, "create dir for %q", key)
|
||||
}
|
||||
if err := downloadObject(ctx, client, conf.Upload.ImageBucket, key, dst); err != nil {
|
||||
return err
|
||||
}
|
||||
if !aws.ToBool(page.IsTruncated) {
|
||||
return keys, nil
|
||||
}
|
||||
next := aws.ToString(page.NextContinuationToken)
|
||||
if next == "" {
|
||||
return nil, errors.New("truncated R2 list response has no continuation token")
|
||||
}
|
||||
if _, exists := seenTokens[next]; exists {
|
||||
return nil, errors.Errorf("R2 list continuation token %q repeated", next)
|
||||
}
|
||||
seenTokens[next] = struct{}{}
|
||||
continuationToken = aws.String(next)
|
||||
}
|
||||
}
|
||||
|
||||
size := int64(0)
|
||||
if obj.Size != nil {
|
||||
size = *obj.Size
|
||||
}
|
||||
md5 := ""
|
||||
if obj.ETag != nil {
|
||||
md5 = strings.Trim(*obj.ETag, `"`)
|
||||
}
|
||||
*manifest = append(*manifest, imageManifestEntry{Key: key, Size: size, MD5: md5})
|
||||
func planBackupDownloads(imageDir string, keys []string) ([]backupDownloadPlan, error) {
|
||||
plans := make([]backupDownloadPlan, 0, len(keys))
|
||||
files := make(map[string]string, len(keys))
|
||||
directories := make(map[string]string)
|
||||
|
||||
*count++
|
||||
*totalBytes += size
|
||||
root, err := filepath.Abs(imageDir)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "resolve images directory")
|
||||
}
|
||||
directories[canonicalLocalPath(root)] = ""
|
||||
|
||||
for _, key := range keys {
|
||||
destination, err := localImagePath(root, key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
canonicalDestination := canonicalLocalPath(destination)
|
||||
if previous, exists := files[canonicalDestination]; exists {
|
||||
return nil, errors.Errorf("objects %q and %q map to the same local file", previous, key)
|
||||
}
|
||||
if previous, exists := directories[canonicalDestination]; exists {
|
||||
return nil, errors.Errorf("file/directory collision between objects %q and %q", previous, key)
|
||||
}
|
||||
|
||||
for parent := filepath.Dir(destination); pathWithinDirectory(root, parent); parent = filepath.Dir(parent) {
|
||||
canonicalParent := canonicalLocalPath(parent)
|
||||
if previous, exists := files[canonicalParent]; exists {
|
||||
return nil, errors.Errorf("file/directory collision between objects %q and %q", previous, key)
|
||||
}
|
||||
if _, exists := directories[canonicalParent]; !exists {
|
||||
directories[canonicalParent] = key
|
||||
}
|
||||
if parent == root {
|
||||
break
|
||||
}
|
||||
}
|
||||
files[canonicalDestination] = key
|
||||
plans = append(plans, backupDownloadPlan{key: key, destination: destination})
|
||||
}
|
||||
|
||||
if err := prepareBackupImageRoot(root); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, plan := range plans {
|
||||
if err := preflightBackupLocalPath(root, plan.destination); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return plans, nil
|
||||
}
|
||||
|
||||
func validateObjectKey(key string) error {
|
||||
if key == "" {
|
||||
return errors.New("R2 object key must not be empty")
|
||||
}
|
||||
if strings.ContainsRune(key, 0) {
|
||||
return errors.Errorf("R2 object key %q contains NUL", key)
|
||||
}
|
||||
if strings.Contains(key, `\`) {
|
||||
return errors.Errorf("R2 object key %q contains a backslash", key)
|
||||
}
|
||||
if strings.HasSuffix(key, "/") {
|
||||
return errors.Errorf("R2 object key %q is a directory key", key)
|
||||
}
|
||||
if strings.HasPrefix(key, "/") || filepath.IsAbs(filepath.FromSlash(key)) || filepath.VolumeName(filepath.FromSlash(key)) != "" ||
|
||||
(len(key) >= 2 && ((key[0] >= 'A' && key[0] <= 'Z') || (key[0] >= 'a' && key[0] <= 'z')) && key[1] == ':') {
|
||||
return errors.Errorf("R2 object key %q is an absolute or volume-qualified path", key)
|
||||
}
|
||||
for _, segment := range strings.Split(key, "/") {
|
||||
if segment == "" || segment == "." || segment == ".." {
|
||||
return errors.Errorf("R2 object key %q contains an empty or dot path segment", key)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// backupS3Client 构造用于备份的 R2 客户端。
|
||||
//
|
||||
// 不能直接复用 storage.GetS3Client 做 List：那个客户端把 bucket 路径拼进了 endpoint
|
||||
// （image_endpoint 形如 https://<account>.r2.cloudflarestorage.com/<bucket>），
|
||||
// 它对 GetObject/PutObject 恰好可用（R2 容错），但 ListObjectsV2 会把 URL 拼成
|
||||
// /<bucket>/<bucket>?list-type=2，R2 找不到这个 key 前缀而返回 404 NoSuchKey。
|
||||
//
|
||||
// 这里剥离 endpoint 的 bucket 后缀，得到纯 account endpoint，并启用 path-style：
|
||||
// List/Get 都用 Bucket 参数 + 完整 key（List 返回的 key 自带 bucket 名作前缀，
|
||||
// 如 tamabox/2025/12/<ts>.jpg），这是 R2 上唯一稳定可用的姿势。
|
||||
func backupS3Client() (*s3.Client, error) {
|
||||
endpoint := strings.TrimRight(conf.Upload.ImageEndpoint, "/")
|
||||
// 剥离最后一个路径段（bucket 名）。endpoint 形如
|
||||
// https://<account>.r2.cloudflarestorage.com/<bucket>，去掉 /<bucket>。
|
||||
if i := strings.LastIndex(endpoint, "/"); i > len("https://") {
|
||||
endpoint = endpoint[:i]
|
||||
func localImagePath(imageDir, key string) (string, error) {
|
||||
if err := validateObjectKey(key); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
resolver := aws.EndpointResolverWithOptionsFunc(func(service, region string, options ...interface{}) (aws.Endpoint, error) {
|
||||
return aws.Endpoint{
|
||||
URL: endpoint,
|
||||
HostnameImmutable: true,
|
||||
Source: aws.EndpointSourceCustom,
|
||||
}, nil
|
||||
})
|
||||
cfg, err := config.LoadDefaultConfig(context.Background(),
|
||||
config.WithEndpointResolverWithOptions(resolver),
|
||||
config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(
|
||||
conf.Upload.ImageAccessID,
|
||||
conf.Upload.ImageAccessSecret,
|
||||
"",
|
||||
)),
|
||||
config.WithRegion("auto"),
|
||||
)
|
||||
root, err := filepath.Abs(imageDir)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "load R2 config")
|
||||
return "", errors.Wrap(err, "resolve images directory")
|
||||
}
|
||||
return s3.NewFromConfig(cfg, func(o *s3.Options) {
|
||||
o.UsePathStyle = true
|
||||
}), nil
|
||||
destination := filepath.Join(root, filepath.FromSlash(key))
|
||||
if !pathWithinDirectory(root, destination) || destination == root {
|
||||
return "", errors.Errorf("R2 object key %q escapes the images directory", key)
|
||||
}
|
||||
return destination, nil
|
||||
}
|
||||
|
||||
// downloadObject 下载单个对象到本地文件，失败时删除半成品避免留下损坏文件。
|
||||
func downloadObject(ctx context.Context, client *s3.Client, bucket, key, dst string) error {
|
||||
func pathWithinDirectory(root, candidate string) bool {
|
||||
rel, err := filepath.Rel(root, candidate)
|
||||
if err != nil || filepath.IsAbs(rel) {
|
||||
return false
|
||||
}
|
||||
return rel == "." || rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))
|
||||
}
|
||||
|
||||
func canonicalLocalPath(path string) string {
|
||||
cleaned := filepath.Clean(path)
|
||||
if runtime.GOOS == "windows" {
|
||||
return strings.ToLower(cleaned)
|
||||
}
|
||||
return cleaned
|
||||
}
|
||||
|
||||
func prepareBackupImageRoot(imageDir string) error {
|
||||
info, err := os.Lstat(imageDir)
|
||||
if os.IsNotExist(err) {
|
||||
if err := os.MkdirAll(imageDir, 0o700); err != nil {
|
||||
return errors.Wrap(err, "create images directory")
|
||||
}
|
||||
info, err = os.Lstat(imageDir)
|
||||
}
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "inspect images directory")
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
|
||||
return errors.Errorf("images path %q must be a real directory, not a symlink or file", imageDir)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func preflightBackupLocalPath(imageDir, destination string) error {
|
||||
info, err := os.Lstat(destination)
|
||||
if err == nil && (info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular()) {
|
||||
return errors.Errorf("backup destination %q is not a regular file", destination)
|
||||
}
|
||||
if err != nil && !os.IsNotExist(err) {
|
||||
return errors.Wrapf(err, "inspect backup destination %q", destination)
|
||||
}
|
||||
|
||||
for parent := filepath.Dir(destination); parent != imageDir; parent = filepath.Dir(parent) {
|
||||
if !pathWithinDirectory(imageDir, parent) {
|
||||
return errors.Errorf("backup destination %q escapes images directory", destination)
|
||||
}
|
||||
info, err := os.Lstat(parent)
|
||||
if os.IsNotExist(err) {
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "inspect backup parent %q", parent)
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
|
||||
return errors.Errorf("backup parent %q is not a real directory", parent)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func ensureBackupDestinationDirectory(imageDir, destination string) error {
|
||||
root, err := filepath.Abs(imageDir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
parent := filepath.Dir(destination)
|
||||
rel, err := filepath.Rel(root, parent)
|
||||
if err != nil || !pathWithinDirectory(root, parent) {
|
||||
return errors.Errorf("invalid backup destination %q", destination)
|
||||
}
|
||||
if rel == "." {
|
||||
return nil
|
||||
}
|
||||
|
||||
current := root
|
||||
for _, segment := range strings.Split(rel, string(filepath.Separator)) {
|
||||
current = filepath.Join(current, segment)
|
||||
info, err := os.Lstat(current)
|
||||
if os.IsNotExist(err) {
|
||||
if err := os.Mkdir(current, 0o700); err != nil && !os.IsExist(err) {
|
||||
return errors.Wrapf(err, "create backup directory %q", current)
|
||||
}
|
||||
info, err = os.Lstat(current)
|
||||
}
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "inspect backup directory %q", current)
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
|
||||
return errors.Errorf("backup parent %q is not a real directory", current)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// backupS3Client 与上传、恢复共用 storage.NewS3Client 的 endpoint 语义：
|
||||
// endpoint 可不带路径或只带 /<bucket>，请求始终以 Bucket 参数和原始完整 key 寻址。
|
||||
// List 返回什么 key 就下载什么 key，不额外添加或剥离 bucket 前缀。
|
||||
func backupS3Client(ctx context.Context) (*s3.Client, error) {
|
||||
return storage.NewS3Client(
|
||||
ctx,
|
||||
conf.Upload.ImageEndpoint,
|
||||
conf.Upload.ImageBucket,
|
||||
conf.Upload.ImageAccessID,
|
||||
conf.Upload.ImageAccessSecret,
|
||||
)
|
||||
}
|
||||
|
||||
// downloadObject 原子下载单个对象，并返回实际落盘内容的 SHA-256 和 GET 元数据。
|
||||
// 每次尝试都有独立 60 秒 context；流中断会从头重试整个对象，最多 3 次。
|
||||
func downloadObject(ctx context.Context, client backupS3API, bucket, key, destination string) (downloadedImageObject, error) {
|
||||
var lastErr error
|
||||
for attempt := 0; attempt < imageObjectMaxAttempts; attempt++ {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return downloadedImageObject{}, err
|
||||
}
|
||||
attemptCtx, cancel := context.WithTimeout(ctx, imageObjectAttemptTimeout)
|
||||
downloaded, err := downloadObjectOnce(attemptCtx, client, bucket, key, destination)
|
||||
cancel()
|
||||
if err == nil {
|
||||
return downloaded, nil
|
||||
}
|
||||
lastErr = err
|
||||
if ctx.Err() != nil {
|
||||
return downloadedImageObject{}, ctx.Err()
|
||||
}
|
||||
if attempt == imageObjectMaxAttempts-1 {
|
||||
break
|
||||
}
|
||||
if err := waitForImageObjectRetry(ctx, imageObjectRetryBackoff<<attempt); err != nil {
|
||||
return downloadedImageObject{}, err
|
||||
}
|
||||
}
|
||||
return downloadedImageObject{}, errors.Wrapf(lastErr, "download %q after %d attempts", key, imageObjectMaxAttempts)
|
||||
}
|
||||
|
||||
func downloadObjectOnce(ctx context.Context, client backupS3API, bucket, key, destination string) (downloadedImageObject, error) {
|
||||
out, err := client.GetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(key),
|
||||
})
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "get R2 object %q", key)
|
||||
return downloadedImageObject{}, errors.Wrapf(err, "get R2 object %q", key)
|
||||
}
|
||||
if out == nil || out.Body == nil {
|
||||
return downloadedImageObject{}, errors.Errorf("get R2 object %q returned an empty response body", key)
|
||||
}
|
||||
defer out.Body.Close()
|
||||
|
||||
tmp := dst + ".part"
|
||||
f, err := os.OpenFile(tmp, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "create %q", tmp)
|
||||
hash := sha256.New()
|
||||
var size int64
|
||||
writeErr := writeFileAtomically(destination, func(w io.Writer) error {
|
||||
var err error
|
||||
size, err = io.Copy(io.MultiWriter(w, hash), out.Body)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "download %q", key)
|
||||
}
|
||||
if out.ContentLength != nil && size != aws.ToInt64(out.ContentLength) {
|
||||
return errors.Wrapf(io.ErrUnexpectedEOF, "download %q: GET content length is %d, received %d", key, aws.ToInt64(out.ContentLength), size)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
_ = out.Body.Close()
|
||||
if writeErr != nil {
|
||||
return downloadedImageObject{}, writeErr
|
||||
}
|
||||
return downloadedImageObject{
|
||||
size: size,
|
||||
sha256: hex.EncodeToString(hash.Sum(nil)),
|
||||
etag: aws.ToString(out.ETag),
|
||||
contentType: aws.ToString(out.ContentType),
|
||||
cacheControl: aws.ToString(out.CacheControl),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func waitForImageObjectRetry(ctx context.Context, delay time.Duration) error {
|
||||
timer := time.NewTimer(delay)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
return nil
|
||||
}
|
||||
if _, err := io.Copy(f, out.Body); err != nil {
|
||||
f.Close()
|
||||
os.Remove(tmp)
|
||||
return errors.Wrapf(err, "download %q", key)
|
||||
}
|
||||
if err := f.Close(); err != nil {
|
||||
os.Remove(tmp)
|
||||
return errors.Wrapf(err, "close %q", tmp)
|
||||
}
|
||||
if err := os.Rename(tmp, dst); err != nil {
|
||||
os.Remove(tmp)
|
||||
return errors.Wrapf(err, "finalize %q", dst)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,437 @@
|
||||
// Copyright 2026 tamakyi. All rights reserved.
|
||||
// Use of this source code is governed by a MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/aws/smithy-go"
|
||||
"github.com/pkg/errors"
|
||||
"github.com/sirupsen/logrus"
|
||||
"github.com/urfave/cli/v2"
|
||||
|
||||
"github.com/tamakyi/TamaBox/internal/conf"
|
||||
"github.com/tamakyi/TamaBox/internal/storage"
|
||||
)
|
||||
|
||||
const maxImageManifestBytes = 64 << 20 // 64 MiB，足以容纳大量对象且避免恶意清单耗尽内存。
|
||||
|
||||
// RestoreImages 从 backup 生成的 images-manifest.json 和 images/ 恢复 R2 对象。
|
||||
var RestoreImages = &cli.Command{
|
||||
Name: "restore-images",
|
||||
Usage: "Validate and restore backed-up images to R2",
|
||||
Action: runRestoreImages,
|
||||
Flags: []cli.Flag{
|
||||
&cli.StringFlag{
|
||||
Name: "input",
|
||||
Required: true,
|
||||
Usage: "Backup directory or images-manifest.json path",
|
||||
},
|
||||
&cli.BoolFlag{
|
||||
Name: "dry-run",
|
||||
Usage: "Validate local files and remote targets without writing objects",
|
||||
},
|
||||
&cli.BoolFlag{
|
||||
Name: "yes",
|
||||
Usage: "Confirm actual remote writes",
|
||||
},
|
||||
&cli.BoolFlag{
|
||||
Name: "overwrite",
|
||||
Usage: "Allow replacing objects that already exist",
|
||||
},
|
||||
&cli.StringFlag{
|
||||
Name: "target-prefix",
|
||||
Value: "",
|
||||
Usage: "Prefix prepended to every restored key (for isolated drills)",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
type restoreImagesOptions struct {
|
||||
Input string
|
||||
DryRun bool
|
||||
Yes bool
|
||||
Overwrite bool
|
||||
TargetPrefix string
|
||||
}
|
||||
|
||||
// restoreObjectStore 的 PutObject 是带重试的高层操作，便于离线 fake 精确断言零写入。
|
||||
type restoreObjectStore interface {
|
||||
HeadObject(context.Context, *s3.HeadObjectInput, ...func(*s3.Options)) (*s3.HeadObjectOutput, error)
|
||||
PutObject(context.Context, *s3.PutObjectInput) error
|
||||
}
|
||||
|
||||
type r2RestoreObjectStore struct {
|
||||
client *s3.Client
|
||||
}
|
||||
|
||||
func (s *r2RestoreObjectStore) HeadObject(ctx context.Context, input *s3.HeadObjectInput, options ...func(*s3.Options)) (*s3.HeadObjectOutput, error) {
|
||||
return s.client.HeadObject(ctx, input, options...)
|
||||
}
|
||||
|
||||
func (s *r2RestoreObjectStore) PutObject(ctx context.Context, input *s3.PutObjectInput) error {
|
||||
// storage.PutObject 提供每次独立 60 秒超时、最多 3 次重试、Body 回绕和可取消退避。
|
||||
return storage.PutObject(ctx, s.client, input)
|
||||
}
|
||||
|
||||
type restorePlannedObject struct {
|
||||
entry imageManifestEntry
|
||||
sourcePath string
|
||||
targetKey string
|
||||
}
|
||||
|
||||
func runRestoreImages(ctx *cli.Context) error {
|
||||
opts := restoreImagesOptions{
|
||||
Input: ctx.String("input"),
|
||||
DryRun: ctx.Bool("dry-run"),
|
||||
Yes: ctx.Bool("yes"),
|
||||
Overwrite: ctx.Bool("overwrite"),
|
||||
TargetPrefix: ctx.String("target-prefix"),
|
||||
}
|
||||
if opts.Input == "" {
|
||||
return errors.New("--input must not be empty")
|
||||
}
|
||||
if !opts.DryRun && !opts.Yes {
|
||||
return errors.New("actual restore writes require --yes (or use --dry-run)")
|
||||
}
|
||||
if err := conf.Init(); err != nil {
|
||||
return errors.Wrap(err, "load configuration")
|
||||
}
|
||||
|
||||
client, err := backupS3Client(ctx.Context)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "get R2 client for image restore")
|
||||
}
|
||||
return restoreImagesWithStore(ctx.Context, &r2RestoreObjectStore{client: client}, conf.Upload.ImageBucket, opts)
|
||||
}
|
||||
|
||||
func restoreImagesWithStore(ctx context.Context, store restoreObjectStore, bucket string, opts restoreImagesOptions) error {
|
||||
if opts.Input == "" {
|
||||
return errors.New("restore input must not be empty")
|
||||
}
|
||||
if !opts.DryRun && !opts.Yes {
|
||||
return errors.New("actual restore writes require --yes (or use --dry-run)")
|
||||
}
|
||||
|
||||
manifestPath, imageDir, err := resolveRestoreInput(opts.Input)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
entries, err := readImageManifest(manifestPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 关键安全顺序：先验证所有 manifest 条目和全部本地文件，再做任何 Head/Put。
|
||||
// 因此重复 key、非法 hash/size、缺文件、篡改或 symlink 都不会造成部分上传。
|
||||
objects := make([]restorePlannedObject, 0, len(entries))
|
||||
for _, entry := range entries {
|
||||
targetKey, err := restoreTargetKey(opts.TargetPrefix, entry.Key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
sourcePath, err := localImagePath(imageDir, entry.Key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := validateRestoreLocalFile(imageDir, sourcePath, entry); err != nil {
|
||||
return err
|
||||
}
|
||||
objects = append(objects, restorePlannedObject{
|
||||
entry: entry,
|
||||
sourcePath: sourcePath,
|
||||
targetKey: targetKey,
|
||||
})
|
||||
}
|
||||
|
||||
// 默认先 Head 所有目标。任何目标已存在都会在 Put 循环开始前失败，保证零 Put。
|
||||
if !opts.Overwrite {
|
||||
for _, object := range objects {
|
||||
headCtx, cancel := context.WithTimeout(ctx, imageObjectAttemptTimeout)
|
||||
_, err := store.HeadObject(headCtx, &s3.HeadObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(object.targetKey),
|
||||
})
|
||||
cancel()
|
||||
if err == nil {
|
||||
return errors.Errorf("restore target %q already exists; no objects were uploaded (use --overwrite to replace)", object.targetKey)
|
||||
}
|
||||
if !isS3ObjectNotFound(err) {
|
||||
return errors.Wrapf(err, "check restore target %q; no objects were uploaded", object.targetKey)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if opts.DryRun {
|
||||
logrus.WithFields(logrus.Fields{
|
||||
"objects": len(objects),
|
||||
"prefix": opts.TargetPrefix,
|
||||
}).Info("图片恢复演练完成，未写入对象")
|
||||
return nil
|
||||
}
|
||||
|
||||
succeeded := 0
|
||||
for _, object := range objects {
|
||||
file, err := openRestoreFileForPut(object.sourcePath, object.entry.Size)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "open validated restore file for %q after %d successful uploads; restore is not transactional", object.entry.Key, succeeded)
|
||||
}
|
||||
input := &s3.PutObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(object.targetKey),
|
||||
Body: file,
|
||||
ContentLength: aws.Int64(object.entry.Size),
|
||||
ContentType: nonEmptyStringPointer(object.entry.ContentType),
|
||||
CacheControl: nonEmptyStringPointer(object.entry.CacheControl),
|
||||
}
|
||||
putErr := store.PutObject(ctx, input)
|
||||
closeErr := file.Close()
|
||||
if putErr != nil {
|
||||
return errors.Wrapf(putErr, "restore object %q to %q failed after %d successful uploads; restore is not transactional", object.entry.Key, object.targetKey, succeeded)
|
||||
}
|
||||
if closeErr != nil {
|
||||
return errors.Wrapf(closeErr, "close restored file %q (remote object was uploaded); restore is not transactional", object.sourcePath)
|
||||
}
|
||||
succeeded++
|
||||
}
|
||||
|
||||
logrus.WithFields(logrus.Fields{
|
||||
"objects": succeeded,
|
||||
"prefix": opts.TargetPrefix,
|
||||
}).Info("图片恢复完成")
|
||||
return nil
|
||||
}
|
||||
|
||||
func resolveRestoreInput(input string) (manifestPath, imageDir string, err error) {
|
||||
info, err := os.Lstat(input)
|
||||
if err != nil {
|
||||
return "", "", errors.Wrapf(err, "inspect restore input %q", input)
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink != 0 {
|
||||
return "", "", errors.Errorf("restore input %q must not be a symlink", input)
|
||||
}
|
||||
|
||||
manifestPath = input
|
||||
if info.IsDir() {
|
||||
manifestPath = filepath.Join(input, "images-manifest.json")
|
||||
} else if !info.Mode().IsRegular() {
|
||||
return "", "", errors.Errorf("restore input %q is neither a directory nor a regular manifest file", input)
|
||||
}
|
||||
return manifestPath, filepath.Join(filepath.Dir(manifestPath), "images"), nil
|
||||
}
|
||||
|
||||
func readImageManifest(manifestPath string) ([]imageManifestEntry, error) {
|
||||
info, err := os.Lstat(manifestPath)
|
||||
if err != nil {
|
||||
return nil, errors.Wrapf(err, "inspect image manifest %q", manifestPath)
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
|
||||
return nil, errors.Errorf("image manifest %q must be a regular non-symlink file", manifestPath)
|
||||
}
|
||||
if info.Size() > maxImageManifestBytes {
|
||||
return nil, errors.Errorf("image manifest %q exceeds the %d-byte limit", manifestPath, maxImageManifestBytes)
|
||||
}
|
||||
|
||||
file, err := os.Open(manifestPath)
|
||||
if err != nil {
|
||||
return nil, errors.Wrapf(err, "open image manifest %q", manifestPath)
|
||||
}
|
||||
data, readErr := io.ReadAll(io.LimitReader(file, maxImageManifestBytes+1))
|
||||
closeErr := file.Close()
|
||||
if readErr != nil {
|
||||
return nil, errors.Wrap(readErr, "read image manifest")
|
||||
}
|
||||
if closeErr != nil {
|
||||
return nil, errors.Wrap(closeErr, "close image manifest")
|
||||
}
|
||||
if len(data) > maxImageManifestBytes {
|
||||
return nil, errors.Errorf("image manifest %q exceeds the %d-byte limit", manifestPath, maxImageManifestBytes)
|
||||
}
|
||||
|
||||
decoder := json.NewDecoder(bytes.NewReader(data))
|
||||
var entries []imageManifestEntry
|
||||
if err := decoder.Decode(&entries); err != nil {
|
||||
return nil, errors.Wrap(err, "decode image manifest")
|
||||
}
|
||||
if entries == nil {
|
||||
return nil, errors.New("image manifest must be a JSON array, not null")
|
||||
}
|
||||
var trailing json.RawMessage
|
||||
if err := decoder.Decode(&trailing); err != io.EOF {
|
||||
if err == nil {
|
||||
return nil, errors.New("image manifest contains trailing JSON")
|
||||
}
|
||||
return nil, errors.Wrap(err, "image manifest contains trailing data")
|
||||
}
|
||||
|
||||
seen := make(map[string]struct{}, len(entries))
|
||||
for index := range entries {
|
||||
entry := &entries[index]
|
||||
if err := validateObjectKey(entry.Key); err != nil {
|
||||
return nil, errors.Wrapf(err, "invalid manifest entry %d", index)
|
||||
}
|
||||
if _, exists := seen[entry.Key]; exists {
|
||||
return nil, errors.Errorf("image manifest contains duplicate key %q", entry.Key)
|
||||
}
|
||||
seen[entry.Key] = struct{}{}
|
||||
if entry.Size < 0 {
|
||||
return nil, errors.Errorf("manifest entry %q has negative size %d", entry.Key, entry.Size)
|
||||
}
|
||||
decodedHash, err := hex.DecodeString(entry.SHA256)
|
||||
if err != nil || len(decodedHash) != sha256.Size {
|
||||
return nil, errors.Errorf("manifest entry %q has invalid SHA-256 %q", entry.Key, entry.SHA256)
|
||||
}
|
||||
entry.SHA256 = strings.ToLower(entry.SHA256)
|
||||
}
|
||||
return entries, nil
|
||||
}
|
||||
|
||||
func restoreTargetKey(prefix, key string) (string, error) {
|
||||
if prefix != "" && !strings.HasSuffix(prefix, "/") {
|
||||
prefix += "/"
|
||||
}
|
||||
target := prefix + key
|
||||
if err := validateObjectKey(target); err != nil {
|
||||
return "", errors.Wrapf(err, "invalid restore target for source key %q", key)
|
||||
}
|
||||
return target, nil
|
||||
}
|
||||
|
||||
func validateRestoreLocalFile(imageDir, sourcePath string, entry imageManifestEntry) error {
|
||||
root, err := filepath.Abs(imageDir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := validateRestoreDirectoryChain(root, filepath.Dir(sourcePath)); err != nil {
|
||||
return err
|
||||
}
|
||||
info, err := os.Lstat(sourcePath)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "inspect local image for %q", entry.Key)
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
|
||||
return errors.Errorf("local image for %q must be a regular non-symlink file", entry.Key)
|
||||
}
|
||||
if info.Size() != entry.Size {
|
||||
return errors.Errorf("local image size mismatch for %q: manifest=%d actual=%d", entry.Key, entry.Size, info.Size())
|
||||
}
|
||||
|
||||
file, err := os.Open(sourcePath)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "open local image for %q", entry.Key)
|
||||
}
|
||||
openedInfo, statErr := file.Stat()
|
||||
if statErr != nil {
|
||||
_ = file.Close()
|
||||
return errors.Wrapf(statErr, "stat opened local image for %q", entry.Key)
|
||||
}
|
||||
if !openedInfo.Mode().IsRegular() || openedInfo.Size() != entry.Size {
|
||||
_ = file.Close()
|
||||
return errors.Errorf("opened local image for %q is not the expected regular file", entry.Key)
|
||||
}
|
||||
hash := sha256.New()
|
||||
_, hashErr := io.Copy(hash, file)
|
||||
closeErr := file.Close()
|
||||
if hashErr != nil {
|
||||
return errors.Wrapf(hashErr, "hash local image for %q", entry.Key)
|
||||
}
|
||||
if closeErr != nil {
|
||||
return errors.Wrapf(closeErr, "close local image for %q", entry.Key)
|
||||
}
|
||||
actualHash := hex.EncodeToString(hash.Sum(nil))
|
||||
if actualHash != entry.SHA256 {
|
||||
return errors.Errorf("local image SHA-256 mismatch for %q: manifest=%s actual=%s", entry.Key, entry.SHA256, actualHash)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateRestoreDirectoryChain(root, parent string) error {
|
||||
rootInfo, err := os.Lstat(root)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "inspect restore images directory %q", root)
|
||||
}
|
||||
if rootInfo.Mode()&os.ModeSymlink != 0 || !rootInfo.IsDir() {
|
||||
return errors.Errorf("restore images path %q must be a real directory", root)
|
||||
}
|
||||
rel, err := filepath.Rel(root, parent)
|
||||
if err != nil || !pathWithinDirectory(root, parent) {
|
||||
return errors.Errorf("restore file parent %q escapes images directory", parent)
|
||||
}
|
||||
if rel == "." {
|
||||
return nil
|
||||
}
|
||||
current := root
|
||||
for _, segment := range strings.Split(rel, string(filepath.Separator)) {
|
||||
current = filepath.Join(current, segment)
|
||||
info, err := os.Lstat(current)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "inspect restore parent directory %q", current)
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
|
||||
return errors.Errorf("restore parent %q must be a real directory", current)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func openRestoreFileForPut(sourcePath string, expectedSize int64) (*os.File, error) {
|
||||
info, err := os.Lstat(sourcePath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() || info.Size() != expectedSize {
|
||||
return nil, fmt.Errorf("local file %q changed after validation", sourcePath)
|
||||
}
|
||||
file, err := os.Open(sourcePath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
openedInfo, err := file.Stat()
|
||||
if err != nil {
|
||||
_ = file.Close()
|
||||
return nil, err
|
||||
}
|
||||
if !openedInfo.Mode().IsRegular() || openedInfo.Size() != expectedSize {
|
||||
_ = file.Close()
|
||||
return nil, fmt.Errorf("local file %q changed after validation", sourcePath)
|
||||
}
|
||||
return file, nil
|
||||
}
|
||||
|
||||
func isS3ObjectNotFound(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var statusErr interface{ HTTPStatusCode() int }
|
||||
if errors.As(err, &statusErr) && statusErr.HTTPStatusCode() == 404 {
|
||||
return true
|
||||
}
|
||||
var apiErr smithy.APIError
|
||||
if errors.As(err, &apiErr) {
|
||||
switch apiErr.ErrorCode() {
|
||||
case "NotFound", "NoSuchKey", "404":
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func nonEmptyStringPointer(value string) *string {
|
||||
if value == "" {
|
||||
return nil
|
||||
}
|
||||
return aws.String(value)
|
||||
}
|
||||
@@ -0,0 +1,184 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/tamakyi/TamaBox/internal/conf"
|
||||
"github.com/tamakyi/TamaBox/internal/storage"
|
||||
)
|
||||
|
||||
// TestR2BackupRestoreIntegration 默认跳过。只有显式设置 TB_R2_INTEGRATION=1，
|
||||
// 并通过 NEKOBOX_CONFIG_PATH 提供真实 R2 配置时才运行。
|
||||
// 测试只访问 tamabox-backup-test/<随机>/ 下的 source/ 与 target/ 两个专属前缀，
|
||||
// 绝不执行全桶 List，也不会读取、覆盖或删除其他 key。
|
||||
func TestR2BackupRestoreIntegration(t *testing.T) {
|
||||
if os.Getenv("TB_R2_INTEGRATION") != "1" {
|
||||
t.Skip("set TB_R2_INTEGRATION=1 to run the real R2 integration test")
|
||||
}
|
||||
if os.Getenv("NEKOBOX_CONFIG_PATH") == "" {
|
||||
t.Fatal("NEKOBOX_CONFIG_PATH must point to the R2 integration configuration")
|
||||
}
|
||||
require.NoError(t, conf.Init())
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
|
||||
defer cancel()
|
||||
client, err := storage.NewS3Client(
|
||||
ctx,
|
||||
conf.Upload.ImageEndpoint,
|
||||
conf.Upload.ImageBucket,
|
||||
conf.Upload.ImageAccessID,
|
||||
conf.Upload.ImageAccessSecret,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
random := make([]byte, 12)
|
||||
_, err = rand.Read(random)
|
||||
require.NoError(t, err)
|
||||
basePrefix := fmt.Sprintf("tamabox-backup-test/%d-%s/", time.Now().UTC().UnixNano(), hex.EncodeToString(random))
|
||||
sourcePrefix := basePrefix + "source/"
|
||||
targetPrefix := basePrefix + "target/"
|
||||
|
||||
objects := []struct {
|
||||
relativeKey string
|
||||
body []byte
|
||||
contentType string
|
||||
cacheControl string
|
||||
}{
|
||||
{
|
||||
relativeKey: "2026/11/04/root.webp",
|
||||
body: []byte("tamabox-r2-integration-root"),
|
||||
contentType: "image/webp",
|
||||
cacheControl: "public, max-age=60",
|
||||
},
|
||||
{
|
||||
relativeKey: "picture/2026/11/04/picture.webp",
|
||||
body: []byte("tamabox-r2-integration-picture"),
|
||||
contentType: "image/webp",
|
||||
cacheControl: "public, max-age=120",
|
||||
},
|
||||
}
|
||||
|
||||
// created 只记录本测试成功创建的精确 key。任何失败清理也只删除这些 key。
|
||||
created := make(map[string]bool)
|
||||
t.Cleanup(func() {
|
||||
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||||
defer cleanupCancel()
|
||||
for key := range created {
|
||||
if err := storage.DeleteObject(cleanupCtx, client, conf.Upload.ImageBucket, key); err != nil {
|
||||
t.Errorf("clean up own R2 key %q: %v", key, err)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
for _, object := range objects {
|
||||
key := sourcePrefix + object.relativeKey
|
||||
err := storage.PutObject(ctx, client, &s3.PutObjectInput{
|
||||
Bucket: aws.String(conf.Upload.ImageBucket),
|
||||
Key: aws.String(key),
|
||||
Body: bytes.NewReader(object.body),
|
||||
ContentLength: aws.Int64(int64(len(object.body))),
|
||||
ContentType: aws.String(object.contentType),
|
||||
CacheControl: aws.String(object.cacheControl),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
created[key] = true
|
||||
}
|
||||
|
||||
backupDir := t.TempDir()
|
||||
require.NoError(t, backupImagesWithClient(ctx, client, conf.Upload.ImageBucket, backupDir, sourcePrefix))
|
||||
manifest, err := readImageManifest(backupDir + string(os.PathSeparator) + "images-manifest.json")
|
||||
require.NoError(t, err)
|
||||
require.Len(t, manifest, len(objects))
|
||||
manifestKeys := make(map[string]bool, len(manifest))
|
||||
for _, entry := range manifest {
|
||||
manifestKeys[entry.Key] = true
|
||||
}
|
||||
for _, object := range objects {
|
||||
require.True(t, manifestKeys[sourcePrefix+object.relativeKey])
|
||||
}
|
||||
|
||||
for _, object := range objects {
|
||||
sourceKey := sourcePrefix + object.relativeKey
|
||||
require.NoError(t, storage.DeleteObject(ctx, client, conf.Upload.ImageBucket, sourceKey))
|
||||
delete(created, sourceKey)
|
||||
}
|
||||
|
||||
store := &integrationTrackingRestoreStore{
|
||||
base: &r2RestoreObjectStore{client: client},
|
||||
onUpload: func(key string) {
|
||||
created[key] = true
|
||||
},
|
||||
}
|
||||
require.NoError(t, restoreImagesWithStore(ctx, store, conf.Upload.ImageBucket, restoreImagesOptions{
|
||||
Input: backupDir,
|
||||
Yes: true,
|
||||
TargetPrefix: targetPrefix,
|
||||
}))
|
||||
|
||||
for _, object := range objects {
|
||||
targetKey := targetPrefix + sourcePrefix + object.relativeKey
|
||||
require.True(t, created[targetKey], "restore must report its own successful Put")
|
||||
output, err := client.GetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: aws.String(conf.Upload.ImageBucket),
|
||||
Key: aws.String(targetKey),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
body, readErr := io.ReadAll(output.Body)
|
||||
closeErr := output.Body.Close()
|
||||
require.NoError(t, readErr)
|
||||
require.NoError(t, closeErr)
|
||||
digest := sha256.Sum256(object.body)
|
||||
require.Equal(t, fmt.Sprintf("%x", digest), fmt.Sprintf("%x", sha256.Sum256(body)))
|
||||
require.Equal(t, object.contentType, aws.ToString(output.ContentType))
|
||||
require.Equal(t, object.cacheControl, aws.ToString(output.CacheControl))
|
||||
}
|
||||
|
||||
for key := range created {
|
||||
require.NoError(t, storage.DeleteObject(ctx, client, conf.Upload.ImageBucket, key))
|
||||
delete(created, key)
|
||||
}
|
||||
requireOwnPrefixEmpty(ctx, t, client, conf.Upload.ImageBucket, sourcePrefix)
|
||||
requireOwnPrefixEmpty(ctx, t, client, conf.Upload.ImageBucket, targetPrefix)
|
||||
}
|
||||
|
||||
type integrationTrackingRestoreStore struct {
|
||||
base restoreObjectStore
|
||||
onUpload func(string)
|
||||
}
|
||||
|
||||
func (s *integrationTrackingRestoreStore) HeadObject(ctx context.Context, input *s3.HeadObjectInput, options ...func(*s3.Options)) (*s3.HeadObjectOutput, error) {
|
||||
return s.base.HeadObject(ctx, input, options...)
|
||||
}
|
||||
|
||||
func (s *integrationTrackingRestoreStore) PutObject(ctx context.Context, input *s3.PutObjectInput) error {
|
||||
if err := s.base.PutObject(ctx, input); err != nil {
|
||||
return err
|
||||
}
|
||||
s.onUpload(aws.ToString(input.Key))
|
||||
return nil
|
||||
}
|
||||
|
||||
func requireOwnPrefixEmpty(ctx context.Context, t *testing.T, client *s3.Client, bucket, prefix string) {
|
||||
t.Helper()
|
||||
require.NotEmpty(t, prefix)
|
||||
output, err := client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
|
||||
Bucket: aws.String(bucket),
|
||||
Prefix: aws.String(prefix),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, output.Contents, "专属 prefix 清理后必须为空：%q", prefix)
|
||||
require.False(t, aws.ToBool(output.IsTruncated), "随机专属 prefix 下对象数不应触发分页")
|
||||
}
|
||||
@@ -0,0 +1,603 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3/types"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type fakeBackupObject struct {
|
||||
body []byte
|
||||
etag string
|
||||
contentType string
|
||||
cacheControl string
|
||||
interruptedAttempts int
|
||||
wrongLengthAttempts int
|
||||
}
|
||||
|
||||
type fakeBackupS3 struct {
|
||||
pages []*s3.ListObjectsV2Output
|
||||
listInputs []*s3.ListObjectsV2Input
|
||||
objects map[string]fakeBackupObject
|
||||
getCalls map[string]int
|
||||
getOrder []string
|
||||
}
|
||||
|
||||
func (f *fakeBackupS3) ListObjectsV2(_ context.Context, input *s3.ListObjectsV2Input, _ ...func(*s3.Options)) (*s3.ListObjectsV2Output, error) {
|
||||
f.listInputs = append(f.listInputs, input)
|
||||
index := len(f.listInputs) - 1
|
||||
if index >= len(f.pages) {
|
||||
return nil, fmt.Errorf("unexpected ListObjectsV2 call %d", index+1)
|
||||
}
|
||||
return f.pages[index], nil
|
||||
}
|
||||
|
||||
func (f *fakeBackupS3) GetObject(_ context.Context, input *s3.GetObjectInput, _ ...func(*s3.Options)) (*s3.GetObjectOutput, error) {
|
||||
key := aws.ToString(input.Key)
|
||||
if f.getCalls == nil {
|
||||
f.getCalls = make(map[string]int)
|
||||
}
|
||||
f.getCalls[key]++
|
||||
f.getOrder = append(f.getOrder, key)
|
||||
object, exists := f.objects[key]
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("unexpected GetObject key %q", key)
|
||||
}
|
||||
|
||||
var body io.ReadCloser
|
||||
if f.getCalls[key] <= object.interruptedAttempts {
|
||||
body = io.NopCloser(&interruptedReader{data: append([]byte(nil), object.body...)})
|
||||
} else {
|
||||
body = io.NopCloser(strings.NewReader(string(object.body)))
|
||||
}
|
||||
contentLength := int64(len(object.body))
|
||||
if f.getCalls[key] <= object.wrongLengthAttempts {
|
||||
contentLength++
|
||||
}
|
||||
return &s3.GetObjectOutput{
|
||||
Body: body,
|
||||
ContentLength: aws.Int64(contentLength),
|
||||
ETag: aws.String(object.etag),
|
||||
ContentType: aws.String(object.contentType),
|
||||
CacheControl: aws.String(object.cacheControl),
|
||||
}, nil
|
||||
}
|
||||
|
||||
type interruptedReader struct {
|
||||
data []byte
|
||||
}
|
||||
|
||||
func (r *interruptedReader) Read(p []byte) (int, error) {
|
||||
if len(r.data) == 0 {
|
||||
return 0, io.ErrUnexpectedEOF
|
||||
}
|
||||
n := copy(p, r.data)
|
||||
r.data = r.data[n:]
|
||||
if len(r.data) == 0 {
|
||||
return n, io.ErrUnexpectedEOF
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func backupListPage(keys []string, truncated bool, nextToken string) *s3.ListObjectsV2Output {
|
||||
objects := make([]types.Object, 0, len(keys))
|
||||
for _, key := range keys {
|
||||
objects = append(objects, types.Object{
|
||||
Key: aws.String(key),
|
||||
Size: aws.Int64(999999), // 故意不可信；manifest 必须记录实际下载字节数。
|
||||
ETag: aws.String("list-etag-must-not-be-used"),
|
||||
})
|
||||
}
|
||||
output := &s3.ListObjectsV2Output{Contents: objects, IsTruncated: aws.Bool(truncated)}
|
||||
if nextToken != "" {
|
||||
output.NextContinuationToken = aws.String(nextToken)
|
||||
}
|
||||
return output
|
||||
}
|
||||
|
||||
func TestBackupImagesWithClientCoversAllKeysAndPages(t *testing.T) {
|
||||
keys := []string{
|
||||
"2026/11/04/root.webp",
|
||||
"picture/2026/11/04/picture.webp",
|
||||
"tamabox/2025/12/historical.jpg",
|
||||
"static/site.css",
|
||||
}
|
||||
objects := make(map[string]fakeBackupObject, len(keys))
|
||||
for index, key := range keys {
|
||||
objects[key] = fakeBackupObject{
|
||||
body: []byte(fmt.Sprintf("actual-body-%d", index)),
|
||||
etag: fmt.Sprintf("get-etag-%d", index),
|
||||
contentType: []string{"image/webp", "image/webp", "image/jpeg", "text/css"}[index],
|
||||
cacheControl: fmt.Sprintf("public, max-age=%d", index+60),
|
||||
}
|
||||
}
|
||||
client := &fakeBackupS3{
|
||||
pages: []*s3.ListObjectsV2Output{
|
||||
backupListPage(keys[:2], true, "page-2"),
|
||||
backupListPage(keys[2:], false, ""),
|
||||
},
|
||||
objects: objects,
|
||||
}
|
||||
outDir := t.TempDir()
|
||||
require.NoError(t, backupImagesWithClient(context.Background(), client, "configured-bucket", outDir, ""))
|
||||
|
||||
require.Len(t, client.listInputs, 2)
|
||||
for index, input := range client.listInputs {
|
||||
require.Equal(t, "configured-bucket", aws.ToString(input.Bucket))
|
||||
require.Equal(t, "", aws.ToString(input.Prefix), "生产默认必须列举整个 bucket")
|
||||
if index == 0 {
|
||||
require.Nil(t, input.ContinuationToken)
|
||||
} else {
|
||||
require.Equal(t, "page-2", aws.ToString(input.ContinuationToken))
|
||||
}
|
||||
}
|
||||
require.ElementsMatch(t, keys, client.getOrder)
|
||||
|
||||
manifestBytes, err := os.ReadFile(filepath.Join(outDir, "images-manifest.json"))
|
||||
require.NoError(t, err)
|
||||
var manifest []imageManifestEntry
|
||||
require.NoError(t, json.Unmarshal(manifestBytes, &manifest))
|
||||
require.Len(t, manifest, len(keys))
|
||||
for index, entry := range manifest {
|
||||
require.Equal(t, keys[index], entry.Key)
|
||||
object := objects[entry.Key]
|
||||
require.Equal(t, int64(len(object.body)), entry.Size)
|
||||
digest := sha256.Sum256(object.body)
|
||||
require.Equal(t, fmt.Sprintf("%x", digest), entry.SHA256)
|
||||
require.Equal(t, object.etag, entry.ETag)
|
||||
require.Equal(t, object.contentType, entry.ContentType)
|
||||
require.Equal(t, object.cacheControl, entry.CacheControl)
|
||||
|
||||
path, err := localImagePath(filepath.Join(outDir, "images"), entry.Key)
|
||||
require.NoError(t, err)
|
||||
content, err := os.ReadFile(path)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, object.body, content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackupImagesWithClientSupportsIsolatedPrefixAndEmptyManifest(t *testing.T) {
|
||||
t.Run("prefix", func(t *testing.T) {
|
||||
key := "tamabox-backup-test/unit/source/root.webp"
|
||||
client := &fakeBackupS3{
|
||||
pages: []*s3.ListObjectsV2Output{backupListPage([]string{key}, false, "")},
|
||||
objects: map[string]fakeBackupObject{
|
||||
key: {body: []byte("isolated"), contentType: "image/webp"},
|
||||
},
|
||||
}
|
||||
require.NoError(t, backupImagesWithClient(context.Background(), client, "bucket", t.TempDir(), "tamabox-backup-test/unit/source/"))
|
||||
require.Len(t, client.listInputs, 1)
|
||||
require.Equal(t, "tamabox-backup-test/unit/source/", aws.ToString(client.listInputs[0].Prefix))
|
||||
require.Equal(t, []string{key}, client.getOrder)
|
||||
})
|
||||
|
||||
t.Run("empty array", func(t *testing.T) {
|
||||
client := &fakeBackupS3{pages: []*s3.ListObjectsV2Output{backupListPage(nil, false, "")}}
|
||||
outDir := t.TempDir()
|
||||
require.NoError(t, backupImagesWithClient(context.Background(), client, "bucket", outDir, ""))
|
||||
content, err := os.ReadFile(filepath.Join(outDir, "images-manifest.json"))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "[]", strings.TrimSpace(string(content)))
|
||||
})
|
||||
}
|
||||
|
||||
func TestBackupRejectsPathAttacksBeforeAnyDownload(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
key *string
|
||||
}{
|
||||
{name: "nil key", key: nil},
|
||||
{name: "empty key", key: aws.String("")},
|
||||
{name: "absolute POSIX", key: aws.String("/outside")},
|
||||
{name: "absolute Windows", key: aws.String("C:/outside")},
|
||||
{name: "volume relative", key: aws.String("C:outside")},
|
||||
{name: "parent", key: aws.String("../outside")},
|
||||
{name: "nested parent", key: aws.String("picture/../outside")},
|
||||
{name: "backslash", key: aws.String(`picture\outside`)},
|
||||
{name: "NUL", key: aws.String("picture/\x00outside")},
|
||||
{name: "empty segment", key: aws.String("picture//outside")},
|
||||
{name: "leading dot segment", key: aws.String("./outside")},
|
||||
{name: "inner dot segment", key: aws.String("picture/./outside")},
|
||||
{name: "directory", key: aws.String("picture/")},
|
||||
{name: "dot only", key: aws.String(".")},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
page := backupListPage(nil, false, "")
|
||||
page.Contents = []types.Object{{Key: tt.key}}
|
||||
client := &fakeBackupS3{pages: []*s3.ListObjectsV2Output{page}}
|
||||
err := backupImagesWithClient(context.Background(), client, "bucket", t.TempDir(), "")
|
||||
require.Error(t, err)
|
||||
require.Empty(t, client.getOrder, "非法 key 必须在任何下载前被拒绝")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackupPreflightsAllPagesAndLocalCollisionsBeforeDownload(t *testing.T) {
|
||||
t.Run("invalid key on later page", func(t *testing.T) {
|
||||
client := &fakeBackupS3{
|
||||
pages: []*s3.ListObjectsV2Output{
|
||||
backupListPage([]string{"picture/valid.webp"}, true, "next"),
|
||||
backupListPage([]string{"../escape"}, false, ""),
|
||||
},
|
||||
objects: map[string]fakeBackupObject{"picture/valid.webp": {body: []byte("valid")}},
|
||||
}
|
||||
require.Error(t, backupImagesWithClient(context.Background(), client, "bucket", t.TempDir(), ""))
|
||||
require.Len(t, client.listInputs, 2)
|
||||
require.Empty(t, client.getOrder)
|
||||
})
|
||||
|
||||
for _, keys := range [][]string{{"a", "a/b"}, {"a/b", "a"}, {"same", "same"}} {
|
||||
t.Run(fmt.Sprintf("collision %q", keys), func(t *testing.T) {
|
||||
client := &fakeBackupS3{pages: []*s3.ListObjectsV2Output{backupListPage(keys, false, "")}}
|
||||
require.Error(t, backupImagesWithClient(context.Background(), client, "bucket", t.TempDir(), ""))
|
||||
require.Empty(t, client.getOrder)
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("existing file blocks planned directory", func(t *testing.T) {
|
||||
outDir := t.TempDir()
|
||||
require.NoError(t, os.MkdirAll(filepath.Join(outDir, "images"), 0o700))
|
||||
require.NoError(t, os.WriteFile(filepath.Join(outDir, "images", "a"), []byte("existing"), 0o600))
|
||||
client := &fakeBackupS3{pages: []*s3.ListObjectsV2Output{backupListPage([]string{"a/b"}, false, "")}}
|
||||
require.Error(t, backupImagesWithClient(context.Background(), client, "bucket", outDir, ""))
|
||||
require.Empty(t, client.getOrder)
|
||||
})
|
||||
|
||||
t.Run("existing directory blocks planned file", func(t *testing.T) {
|
||||
outDir := t.TempDir()
|
||||
require.NoError(t, os.MkdirAll(filepath.Join(outDir, "images", "a"), 0o700))
|
||||
client := &fakeBackupS3{pages: []*s3.ListObjectsV2Output{backupListPage([]string{"a"}, false, "")}}
|
||||
require.Error(t, backupImagesWithClient(context.Background(), client, "bucket", outDir, ""))
|
||||
require.Empty(t, client.getOrder)
|
||||
})
|
||||
}
|
||||
|
||||
func TestBackupRetriesInterruptedStreamsAndLeavesNoPart(t *testing.T) {
|
||||
t.Run("eventual success", func(t *testing.T) {
|
||||
key := "picture/retry.webp"
|
||||
body := []byte("complete-object-after-retries")
|
||||
client := &fakeBackupS3{
|
||||
pages: []*s3.ListObjectsV2Output{backupListPage([]string{key}, false, "")},
|
||||
objects: map[string]fakeBackupObject{
|
||||
key: {body: body, interruptedAttempts: 2, etag: "successful-get-etag"},
|
||||
},
|
||||
}
|
||||
outDir := t.TempDir()
|
||||
require.NoError(t, backupImagesWithClient(context.Background(), client, "bucket", outDir, ""))
|
||||
require.Equal(t, 3, client.getCalls[key])
|
||||
content, err := os.ReadFile(filepath.Join(outDir, "images", "picture", "retry.webp"))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, body, content)
|
||||
requireNoPartFiles(t, filepath.Join(outDir, "images", "picture"))
|
||||
})
|
||||
|
||||
t.Run("short body relative to GET content length", func(t *testing.T) {
|
||||
key := "short-body.webp"
|
||||
client := &fakeBackupS3{
|
||||
pages: []*s3.ListObjectsV2Output{backupListPage([]string{key}, false, "")},
|
||||
objects: map[string]fakeBackupObject{
|
||||
key: {body: []byte("complete-on-second-attempt"), wrongLengthAttempts: 1},
|
||||
},
|
||||
}
|
||||
outDir := t.TempDir()
|
||||
require.NoError(t, backupImagesWithClient(context.Background(), client, "bucket", outDir, ""))
|
||||
require.Equal(t, 2, client.getCalls[key])
|
||||
requireNoPartFiles(t, filepath.Join(outDir, "images"))
|
||||
})
|
||||
|
||||
t.Run("final failure", func(t *testing.T) {
|
||||
key := "retry.txt"
|
||||
client := &fakeBackupS3{
|
||||
pages: []*s3.ListObjectsV2Output{backupListPage([]string{key}, false, "")},
|
||||
objects: map[string]fakeBackupObject{
|
||||
key: {body: []byte("never-complete"), interruptedAttempts: 3},
|
||||
},
|
||||
}
|
||||
outDir := t.TempDir()
|
||||
err := backupImagesWithClient(context.Background(), client, "bucket", outDir, "")
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "after 3 attempts")
|
||||
require.Equal(t, 3, client.getCalls[key])
|
||||
require.NoFileExists(t, filepath.Join(outDir, "images", key))
|
||||
requireNoPartFiles(t, filepath.Join(outDir, "images"))
|
||||
require.NoFileExists(t, filepath.Join(outDir, "images-manifest.json"))
|
||||
})
|
||||
}
|
||||
|
||||
type capturedRestorePut struct {
|
||||
key string
|
||||
body []byte
|
||||
contentType string
|
||||
cacheControl string
|
||||
size int64
|
||||
}
|
||||
|
||||
type fakeRestoreStore struct {
|
||||
existing map[string]bool
|
||||
headErrors map[string]error
|
||||
putFailures map[string]error
|
||||
headCalls []string
|
||||
puts []capturedRestorePut
|
||||
}
|
||||
|
||||
func (f *fakeRestoreStore) HeadObject(_ context.Context, input *s3.HeadObjectInput, _ ...func(*s3.Options)) (*s3.HeadObjectOutput, error) {
|
||||
key := aws.ToString(input.Key)
|
||||
f.headCalls = append(f.headCalls, key)
|
||||
if err := f.headErrors[key]; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if f.existing[key] {
|
||||
return &s3.HeadObjectOutput{}, nil
|
||||
}
|
||||
return nil, &types.NotFound{}
|
||||
}
|
||||
|
||||
func (f *fakeRestoreStore) PutObject(_ context.Context, input *s3.PutObjectInput) error {
|
||||
body, err := io.ReadAll(input.Body)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
key := aws.ToString(input.Key)
|
||||
f.puts = append(f.puts, capturedRestorePut{
|
||||
key: key,
|
||||
body: body,
|
||||
contentType: aws.ToString(input.ContentType),
|
||||
cacheControl: aws.ToString(input.CacheControl),
|
||||
size: aws.ToInt64(input.ContentLength),
|
||||
})
|
||||
return f.putFailures[key]
|
||||
}
|
||||
|
||||
type restoreFixtureObject struct {
|
||||
key string
|
||||
body []byte
|
||||
storedBody []byte
|
||||
manifestSize *int64
|
||||
manifestSHA string
|
||||
contentType string
|
||||
cacheControl string
|
||||
}
|
||||
|
||||
func createRestoreFixture(t *testing.T, objects []restoreFixtureObject) string {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
manifest := make([]imageManifestEntry, 0, len(objects))
|
||||
for _, object := range objects {
|
||||
storedBody := object.storedBody
|
||||
if storedBody == nil {
|
||||
storedBody = object.body
|
||||
}
|
||||
path := filepath.Join(dir, "images", filepath.FromSlash(object.key))
|
||||
require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700))
|
||||
require.NoError(t, os.WriteFile(path, storedBody, 0o600))
|
||||
|
||||
size := int64(len(object.body))
|
||||
if object.manifestSize != nil {
|
||||
size = *object.manifestSize
|
||||
}
|
||||
digest := sha256.Sum256(object.body)
|
||||
hash := fmt.Sprintf("%x", digest)
|
||||
if object.manifestSHA != "" {
|
||||
hash = object.manifestSHA
|
||||
}
|
||||
manifest = append(manifest, imageManifestEntry{
|
||||
Key: object.key,
|
||||
Size: size,
|
||||
SHA256: hash,
|
||||
ETag: "source-get-etag",
|
||||
ContentType: object.contentType,
|
||||
CacheControl: object.cacheControl,
|
||||
})
|
||||
}
|
||||
data, err := json.MarshalIndent(manifest, "", " ")
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, os.WriteFile(filepath.Join(dir, "images-manifest.json"), data, 0o600))
|
||||
return dir
|
||||
}
|
||||
|
||||
func TestRestoreImagesPreservesKeysContentAndMetadataWithTargetPrefix(t *testing.T) {
|
||||
objects := []restoreFixtureObject{
|
||||
{key: "2026/11/04/root.webp", body: []byte("root-body"), contentType: "image/webp", cacheControl: "public, max-age=60"},
|
||||
{key: "picture/2026/11/04/pic.webp", body: []byte("picture-body"), contentType: "image/webp", cacheControl: "public, max-age=120"},
|
||||
{key: "static/site.css", body: []byte("body{}"), contentType: "text/css", cacheControl: "no-cache"},
|
||||
}
|
||||
input := createRestoreFixture(t, objects)
|
||||
store := &fakeRestoreStore{}
|
||||
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||
Input: input,
|
||||
Yes: true,
|
||||
TargetPrefix: "isolated-restore", // 不带斜杠也应规范化为恰好一个路径分隔符。
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, store.headCalls, len(objects))
|
||||
require.Len(t, store.puts, len(objects))
|
||||
|
||||
byKey := make(map[string]capturedRestorePut, len(store.puts))
|
||||
for _, put := range store.puts {
|
||||
byKey[put.key] = put
|
||||
}
|
||||
for _, object := range objects {
|
||||
target := "isolated-restore/" + object.key
|
||||
put := byKey[target]
|
||||
require.Equal(t, object.body, put.body)
|
||||
require.Equal(t, object.contentType, put.contentType)
|
||||
require.Equal(t, object.cacheControl, put.cacheControl)
|
||||
require.Equal(t, int64(len(object.body)), put.size)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreImagesDefaultRefusesOverwriteWithZeroPut(t *testing.T) {
|
||||
objects := []restoreFixtureObject{
|
||||
{key: "first.webp", body: []byte("first")},
|
||||
{key: "second.webp", body: []byte("second")},
|
||||
}
|
||||
store := &fakeRestoreStore{existing: map[string]bool{"second.webp": true}}
|
||||
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||
Input: createRestoreFixture(t, objects),
|
||||
Yes: true,
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "no objects were uploaded")
|
||||
require.Empty(t, store.puts)
|
||||
}
|
||||
|
||||
func TestRestoreImagesOverwriteAndDryRun(t *testing.T) {
|
||||
t.Run("overwrite", func(t *testing.T) {
|
||||
store := &fakeRestoreStore{existing: map[string]bool{"picture.webp": true}}
|
||||
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||
Input: createRestoreFixture(t, []restoreFixtureObject{{key: "picture.webp", body: []byte("replacement")}}),
|
||||
Yes: true,
|
||||
Overwrite: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, store.headCalls, "显式 overwrite 不需要默认防覆盖 Head")
|
||||
require.Len(t, store.puts, 1)
|
||||
require.Equal(t, []byte("replacement"), store.puts[0].body)
|
||||
})
|
||||
|
||||
t.Run("dry run without yes", func(t *testing.T) {
|
||||
store := &fakeRestoreStore{}
|
||||
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||
Input: createRestoreFixture(t, []restoreFixtureObject{{key: "picture.webp", body: []byte("dry-run")}}),
|
||||
DryRun: true,
|
||||
TargetPrefix: "drill/",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"drill/picture.webp"}, store.headCalls)
|
||||
require.Empty(t, store.puts)
|
||||
})
|
||||
|
||||
t.Run("actual write requires yes", func(t *testing.T) {
|
||||
store := &fakeRestoreStore{}
|
||||
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||
Input: createRestoreFixture(t, []restoreFixtureObject{{key: "picture.webp", body: []byte("write")}}),
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "--yes")
|
||||
require.Empty(t, store.headCalls)
|
||||
require.Empty(t, store.puts)
|
||||
})
|
||||
}
|
||||
|
||||
func TestRestoreRejectsTamperedFilesBeforeAnyRemoteCall(t *testing.T) {
|
||||
wrongSize := int64(999)
|
||||
tests := []struct {
|
||||
name string
|
||||
second restoreFixtureObject
|
||||
}{
|
||||
{
|
||||
name: "SHA256 tamper with unchanged size",
|
||||
second: restoreFixtureObject{
|
||||
key: "second.webp", body: []byte("second"), storedBody: []byte("SECOND"),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "size tamper",
|
||||
second: restoreFixtureObject{
|
||||
key: "second.webp", body: []byte("second"), manifestSize: &wrongSize,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "manifest hash tamper",
|
||||
second: restoreFixtureObject{
|
||||
key: "second.webp", body: []byte("second"), manifestSHA: strings.Repeat("0", 64),
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
store := &fakeRestoreStore{}
|
||||
input := createRestoreFixture(t, []restoreFixtureObject{
|
||||
{key: "first.webp", body: []byte("first")},
|
||||
tt.second,
|
||||
})
|
||||
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{Input: input, Yes: true})
|
||||
require.Error(t, err)
|
||||
require.Empty(t, store.headCalls, "全部本地文件应在任何远端操作前完成校验")
|
||||
require.Empty(t, store.puts, "首个有效文件也不能在后续文件校验失败前上传")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadImageManifestRejectsInvalidStructure(t *testing.T) {
|
||||
validHash := strings.Repeat("a", 64)
|
||||
tests := []struct {
|
||||
name string
|
||||
json string
|
||||
}{
|
||||
{name: "null", json: `null`},
|
||||
{name: "trailing JSON", json: `[] {}`},
|
||||
{name: "trailing garbage", json: `[] garbage`},
|
||||
{name: "duplicate key", json: fmt.Sprintf(`[{"key":"same","size":0,"sha256":%q},{"key":"same","size":0,"sha256":%q}]`, validHash, validHash)},
|
||||
{name: "short hash", json: `[{"key":"image","size":1,"sha256":"abc"}]`},
|
||||
{name: "non hex hash", json: fmt.Sprintf(`[{"key":"image","size":1,"sha256":%q}]`, strings.Repeat("z", 64))},
|
||||
{name: "negative size", json: fmt.Sprintf(`[{"key":"image","size":-1,"sha256":%q}]`, validHash)},
|
||||
{name: "path traversal", json: fmt.Sprintf(`[{"key":"../image","size":1,"sha256":%q}]`, validHash)},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "images-manifest.json")
|
||||
// 非数组用例直接保留原 JSON；其余用例同样应被解析层拒绝。
|
||||
require.NoError(t, os.WriteFile(path, []byte(tt.json), 0o600))
|
||||
_, err := readImageManifest(path)
|
||||
require.Error(t, err)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreRejectsSymlinkBeforeRemoteCalls(t *testing.T) {
|
||||
input := createRestoreFixture(t, []restoreFixtureObject{{key: "picture.webp", body: []byte("content")}})
|
||||
source := filepath.Join(input, "images", "picture.webp")
|
||||
outside := filepath.Join(t.TempDir(), "outside.webp")
|
||||
require.NoError(t, os.WriteFile(outside, []byte("content"), 0o600))
|
||||
require.NoError(t, os.Remove(source))
|
||||
if err := os.Symlink(outside, source); err != nil {
|
||||
t.Skipf("当前平台无法创建 symlink：%v", err)
|
||||
}
|
||||
|
||||
store := &fakeRestoreStore{}
|
||||
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{Input: input, DryRun: true})
|
||||
require.Error(t, err)
|
||||
require.Empty(t, store.headCalls)
|
||||
require.Empty(t, store.puts)
|
||||
}
|
||||
|
||||
func TestRestoreReportsPartialFailureWithoutClaimingTransaction(t *testing.T) {
|
||||
store := &fakeRestoreStore{putFailures: map[string]error{"second.webp": errors.New("injected put failure")}}
|
||||
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||
Input: createRestoreFixture(t, []restoreFixtureObject{
|
||||
{key: "first.webp", body: []byte("first")},
|
||||
{key: "second.webp", body: []byte("second")},
|
||||
}),
|
||||
Yes: true,
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "1 successful uploads")
|
||||
require.Contains(t, err.Error(), "not transactional")
|
||||
require.Len(t, store.puts, 2)
|
||||
}
|
||||
|
||||
func TestRestoreRejectsUnsafeTargetPrefixBeforeRemoteCalls(t *testing.T) {
|
||||
store := &fakeRestoreStore{}
|
||||
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||
Input: createRestoreFixture(t, []restoreFixtureObject{{key: "picture.webp", body: []byte("content")}}),
|
||||
Yes: true,
|
||||
TargetPrefix: "../outside/",
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Empty(t, store.headCalls)
|
||||
require.Empty(t, store.puts)
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"flag"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/urfave/cli/v2"
|
||||
)
|
||||
|
||||
func TestRunBackupRejectsMutuallyExclusiveModes(t *testing.T) {
|
||||
set := flag.NewFlagSet("backup", flag.ContinueOnError)
|
||||
set.Bool("db-only", true, "")
|
||||
set.Bool("images-only", true, "")
|
||||
|
||||
err := runBackup(cli.NewContext(nil, set, nil))
|
||||
require.EqualError(t, err, "--db-only and --images-only cannot be used together")
|
||||
}
|
||||
|
||||
func TestWriteFileAtomically(t *testing.T) {
|
||||
t.Run("success", func(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
target := filepath.Join(dir, "backup.sql")
|
||||
require.NoError(t, writeFileAtomically(target, func(w io.Writer) error {
|
||||
_, err := io.WriteString(w, "complete backup")
|
||||
return err
|
||||
}))
|
||||
|
||||
content, err := os.ReadFile(target)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "complete backup", string(content))
|
||||
requireNoPartFiles(t, dir)
|
||||
})
|
||||
|
||||
t.Run("failure preserves previous backup", func(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
target := filepath.Join(dir, "backup.sql")
|
||||
require.NoError(t, os.WriteFile(target, []byte("known good"), 0o600))
|
||||
writeErr := errors.New("dump failed")
|
||||
|
||||
err := writeFileAtomically(target, func(w io.Writer) error {
|
||||
_, _ = io.WriteString(w, "partial")
|
||||
return writeErr
|
||||
})
|
||||
require.ErrorIs(t, err, writeErr)
|
||||
|
||||
content, err := os.ReadFile(target)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "known good", string(content))
|
||||
requireNoPartFiles(t, dir)
|
||||
})
|
||||
|
||||
t.Run("failure leaves no new backup", func(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
target := filepath.Join(dir, "backup.sql")
|
||||
writeErr := errors.New("dump failed")
|
||||
|
||||
err := writeFileAtomically(target, func(w io.Writer) error {
|
||||
_, _ = io.WriteString(w, "partial")
|
||||
return writeErr
|
||||
})
|
||||
require.ErrorIs(t, err, writeErr)
|
||||
require.NoFileExists(t, target)
|
||||
requireNoPartFiles(t, dir)
|
||||
})
|
||||
}
|
||||
|
||||
func requireNoPartFiles(t *testing.T, dir string) {
|
||||
t.Helper()
|
||||
matches, err := filepath.Glob(filepath.Join(dir, "*.part"))
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, matches)
|
||||
}
|
||||
|
||||
func TestImageManifestEntryIntegrityFields(t *testing.T) {
|
||||
entry := imageManifestEntry{
|
||||
Key: "tamabox/image.webp",
|
||||
Size: 3,
|
||||
SHA256: "actual-content-sha256",
|
||||
ETag: "provider-etag",
|
||||
}
|
||||
encoded, err := json.Marshal(entry)
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, `{
|
||||
"key":"tamabox/image.webp",
|
||||
"size":3,
|
||||
"sha256":"actual-content-sha256",
|
||||
"etag":"provider-etag"
|
||||
}`, string(encoded))
|
||||
require.NotContains(t, string(encoded), "md5")
|
||||
}
|
||||
+3
-1
@@ -95,7 +95,9 @@ func runBan(ctx *cli.Context) error {
|
||||
return errors.Wrap(err, "ban user and cleanup")
|
||||
}
|
||||
|
||||
logrus.WithField("user_id", user.ID).Info("用户已封禁，关联数据已清理")
|
||||
// 数据库事务已经提交，先明确输出这一事实，再进行非事务性的 R2 清理。
|
||||
logrus.WithField("user_id", user.ID).Info("数据库封禁与关联数据清理已提交")
|
||||
drainR2OutboxAfterBan(ctx.Context)
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,365 @@
|
||||
// Copyright 2026 tamakyi. All rights reserved.
|
||||
// Use of this source code is governed by a MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/sirupsen/logrus"
|
||||
"github.com/urfave/cli/v2"
|
||||
|
||||
"github.com/tamakyi/TamaBox/internal/conf"
|
||||
"github.com/tamakyi/TamaBox/internal/db"
|
||||
"github.com/tamakyi/TamaBox/internal/storage"
|
||||
)
|
||||
|
||||
const (
|
||||
r2OutboxDefaultLimit = 100
|
||||
r2OutboxWorkerInterval = time.Minute
|
||||
r2OutboxInitialBackoff = time.Minute
|
||||
r2OutboxMaxBackoff = time.Hour
|
||||
r2OutboxMaxLastErrorSize = 1024
|
||||
)
|
||||
|
||||
var errR2OutboxNotConfigured = errors.New("R2 configuration is incomplete")
|
||||
|
||||
// R2Outbox 有界处理一批持久化的 R2 对象删除任务。
|
||||
var R2Outbox = &cli.Command{
|
||||
Name: "r2-outbox",
|
||||
Usage: "Drain a bounded batch of pending R2 object deletion tasks",
|
||||
Action: runR2Outbox,
|
||||
Flags: []cli.Flag{
|
||||
&cli.IntFlag{
|
||||
Name: "limit",
|
||||
Value: r2OutboxDefaultLimit,
|
||||
Usage: "Maximum number of due tasks to process",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
type r2OutboxTaskStore interface {
|
||||
ListDue(ctx context.Context, limit int) ([]*db.R2DeleteTask, error)
|
||||
Delete(id uint) error
|
||||
MarkFailed(id uint, nextRetry time.Time, lastError string) error
|
||||
HasReferences(ctx context.Context, bucket, key string) (bool, error)
|
||||
}
|
||||
|
||||
type r2ObjectDeleter interface {
|
||||
DeleteObject(ctx context.Context, bucket, key string) error
|
||||
}
|
||||
|
||||
type r2S3ObjectDeleter struct {
|
||||
client *s3.Client
|
||||
}
|
||||
|
||||
func (d *r2S3ObjectDeleter) DeleteObject(ctx context.Context, bucket, key string) error {
|
||||
return storage.DeleteObject(ctx, d.client, bucket, key)
|
||||
}
|
||||
|
||||
type r2OutboxResult struct {
|
||||
Listed int
|
||||
Removed int
|
||||
ObjectsDeleted int
|
||||
ReferenceOnly int
|
||||
Failed int
|
||||
Skipped int
|
||||
Remaining int
|
||||
}
|
||||
|
||||
// processR2Outbox 串行处理一批到期任务。单个任务失败不会中断同批后续任务；
|
||||
// context 取消则立即停止，且不会把关闭过程中的取消误记为一次任务失败。
|
||||
func processR2Outbox(ctx context.Context, store r2OutboxTaskStore, deleter r2ObjectDeleter, limit int) (result r2OutboxResult, err error) {
|
||||
defer func() {
|
||||
result.Remaining = result.Listed - result.Removed
|
||||
}()
|
||||
|
||||
tasks, err := store.ListDue(ctx, limit)
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return result, ctx.Err()
|
||||
}
|
||||
return result, fmt.Errorf("list R2 outbox tasks: %s", sanitizeR2OutboxError(err))
|
||||
}
|
||||
result.Listed = len(tasks)
|
||||
|
||||
for _, task := range tasks {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return result, err
|
||||
}
|
||||
if task == nil {
|
||||
result.Failed++
|
||||
logrus.Warn("R2 outbox 返回了空任务，已跳过")
|
||||
continue
|
||||
}
|
||||
|
||||
logger := logrus.WithField("task_id", task.ID)
|
||||
if task.Bucket != conf.Upload.ImageBucket {
|
||||
result.Skipped++
|
||||
logger.Warn("R2 outbox 任务 bucket 与当前配置不匹配，任务已保留")
|
||||
continue
|
||||
}
|
||||
|
||||
referenced, refErr := store.HasReferences(ctx, task.Bucket, task.Key)
|
||||
if ctx.Err() != nil {
|
||||
return result, ctx.Err()
|
||||
}
|
||||
if refErr != nil {
|
||||
result.Failed++
|
||||
markR2OutboxTaskFailed(store, task, refErr)
|
||||
continue
|
||||
}
|
||||
if referenced {
|
||||
// 还有引用时只清理过期 outbox 记录，绝不删除对象。最后一个引用被
|
||||
// 删除时，对应数据库事务会重新入队一条任务。
|
||||
if err := store.Delete(task.ID); err != nil {
|
||||
result.Failed++
|
||||
logger.WithField("error", sanitizeR2OutboxError(err)).Warn("R2 outbox 引用任务清理失败，任务已保留")
|
||||
continue
|
||||
}
|
||||
result.Removed++
|
||||
result.ReferenceOnly++
|
||||
continue
|
||||
}
|
||||
|
||||
deleteErr := deleter.DeleteObject(ctx, task.Bucket, task.Key)
|
||||
if deleteErr != nil && !isS3ObjectNotFound(deleteErr) {
|
||||
if ctx.Err() != nil {
|
||||
return result, ctx.Err()
|
||||
}
|
||||
result.Failed++
|
||||
markR2OutboxTaskFailed(store, task, deleteErr)
|
||||
continue
|
||||
}
|
||||
|
||||
// storage.DeleteObject 已将 404 归一为成功；这里再次识别 404，保证其他
|
||||
// deleter 实现也遵循相同的幂等语义。只有任务记录也删除后才算完成。
|
||||
if err := store.Delete(task.ID); err != nil {
|
||||
result.Failed++
|
||||
logger.WithField("error", sanitizeR2OutboxError(err)).Warn("R2 对象已删除，但 outbox 任务清理失败；后续将幂等重试")
|
||||
continue
|
||||
}
|
||||
result.Removed++
|
||||
result.ObjectsDeleted++
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func markR2OutboxTaskFailed(store r2OutboxTaskStore, task *db.R2DeleteTask, cause error) {
|
||||
nextRetry := time.Now().Add(r2OutboxBackoff(task.Attempts))
|
||||
lastError := sanitizeR2OutboxError(cause)
|
||||
if err := store.MarkFailed(task.ID, nextRetry, lastError); err != nil {
|
||||
logrus.WithFields(logrus.Fields{
|
||||
"task_id": task.ID,
|
||||
"processing_error": lastError,
|
||||
"persist_error": sanitizeR2OutboxError(err),
|
||||
}).Warn("R2 outbox 任务失败且重试状态写入失败，任务已保留")
|
||||
return
|
||||
}
|
||||
logrus.WithFields(logrus.Fields{
|
||||
"task_id": task.ID,
|
||||
"error": lastError,
|
||||
"next_retry": nextRetry,
|
||||
}).Warn("R2 outbox 任务处理失败，已安排重试")
|
||||
}
|
||||
|
||||
func r2OutboxBackoff(attempts int) time.Duration {
|
||||
if attempts < 0 {
|
||||
attempts = 0
|
||||
}
|
||||
// 1、2、4、8、16、32 分钟后，下一次本应到 64 分钟，直接封顶一小时。
|
||||
if attempts >= 6 {
|
||||
return r2OutboxMaxBackoff
|
||||
}
|
||||
return r2OutboxInitialBackoff << attempts
|
||||
}
|
||||
|
||||
// sanitizeR2OutboxError 在持久化或记录错误前移除已知配置与凭据，并限制为
|
||||
// 1024 字节。截断保持 UTF-8 边界，避免向 text 列写入半个多字节字符。
|
||||
func sanitizeR2OutboxError(err error) string {
|
||||
if err == nil {
|
||||
return ""
|
||||
}
|
||||
message := err.Error()
|
||||
sensitiveValues := []string{
|
||||
conf.Database.DSN,
|
||||
conf.Database.Password,
|
||||
conf.Upload.ImageEndpoint,
|
||||
conf.Upload.ImageAccessID,
|
||||
conf.Upload.ImageAccessSecret,
|
||||
conf.Upload.ImageBucket,
|
||||
conf.Upload.ImageBucketCDNHost,
|
||||
}
|
||||
for _, value := range sensitiveValues {
|
||||
if value != "" {
|
||||
message = strings.ReplaceAll(message, value, "[redacted]")
|
||||
}
|
||||
}
|
||||
if len(message) <= r2OutboxMaxLastErrorSize {
|
||||
return message
|
||||
}
|
||||
end := r2OutboxMaxLastErrorSize
|
||||
for end > 0 && !utf8.ValidString(message[:end]) {
|
||||
end--
|
||||
}
|
||||
return message[:end]
|
||||
}
|
||||
|
||||
func r2OutboxConfigured() bool {
|
||||
return strings.TrimSpace(conf.Upload.ImageEndpoint) != "" &&
|
||||
strings.TrimSpace(conf.Upload.ImageAccessID) != "" &&
|
||||
strings.TrimSpace(conf.Upload.ImageAccessSecret) != "" &&
|
||||
strings.TrimSpace(conf.Upload.ImageBucket) != ""
|
||||
}
|
||||
|
||||
func newR2OutboxObjectDeleter(ctx context.Context) (r2ObjectDeleter, error) {
|
||||
if !r2OutboxConfigured() {
|
||||
return nil, errR2OutboxNotConfigured
|
||||
}
|
||||
client, err := storage.NewS3Client(
|
||||
ctx,
|
||||
conf.Upload.ImageEndpoint,
|
||||
conf.Upload.ImageBucket,
|
||||
conf.Upload.ImageAccessID,
|
||||
conf.Upload.ImageAccessSecret,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("initialize R2 client: %s", sanitizeR2OutboxError(err))
|
||||
}
|
||||
return &r2S3ObjectDeleter{client: client}, nil
|
||||
}
|
||||
|
||||
// startR2OutboxWorker 只能在 db.Init 成功后调用。goroutine 由调用方提供的
|
||||
// WaitGroup 跟踪，Web 退出时先取消 ctx，再等待正在执行的一轮收尾。
|
||||
func startR2OutboxWorker(ctx context.Context, wg *sync.WaitGroup) {
|
||||
if !r2OutboxConfigured() {
|
||||
return
|
||||
}
|
||||
deleter, err := newR2OutboxObjectDeleter(ctx)
|
||||
if err != nil {
|
||||
logrus.WithField("error", sanitizeR2OutboxError(err)).Warn("R2 outbox worker 启动失败，Web 服务继续运行")
|
||||
return
|
||||
}
|
||||
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
runR2OutboxWorker(ctx, db.R2DeleteTasks, deleter)
|
||||
}()
|
||||
}
|
||||
|
||||
func runR2OutboxWorker(ctx context.Context, store r2OutboxTaskStore, deleter r2ObjectDeleter) {
|
||||
ticker := time.NewTicker(r2OutboxWorkerInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
// 启动先立即跑一轮；处理期间不启动并发轮次。若一轮超过一分钟，
|
||||
// ticker 只会合并唤醒，不会形成多个 worker 同时消费。
|
||||
_, err := processR2Outbox(ctx, store, deleter, r2OutboxDefaultLimit)
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
logrus.WithField("error", sanitizeR2OutboxError(err)).Warn("R2 outbox 本轮处理失败，将在下一分钟重试")
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func drainR2Outbox(ctx context.Context, limit int) (r2OutboxResult, error) {
|
||||
deleter, err := newR2OutboxObjectDeleter(ctx)
|
||||
if err != nil {
|
||||
return r2OutboxResult{}, err
|
||||
}
|
||||
return processR2Outbox(ctx, db.R2DeleteTasks, deleter, limit)
|
||||
}
|
||||
|
||||
func runR2Outbox(ctx *cli.Context) error {
|
||||
limit := ctx.Int("limit")
|
||||
if limit <= 0 {
|
||||
return errors.New("--limit must be greater than zero")
|
||||
}
|
||||
if err := conf.Init(); err != nil {
|
||||
return fmt.Errorf("load configuration: %w", err)
|
||||
}
|
||||
if !r2OutboxConfigured() {
|
||||
return errR2OutboxNotConfigured
|
||||
}
|
||||
|
||||
dbType := conf.Database.Type
|
||||
var dsn string
|
||||
switch dbType {
|
||||
case "mysql", "":
|
||||
dsn = conf.MySQLDsn()
|
||||
case "postgres":
|
||||
dsn = conf.PostgresDsn()
|
||||
default:
|
||||
return fmt.Errorf("unknown database type: %q", dbType)
|
||||
}
|
||||
conf.Database.DSN = dsn
|
||||
gdb, err := db.Init(dbType, dsn)
|
||||
if err != nil {
|
||||
return fmt.Errorf("connect to database: %w", err)
|
||||
}
|
||||
if sqlDB, dbErr := gdb.DB(); dbErr == nil {
|
||||
defer func() {
|
||||
if err := sqlDB.Close(); err != nil {
|
||||
logrus.WithField("error", sanitizeR2OutboxError(err)).Warn("关闭数据库连接池失败")
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
result, err := drainR2Outbox(ctx.Context, limit)
|
||||
if err != nil {
|
||||
return fmt.Errorf("drain R2 outbox: %s", sanitizeR2OutboxError(err))
|
||||
}
|
||||
logR2OutboxResult("R2 outbox 本轮处理完成", result)
|
||||
return nil
|
||||
}
|
||||
|
||||
func logR2OutboxResult(message string, result r2OutboxResult) {
|
||||
entry := logrus.WithFields(logrus.Fields{
|
||||
"listed": result.Listed,
|
||||
"removed": result.Removed,
|
||||
"object_deletes": result.ObjectsDeleted,
|
||||
"reference_only": result.ReferenceOnly,
|
||||
"failed": result.Failed,
|
||||
"skipped": result.Skipped,
|
||||
"remaining": result.Remaining,
|
||||
})
|
||||
if result.Remaining > 0 {
|
||||
entry.Warn(message + "，仍有任务待重试或保留")
|
||||
return
|
||||
}
|
||||
entry.Info(message)
|
||||
}
|
||||
|
||||
// drainR2OutboxAfterBan 在封禁事务提交后执行。R2 失败只影响对象清理进度，
|
||||
// 不向 runBan 返回错误，避免运维误以为需要重新执行不可逆的数据库封禁。
|
||||
func drainR2OutboxAfterBan(ctx context.Context) {
|
||||
result, err := drainR2Outbox(ctx, r2OutboxDefaultLimit)
|
||||
if err != nil {
|
||||
logrus.WithField("error", sanitizeR2OutboxError(err)).Warn(
|
||||
"数据库封禁与清理已提交，但 R2 outbox 本轮未能处理；任务仍保留，无需重跑 ban，可稍后运行 r2-outbox",
|
||||
)
|
||||
return
|
||||
}
|
||||
if result.Remaining > 0 {
|
||||
logR2OutboxResult("数据库封禁与清理已提交；R2 对象清理将在后续重试，无需重跑 ban", result)
|
||||
return
|
||||
}
|
||||
logR2OutboxResult("数据库封禁与清理已提交；R2 outbox 本轮处理完成", result)
|
||||
}
|
||||
@@ -0,0 +1,270 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/tamakyi/TamaBox/internal/conf"
|
||||
"github.com/tamakyi/TamaBox/internal/db"
|
||||
)
|
||||
|
||||
type fakeR2OutboxFailure struct {
|
||||
nextRetry time.Time
|
||||
lastError string
|
||||
}
|
||||
|
||||
type fakeR2OutboxStore struct {
|
||||
tasks []*db.R2DeleteTask
|
||||
references map[string]bool
|
||||
referenceErrs map[string]error
|
||||
referenceCalls []string
|
||||
deleted map[uint]bool
|
||||
deleteErr error
|
||||
failures map[uint]fakeR2OutboxFailure
|
||||
markFailedErr error
|
||||
}
|
||||
|
||||
func newFakeR2OutboxStore(tasks ...*db.R2DeleteTask) *fakeR2OutboxStore {
|
||||
return &fakeR2OutboxStore{
|
||||
tasks: tasks,
|
||||
references: make(map[string]bool),
|
||||
referenceErrs: make(map[string]error),
|
||||
deleted: make(map[uint]bool),
|
||||
failures: make(map[uint]fakeR2OutboxFailure),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *fakeR2OutboxStore) ListDue(ctx context.Context, limit int) ([]*db.R2DeleteTask, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tasks := s.tasks
|
||||
if limit >= 0 && len(tasks) > limit {
|
||||
tasks = tasks[:limit]
|
||||
}
|
||||
return append([]*db.R2DeleteTask(nil), tasks...), nil
|
||||
}
|
||||
|
||||
func (s *fakeR2OutboxStore) Delete(id uint) error {
|
||||
if s.deleteErr != nil {
|
||||
return s.deleteErr
|
||||
}
|
||||
s.deleted[id] = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *fakeR2OutboxStore) MarkFailed(id uint, nextRetry time.Time, lastError string) error {
|
||||
if s.markFailedErr != nil {
|
||||
return s.markFailedErr
|
||||
}
|
||||
s.failures[id] = fakeR2OutboxFailure{nextRetry: nextRetry, lastError: lastError}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *fakeR2OutboxStore) HasReferences(ctx context.Context, _, key string) (bool, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
s.referenceCalls = append(s.referenceCalls, key)
|
||||
return s.references[key], s.referenceErrs[key]
|
||||
}
|
||||
|
||||
type fakeR2DeleteCall struct {
|
||||
bucket string
|
||||
key string
|
||||
}
|
||||
|
||||
type fakeR2ObjectDeleter struct {
|
||||
calls []fakeR2DeleteCall
|
||||
errs map[string]error
|
||||
onDelete func(context.Context, string, string) error
|
||||
}
|
||||
|
||||
func (d *fakeR2ObjectDeleter) DeleteObject(ctx context.Context, bucket, key string) error {
|
||||
d.calls = append(d.calls, fakeR2DeleteCall{bucket: bucket, key: key})
|
||||
if d.onDelete != nil {
|
||||
return d.onDelete(ctx, bucket, key)
|
||||
}
|
||||
return d.errs[key]
|
||||
}
|
||||
|
||||
type fakeR2HTTPStatusError int
|
||||
|
||||
func (e fakeR2HTTPStatusError) Error() string {
|
||||
return fmt.Sprintf("R2 returned HTTP %d", int(e))
|
||||
}
|
||||
|
||||
func (e fakeR2HTTPStatusError) HTTPStatusCode() int {
|
||||
return int(e)
|
||||
}
|
||||
|
||||
func setR2OutboxTestConfig(t *testing.T) {
|
||||
t.Helper()
|
||||
old := conf.Upload
|
||||
conf.Upload.ImageEndpoint = "https://r2.example.test"
|
||||
conf.Upload.ImageAccessID = "test-access-id"
|
||||
conf.Upload.ImageAccessSecret = "test-secret"
|
||||
conf.Upload.ImageBucket = "configured-bucket"
|
||||
conf.Upload.ImageBucketCDNHost = "cdn.example.test"
|
||||
t.Cleanup(func() {
|
||||
conf.Upload = old
|
||||
})
|
||||
}
|
||||
|
||||
func TestProcessR2OutboxSuccessAndNotFound(t *testing.T) {
|
||||
setR2OutboxTestConfig(t)
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
deleteErr error
|
||||
}{
|
||||
{name: "success"},
|
||||
{name: "not found is success", deleteErr: fakeR2HTTPStatusError(404)},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
first := &db.R2DeleteTask{ID: 1, Bucket: conf.Upload.ImageBucket, Key: "first.webp"}
|
||||
second := &db.R2DeleteTask{ID: 2, Bucket: conf.Upload.ImageBucket, Key: "second.webp"}
|
||||
store := newFakeR2OutboxStore(first, second)
|
||||
deleter := &fakeR2ObjectDeleter{errs: map[string]error{"second.webp": tt.deleteErr}}
|
||||
|
||||
result, err := processR2Outbox(context.Background(), store, deleter, 100)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, result.Listed)
|
||||
require.Equal(t, 2, result.Removed)
|
||||
require.Equal(t, 2, result.ObjectsDeleted)
|
||||
require.Zero(t, result.Failed)
|
||||
require.Zero(t, result.Remaining)
|
||||
require.Equal(t, []fakeR2DeleteCall{
|
||||
{bucket: conf.Upload.ImageBucket, key: "first.webp"},
|
||||
{bucket: conf.Upload.ImageBucket, key: "second.webp"},
|
||||
}, deleter.calls, "任务必须按 ListDue 顺序串行处理")
|
||||
require.True(t, store.deleted[first.ID])
|
||||
require.True(t, store.deleted[second.ID])
|
||||
require.Empty(t, store.failures)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessR2OutboxReferencedTaskDeletesNoObject(t *testing.T) {
|
||||
setR2OutboxTestConfig(t)
|
||||
task := &db.R2DeleteTask{ID: 11, Bucket: conf.Upload.ImageBucket, Key: "shared.webp"}
|
||||
store := newFakeR2OutboxStore(task)
|
||||
store.references[task.Key] = true
|
||||
deleter := &fakeR2ObjectDeleter{}
|
||||
|
||||
result, err := processR2Outbox(context.Background(), store, deleter, 100)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, deleter.calls, "仍有引用时绝不能调用 R2 删除")
|
||||
require.True(t, store.deleted[task.ID], "过期任务本身应删除，最后引用移除时会重新入队")
|
||||
require.Equal(t, 1, result.ReferenceOnly)
|
||||
require.Equal(t, 1, result.Removed)
|
||||
require.Zero(t, result.ObjectsDeleted)
|
||||
require.Zero(t, result.Remaining)
|
||||
}
|
||||
|
||||
func TestProcessR2OutboxReferenceErrorSchedulesRetry(t *testing.T) {
|
||||
setR2OutboxTestConfig(t)
|
||||
task := &db.R2DeleteTask{ID: 21, Bucket: conf.Upload.ImageBucket, Key: "unknown.webp"}
|
||||
store := newFakeR2OutboxStore(task)
|
||||
store.referenceErrs[task.Key] = errors.New("reference query failed")
|
||||
deleter := &fakeR2ObjectDeleter{}
|
||||
before := time.Now()
|
||||
|
||||
result, err := processR2Outbox(context.Background(), store, deleter, 100)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, deleter.calls)
|
||||
require.False(t, store.deleted[task.ID])
|
||||
require.Contains(t, store.failures, task.ID)
|
||||
failure := store.failures[task.ID]
|
||||
require.Equal(t, "reference query failed", failure.lastError)
|
||||
require.WithinDuration(t, before.Add(time.Minute), failure.nextRetry, 2*time.Second)
|
||||
require.Equal(t, 1, result.Failed)
|
||||
require.Equal(t, 1, result.Remaining)
|
||||
}
|
||||
|
||||
func TestProcessR2OutboxDeleteFailureUsesBackoffAndBoundsLastError(t *testing.T) {
|
||||
setR2OutboxTestConfig(t)
|
||||
task := &db.R2DeleteTask{ID: 31, Bucket: conf.Upload.ImageBucket, Key: "retry.webp", Attempts: 2}
|
||||
store := newFakeR2OutboxStore(task)
|
||||
longErr := fmt.Errorf("delete failed endpoint=%s credential=%s %s",
|
||||
conf.Upload.ImageEndpoint,
|
||||
conf.Upload.ImageAccessSecret,
|
||||
strings.Repeat("x", 1100),
|
||||
)
|
||||
deleter := &fakeR2ObjectDeleter{errs: map[string]error{task.Key: longErr}}
|
||||
before := time.Now()
|
||||
|
||||
result, err := processR2Outbox(context.Background(), store, deleter, 100)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, deleter.calls, 1)
|
||||
require.False(t, store.deleted[task.ID])
|
||||
require.Contains(t, store.failures, task.ID)
|
||||
failure := store.failures[task.ID]
|
||||
require.WithinDuration(t, before.Add(4*time.Minute), failure.nextRetry, 2*time.Second)
|
||||
require.LessOrEqual(t, len(failure.lastError), r2OutboxMaxLastErrorSize)
|
||||
require.NotContains(t, failure.lastError, conf.Upload.ImageEndpoint)
|
||||
require.NotContains(t, failure.lastError, conf.Upload.ImageAccessSecret)
|
||||
require.Equal(t, 1, result.Failed)
|
||||
require.Equal(t, 1, result.Remaining)
|
||||
}
|
||||
|
||||
func TestProcessR2OutboxBucketMismatchIsRetainedWithoutDeletion(t *testing.T) {
|
||||
setR2OutboxTestConfig(t)
|
||||
task := &db.R2DeleteTask{ID: 41, Bucket: "different-bucket", Key: "foreign.webp"}
|
||||
store := newFakeR2OutboxStore(task)
|
||||
deleter := &fakeR2ObjectDeleter{}
|
||||
|
||||
result, err := processR2Outbox(context.Background(), store, deleter, 100)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, store.referenceCalls, "bucket 不匹配时不应查询当前 bucket 的引用")
|
||||
require.Empty(t, deleter.calls)
|
||||
require.Empty(t, store.failures)
|
||||
require.False(t, store.deleted[task.ID])
|
||||
require.Equal(t, 1, result.Skipped)
|
||||
require.Zero(t, result.Failed)
|
||||
require.Equal(t, 1, result.Remaining)
|
||||
}
|
||||
|
||||
func TestProcessR2OutboxContextCancellationStopsWithoutMarkFailed(t *testing.T) {
|
||||
setR2OutboxTestConfig(t)
|
||||
first := &db.R2DeleteTask{ID: 51, Bucket: conf.Upload.ImageBucket, Key: "first.webp"}
|
||||
second := &db.R2DeleteTask{ID: 52, Bucket: conf.Upload.ImageBucket, Key: "second.webp"}
|
||||
store := newFakeR2OutboxStore(first, second)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
deleter := &fakeR2ObjectDeleter{
|
||||
onDelete: func(ctx context.Context, _, _ string) error {
|
||||
cancel()
|
||||
return ctx.Err()
|
||||
},
|
||||
}
|
||||
|
||||
result, err := processR2Outbox(ctx, store, deleter, 100)
|
||||
require.ErrorIs(t, err, context.Canceled)
|
||||
require.Len(t, deleter.calls, 1, "取消后不得继续处理第二个任务")
|
||||
require.Empty(t, store.failures, "优雅关闭导致的取消不应增加任务失败次数")
|
||||
require.Empty(t, store.deleted)
|
||||
require.Zero(t, result.Failed)
|
||||
require.Equal(t, 2, result.Remaining)
|
||||
}
|
||||
|
||||
func TestR2OutboxBackoffCapsAtOneHour(t *testing.T) {
|
||||
for _, tt := range []struct {
|
||||
attempts int
|
||||
expected time.Duration
|
||||
}{
|
||||
{attempts: -1, expected: time.Minute},
|
||||
{attempts: 0, expected: time.Minute},
|
||||
{attempts: 1, expected: 2 * time.Minute},
|
||||
{attempts: 2, expected: 4 * time.Minute},
|
||||
{attempts: 5, expected: 32 * time.Minute},
|
||||
{attempts: 6, expected: time.Hour},
|
||||
{attempts: 100, expected: time.Hour},
|
||||
} {
|
||||
require.Equal(t, tt.expected, r2OutboxBackoff(tt.attempts))
|
||||
}
|
||||
}
|
||||
+10
-1
@@ -10,6 +10,7 @@ import (
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
@@ -209,7 +210,15 @@ func runWeb(ctx *cli.Context) error {
|
||||
// 「提问已 INSERT、图片还没传完」「回答已写入、配额还没扣」这类半成品。
|
||||
// systemd Restart=always 与 docker restart 都救不了这个。
|
||||
runCtx, stop := signal.NotifyContext(ctx.Context, syscall.SIGTERM, os.Interrupt)
|
||||
defer stop()
|
||||
var r2WorkerWG sync.WaitGroup
|
||||
defer func() {
|
||||
// 先取消 worker 的 context，再等待当前一轮退出，避免无人等待的后台删除。
|
||||
stop()
|
||||
r2WorkerWG.Wait()
|
||||
}()
|
||||
|
||||
// db.Init 已成功；R2 未完整配置时此函数直接返回，不影响 Web 启动。
|
||||
startR2OutboxWorker(runCtx, &r2WorkerWG)
|
||||
|
||||
serveErr := make(chan error, 1)
|
||||
go func() {
|
||||
|
||||
Reference in New Issue
Block a user