From 937e688cfa76f5435449f54583b3834395804c44 Mon Sep 17 00:00:00 2001 From: tamakyi Date: Thu, 5 Nov 2026 08:41:55 +0800 Subject: [PATCH] =?UTF-8?q?=E5=8F=AF=E9=9D=A0=E6=80=A7:=20=E5=AE=8C?= =?UTF-8?q?=E5=96=84=20R2=20=E5=A4=87=E4=BB=BD=E6=81=A2=E5=A4=8D=E4=B8=8E?= =?UTF-8?q?=E5=AF=B9=E8=B1=A1=E7=94=9F=E5=91=BD=E5=91=A8=E6=9C=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cmd/tamabox.go | 2 + docs/cli.md | 61 +- internal/cmd/backup.go | 555 +++++++++++----- internal/cmd/backup_restore.go | 437 +++++++++++++ .../cmd/backup_restore_integration_test.go | 184 ++++++ internal/cmd/backup_restore_test.go | 603 ++++++++++++++++++ internal/cmd/backup_test.go | 96 +++ internal/cmd/ban.go | 4 +- internal/cmd/r2_outbox.go | 365 +++++++++++ internal/cmd/r2_outbox_test.go | 270 ++++++++ internal/cmd/web.go | 11 +- internal/db/db.go | 3 +- internal/db/message_image.go | 3 + internal/db/message_image_test.go | 39 ++ internal/db/r2_delete_task.go | 277 ++++++++ internal/db/r2_delete_task_test.go | 141 ++++ internal/db/upload_images.go | 76 ++- internal/db/upload_images_outbox_test.go | 95 +++ internal/db/users.go | 89 ++- internal/db/users_outbox_test.go | 142 +++++ internal/route/chat/chat_handler.go | 33 +- internal/storage/cdn_url.go | 128 ++++ internal/storage/cdn_url_test.go | 112 ++++ internal/storage/client.go | 31 +- internal/storage/s3.go | 25 +- internal/storage/s3_client.go | 112 ++++ internal/storage/s3_client_test.go | 126 ++++ internal/storage/s3_retry.go | 167 +++++ internal/storage/s3_retry_test.go | 171 +++++ route/question/page.go | 5 +- route/user/profile.go | 17 + tamabox.go | 2 + 32 files changed, 4151 insertions(+), 231 deletions(-) create mode 100644 internal/cmd/backup_restore.go create mode 100644 internal/cmd/backup_restore_integration_test.go create mode 100644 internal/cmd/backup_restore_test.go create mode 100644 internal/cmd/backup_test.go create mode 100644 internal/cmd/r2_outbox.go create mode 100644 internal/cmd/r2_outbox_test.go create mode 100644 internal/db/message_image_test.go create mode 100644 internal/db/r2_delete_task.go create mode 100644 internal/db/r2_delete_task_test.go create mode 100644 internal/db/upload_images_outbox_test.go create mode 100644 internal/db/users_outbox_test.go create mode 100644 internal/storage/cdn_url.go create mode 100644 internal/storage/cdn_url_test.go create mode 100644 internal/storage/s3_client.go create mode 100644 internal/storage/s3_client_test.go create mode 100644 internal/storage/s3_retry.go create mode 100644 internal/storage/s3_retry_test.go diff --git a/cmd/tamabox.go b/cmd/tamabox.go index 7e55ae2..3880010 100644 --- a/cmd/tamabox.go +++ b/cmd/tamabox.go @@ -30,6 +30,8 @@ func main() { cmd.HidePublic, cmd.Honeypot, cmd.Backup, + cmd.RestoreImages, + cmd.R2Outbox, } if err := app.Run(os.Args); err != nil { logrus.WithError(err).Fatal("Failed to start application") diff --git a/docs/cli.md b/docs/cli.md index 6bd7c73..3e2b08f 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -32,7 +32,7 @@ ps -eo pid,lstart,cmd | grep '[t]amabox' # 看进程启动时间，与二进 ```bash ./tamabox --help | sed -n '/COMMANDS/,/GLOBAL/p' -# web / censor / uid / ban / restore / reports / restrict / unrestrict / hide-public / honeypot +# web / censor / uid / ban / restore / reports / restrict / unrestrict / hide-public / honeypot / backup / restore-images ``` > 想确认某个字符串有没有编进二进制，**不要用 `strings`** —— 它默认只输出可打印 ASCII， @@ -77,10 +77,13 @@ pg_dump -Fc > before-$(date +%F-%H%M).dump | `censor` | 重新审核历史内容 | 是 | 审核规则变更后复检 | | `uid` | 为用户生成/补齐 uid | 是 | 数据修复 | | `honeypot` | 扫描器诱捕取证（假登录页） | 部分是（`prune` / `purge-creds`） | 看谁在扫、拿到了什么字典 | +| `backup` | 备份 PostgreSQL 与配置 bucket 的全部 R2 对象 | 写本地备份目录 | 灾难恢复前留档 | +| `restore-images` | 校验并恢复 R2 图片及元数据 | 是（除 `--dry-run`） | 图片灾难恢复、隔离恢复演练 | | `web` | 启动 Web 服务 | — | 生产入口 | -除 `reports`、`honeypot hits/stats/paths/size` 外，所有写操作都要求交互式输入大写 `YES` 确认； +除 `reports`、`honeypot hits/stats/paths/size` 外，既有运营写操作要求交互式输入大写 `YES` 确认； 脚本化时加 `--yes`。注意 `honeypot prune` / `purge-creds` 的确认词是 **`yes`**（小写）。 +`restore-images` 不做交互确认：实际写入必须显式提供 `--yes`，`--dry-run` 除外。 --- @@ -329,7 +332,55 @@ pg_dump -Fc > before-$(date +%F-%H%M).dump --- -## 8. 典型处置流程 +## 8. `backup` / `restore-images` —— 数据库与 R2 图片 + +### 8.1 备份 + +```bash +# 数据库 + 配置 bucket 的全部对象 +./tamabox backup --output /var/backups/tamabox + +# 只备份对象，或只备份数据库 +./tamabox backup --images-only --output /var/backups/tamabox +./tamabox backup --db-only --output /var/backups/tamabox +``` + +图片备份使用**空前缀**分页列举配置 bucket 的全部对象，不再只查 `tamabox/`。根日期 key、 +`picture/`、`tamabox/` 和静态资源都会按原 key 保存到 `/images/`。 +备份会先完成全部分页 List、key 与本地路径预检，再开始下载；非法路径、目录 key、文件/目录碰撞会中止备份。 + +`images-manifest.json` 记录实际下载字节的 `size`、`sha256`，以及 GET 返回的 `etag`、 +`content_type`、`cache_control`。对象与 manifest 都采用临时文件原子落盘；下载流中断会整对象重试， +最终失败不会留下 `.part`。空 bucket 的 manifest 是 `[]`。 + +### 8.2 恢复与隔离演练 + +```bash +# 校验全部本地文件并检查默认防覆盖，不写远端 +./tamabox restore-images --input /var/backups/tamabox --dry-run + +# 实际恢复到原 key；必须显式 --yes +./tamabox restore-images --input /var/backups/tamabox --yes + +# 明确允许替换已存在对象 +./tamabox restore-images --input /var/backups/tamabox --yes --overwrite + +# 在原完整 key 前加隔离前缀，不覆盖生产 key +./tamabox restore-images --input /var/backups/tamabox --yes \ + --target-prefix tamabox-backup-test/manual-restore/ +``` + +`--input` 可指向备份目录，也可直接指向 `images-manifest.json`。恢复会在任何远端操作前校验全部 +manifest 条目和本地文件（普通文件、非 symlink、大小、实际 SHA-256）。默认先 Head 目标；任一目标已存在 +即失败且零 Put，只有 `--overwrite` 才允许替换。`--dry-run` 不写对象，实际写入必须 `--yes`。 + +`--target-prefix` 会加在 manifest 的完整 key 前；非空前缀末尾没有 `/` 时会自动补一个。 +Put 保留内容、原 key 路径（额外添加指定前缀）、`ContentType` 与 `CacheControl`，并使用单对象超时和有限重试。 +**恢复不是事务**：中途失败会明确报告已成功数量，已上传对象不会自动回滚。 + +--- + +## 9. 典型处置流程 1. **发现**：收到告警邮件（同一箱主 24h ≥ 8 条 或 7d ≥ 20 条）→ 或箱主后台出现「被举报 N 次」。 2. **取证**：`reports --domain ` 看条目、占比、历史处置。 @@ -351,7 +402,7 @@ pg_dump -Fc > before-$(date +%F-%H%M).dump --- -## 9. 排错 +## 10. 排错 | 现象 | 原因 | 处理 | |---|---|---| @@ -374,7 +425,7 @@ pg_dump -Fc > before-$(date +%F-%H%M).dump --- -## 10. 明确不做的事 +## 11. 明确不做的事 - **不做自动处置**：举报数达标只发邮件，不自动封禁。这个场景里举报本身就是攻击武器。 - **不做申诉后台**：申诉走邮件，处置结果由 CLI 写入。 diff --git a/internal/cmd/backup.go b/internal/cmd/backup.go index 96008ca..86fb5a5 100644 --- a/internal/cmd/backup.go +++ b/internal/cmd/backup.go @@ -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/（新）或 tamabox/YYYY/MM/.（历史）的形式存在 -// R2 上，数据库里只存 key（upload_images.key）或 URL（users.avatar/background 里的完整 -// CDN 地址）。只备份数据库的话，恢复后这些引用指向的图片若已不在 R2（bucket 误删 / -// 账号丢失 / 图片被清理），页面就全是裂图。因此备份必须把 R2 上的图片一并落盘。 +// 都以根日期 key、picture/ 或 tamabox/ 等形式存在 R2 上，数据库里只存 key 或 CDN URL。 +// 只备份数据库的话，恢复后这些引用指向的图片若已不在 R2，页面就全是裂图。 // // 输出目录结构： // // / // db-<20060102-150405>.sql # pg_dump 全量导出（schema + data，可用 psql 恢复） // images/ # 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/.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://.r2.cloudflarestorage.com/）， -// 它对 GetObject/PutObject 恰好可用（R2 容错），但 ListObjectsV2 会把 URL 拼成 -// //?list-type=2，R2 找不到这个 key 前缀而返回 404 NoSuchKey。 -// -// 这里剥离 endpoint 的 bucket 后缀，得到纯 account endpoint，并启用 path-style： -// List/Get 都用 Bucket 参数 + 完整 key（List 返回的 key 自带 bucket 名作前缀， -// 如 tamabox/2025/12/.jpg），这是 R2 上唯一稳定可用的姿势。 -func backupS3Client() (*s3.Client, error) { - endpoint := strings.TrimRight(conf.Upload.ImageEndpoint, "/") - // 剥离最后一个路径段（bucket 名）。endpoint 形如 - // https://.r2.cloudflarestorage.com/，去掉 /。 - 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 参数和原始完整 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< 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) +} diff --git a/internal/cmd/backup_restore_integration_test.go b/internal/cmd/backup_restore_integration_test.go new file mode 100644 index 0000000..210af44 --- /dev/null +++ b/internal/cmd/backup_restore_integration_test.go @@ -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 下对象数不应触发分页") +} diff --git a/internal/cmd/backup_restore_test.go b/internal/cmd/backup_restore_test.go new file mode 100644 index 0000000..795371a --- /dev/null +++ b/internal/cmd/backup_restore_test.go @@ -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) +} diff --git a/internal/cmd/backup_test.go b/internal/cmd/backup_test.go new file mode 100644 index 0000000..f9f7902 --- /dev/null +++ b/internal/cmd/backup_test.go @@ -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") +} diff --git a/internal/cmd/ban.go b/internal/cmd/ban.go index e59f87a..50507a0 100644 --- a/internal/cmd/ban.go +++ b/internal/cmd/ban.go @@ -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 } diff --git a/internal/cmd/r2_outbox.go b/internal/cmd/r2_outbox.go new file mode 100644 index 0000000..babab2c --- /dev/null +++ b/internal/cmd/r2_outbox.go @@ -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) +} diff --git a/internal/cmd/r2_outbox_test.go b/internal/cmd/r2_outbox_test.go new file mode 100644 index 0000000..a980ef6 --- /dev/null +++ b/internal/cmd/r2_outbox_test.go @@ -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)) + } +} diff --git a/internal/cmd/web.go b/internal/cmd/web.go index 3074e3e..04deef0 100644 --- a/internal/cmd/web.go +++ b/internal/cmd/web.go @@ -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() { diff --git a/internal/db/db.go b/internal/db/db.go index caf963e..84ebb5b 100644 --- a/internal/db/db.go +++ b/internal/db/db.go @@ -36,7 +36,7 @@ func Ping(ctx context.Context) error { var AllTables = []interface{}{ &User{}, &Question{}, &FollowUp{}, &CensorLog{}, &UploadImage{}, &UploadImageQuestion{}, - &Conversation{}, &Message{}, &MessageImage{}, &UserBlock{}, + &Conversation{}, &Message{}, &MessageImage{}, &UserBlock{}, &R2DeleteTask{}, // 匿名投票 &Vote{}, &VoteOption{}, &VoteFormField{}, &VoteSubmission{}, // 功能开关（提问 / 私聊 / 投票 的开关、定时、定量收集） @@ -132,6 +132,7 @@ func Init(typ, dsn string) (*gorm.DB, error) { Conversations = NewConversationStore(db) Messages = NewMessageStore(db) MessageImages = NewMessageImageStore(db) + R2DeleteTasks = NewR2DeleteTaskStore(db) Votes = NewVotesStore(db) VoteOptions = NewVoteOptionsStore(db) VoteFormFields = NewVoteFormFieldsStore(db) diff --git a/internal/db/message_image.go b/internal/db/message_image.go index 9ad7c05..b8c29fc 100644 --- a/internal/db/message_image.go +++ b/internal/db/message_image.go @@ -26,6 +26,7 @@ type MessageImage struct { gorm.Model MessageID uint `gorm:"index" json:"message_id"` URL string `gorm:"size:500" json:"url"` + ObjectKey string `gorm:"type:text" json:"-"` FileName string `gorm:"size:200" json:"file_name"` FileSize int64 `json:"file_size"` } @@ -33,6 +34,7 @@ type MessageImage struct { type CreateMessageImageOptions struct { MessageID uint URL string + ObjectKey string FileName string FileSize int64 } @@ -41,6 +43,7 @@ func (db *messageImages) Create(ctx context.Context, opts CreateMessageImageOpti img := &MessageImage{ MessageID: opts.MessageID, URL: opts.URL, + ObjectKey: opts.ObjectKey, FileName: opts.FileName, FileSize: opts.FileSize, } diff --git a/internal/db/message_image_test.go b/internal/db/message_image_test.go new file mode 100644 index 0000000..c4828dc --- /dev/null +++ b/internal/db/message_image_test.go @@ -0,0 +1,39 @@ +package db + +import ( + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestMessageImageCreatePersistsObjectKey(t *testing.T) { + conn, _ := newTestDB(t) + ctx := context.Background() + store := NewMessageImageStore(conn) + + created, err := store.Create(ctx, CreateMessageImageOptions{ + MessageID: 42, + URL: "https://cdn.example.test/assets/2026/11/private.webp", + ObjectKey: "2026/11/private.webp", + FileName: "private.webp", + FileSize: 1234, + }) + require.NoError(t, err) + require.Equal(t, "2026/11/private.webp", created.ObjectKey) + + var persisted MessageImage + require.NoError(t, conn.WithContext(ctx).First(&persisted, created.ID).Error) + require.Equal(t, created.ObjectKey, persisted.ObjectKey) + byMessage, err := store.GetByMessageIDs(ctx, []uint{42}) + require.NoError(t, err) + require.Len(t, byMessage[42], 1) + require.Equal(t, created.ObjectKey, byMessage[42][0].ObjectKey) + + encoded, err := json.Marshal(&persisted) + require.NoError(t, err) + require.NotContains(t, string(encoded), "ObjectKey") + require.NotContains(t, string(encoded), "object_key") + +} diff --git a/internal/db/r2_delete_task.go b/internal/db/r2_delete_task.go new file mode 100644 index 0000000..fec0282 --- /dev/null +++ b/internal/db/r2_delete_task.go @@ -0,0 +1,277 @@ +package db + +import ( + "context" + "crypto/sha256" + "encoding/hex" + + "strings" + "time" + + "github.com/pkg/errors" + "gorm.io/gorm" + "gorm.io/gorm/clause" + + "github.com/tamakyi/TamaBox/internal/conf" + "github.com/tamakyi/TamaBox/internal/dbutil" + "github.com/tamakyi/TamaBox/internal/storage" +) + +var R2DeleteTasks R2DeleteTaskStore + +var _ R2DeleteTaskStore = (*r2DeleteTasks)(nil) + +// R2DeleteTask 是 R2 对象删除的持久化 outbox 任务。 +// 数据库事务只负责登记任务，实际对象删除由 worker 在事务外执行。 +type R2DeleteTask struct { + ID uint `gorm:"primaryKey"` + Bucket string `gorm:"size:255;not null;uniqueIndex:uk_r2_delete_tasks_bucket_key_hash,priority:1"` + Key string `gorm:"type:text;not null"` + KeyHash string `gorm:"size:64;not null;uniqueIndex:uk_r2_delete_tasks_bucket_key_hash,priority:2"` + Attempts int `gorm:"not null;default:0"` + NextRetryAt time.Time `gorm:"not null;index:idx_r2_delete_tasks_next_retry_at"` + LastError string `gorm:"type:text"` + CreatedAt time.Time + UpdatedAt time.Time +} + +type R2DeleteTaskStore interface { + Enqueue(ctx context.Context, bucket string, keys []string) error + EnqueueValues(ctx context.Context, bucket string, values []string) error + ListDue(ctx context.Context, limit int) ([]*R2DeleteTask, error) + Delete(id uint) error + MarkFailed(id uint, nextRetry time.Time, lastError string) error + // HasReferences 仅供 worker 对当前配置的 bucket 调用。 + // 它只查询数据库，不执行网络请求；查询或旧 URL 解析失败时必须返回错误。 + HasReferences(ctx context.Context, bucket, key string) (bool, error) +} + +func NewR2DeleteTaskStore(db *gorm.DB) R2DeleteTaskStore { + return &r2DeleteTasks{DB: db} +} + +type r2DeleteTasks struct { + *gorm.DB +} + +// EnqueueR2DeleteTasks 在调用方事务中登记对象删除任务。 +// 相同 (bucket,key) 重复入队会重置重试状态，使新事务对应的删除请求立即可见。 +func EnqueueR2DeleteTasks(tx *gorm.DB, bucket string, keys []string) error { + if bucket == "" || len(keys) == 0 { + return nil + } + + now := time.Now() + tasks := make([]*R2DeleteTask, 0, len(keys)) + seen := make(map[string]struct{}, len(keys)) + for _, key := range keys { + if key == "" { + continue + } + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + sum := sha256.Sum256([]byte(key)) + tasks = append(tasks, &R2DeleteTask{ + Bucket: bucket, + Key: key, + KeyHash: hex.EncodeToString(sum[:]), + NextRetryAt: now, + }) + } + if len(tasks) == 0 { + return nil + } + + if err := tx.Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "bucket"}, {Name: "key_hash"}}, + DoUpdates: clause.Assignments(map[string]interface{}{ + "attempts": 0, + "next_retry_at": now, + "last_error": "", + "updated_at": now, + }), + }).Create(&tasks).Error; err != nil { + return errors.Wrap(err, "enqueue R2 delete tasks") + } + return nil +} + +func (db *r2DeleteTasks) Enqueue(ctx context.Context, bucket string, keys []string) error { + return EnqueueR2DeleteTasks(db.WithContext(ctx), bucket, keys) +} + +func (db *r2DeleteTasks) EnqueueValues(ctx context.Context, bucket string, values []string) error { + return enqueueR2DeleteTasksForValues(db.WithContext(ctx), bucket, values) +} + +func (db *r2DeleteTasks) ListDue(ctx context.Context, limit int) ([]*R2DeleteTask, error) { + tasks := make([]*R2DeleteTask, 0) + if limit <= 0 { + return tasks, nil + } + if err := db.WithContext(ctx). + Where("next_retry_at <= ?", time.Now()). + Order("next_retry_at ASC, id ASC"). + Limit(limit). + Find(&tasks).Error; err != nil { + return nil, errors.Wrap(err, "list due R2 delete tasks") + } + return tasks, nil +} + +func (db *r2DeleteTasks) Delete(id uint) error { + if err := db.DB.Delete(&R2DeleteTask{}, id).Error; err != nil { + return errors.Wrap(err, "delete R2 delete task") + } + return nil +} + +func (db *r2DeleteTasks) MarkFailed(id uint, nextRetry time.Time, lastError string) error { + if err := db.Model(&R2DeleteTask{}).Where("id = ?", id).Updates(map[string]interface{}{ + "attempts": gorm.Expr("attempts + 1"), + "next_retry_at": nextRetry, + "last_error": lastError, + "updated_at": time.Now(), + }).Error; err != nil { + return errors.Wrap(err, "mark R2 delete task failed") + } + return nil +} + +func (db *r2DeleteTasks) HasReferences(ctx context.Context, bucket, key string) (bool, error) { + if bucket == "" || key == "" { + return false, nil + } + defaultKey, owned, err := r2ObjectKeyFromValue(key) + if err != nil { + return false, err + } + if owned && isDefaultProfileObject(defaultKey) { + return true, nil + } + + quotedKey := dbutil.QuoteIdentifier(db.Dialector.Name(), "key") + var count int64 + if err := db.WithContext(ctx).Unscoped().Model(&UploadImage{}). + Where(quotedKey+" = ?", key).Count(&count).Error; err != nil { + return false, errors.Wrap(err, "check upload image references") + } + if count > 0 { + return true, nil + } + + count = 0 + if err := db.WithContext(ctx).Unscoped().Model(&MessageImage{}). + Where("object_key = ? OR url = ?", key, key).Count(&count).Error; err != nil { + return false, errors.Wrap(err, "check message image references") + } + if count > 0 { + return true, nil + } + + var uploadURLs []string + if err := db.WithContext(ctx).Unscoped().Model(&UploadImage{}). + Where(quotedKey+" LIKE ? OR "+quotedKey+" LIKE ?", "%://%", "//%"). + Pluck(quotedKey, &uploadURLs).Error; err != nil { + return false, errors.Wrap(err, "list legacy upload image URLs") + } + referenced, err := hasLegacyCDNReference(uploadURLs, key) + if err != nil || referenced { + return referenced, err + } + + var messageURLs []string + if err := db.WithContext(ctx).Unscoped().Model(&MessageImage{}). + Where("url LIKE ? OR url LIKE ?", "%://%", "//%"). + Pluck("url", &messageURLs).Error; err != nil { + return false, errors.Wrap(err, "list legacy message image URLs") + } + referenced, err = hasLegacyCDNReference(messageURLs, key) + if err != nil || referenced { + return referenced, err + } + + var profileUsers []*User + if err := db.WithContext(ctx).Unscoped().Model(&User{}). + Select("avatar", "background", "background_image"). + Where("avatar = ? OR background = ? OR background_image = ? OR avatar LIKE ? OR background LIKE ? OR background_image LIKE ?", key, key, key, "%://%", "%://%", "%://%"). + Find(&profileUsers).Error; err != nil { + return false, errors.Wrap(err, "list profile image references") + } + for _, user := range profileUsers { + referenced, err = hasLegacyCDNReference([]string{user.Avatar, user.Background, user.BackgroundImage}, key) + if err != nil || referenced { + return referenced, err + } + } + return false, nil +} + +func isDefaultProfileObject(key string) bool { + for _, value := range []string{conf.Upload.DefaultAvatarURL, conf.Upload.DefaultBackground, conf.Upload.DefaultBackgroundImageURL} { + defaultKey, owned, err := r2ObjectKeyFromValue(value) + if err == nil && owned && defaultKey == key { + return true + } + } + return false +} + +func hasLegacyCDNReference(values []string, targetKey string) (bool, error) { + for _, value := range values { + key, owned, err := r2ObjectKeyFromValue(value) + if err != nil { + return false, err + } + if owned && key == targetKey { + return true, nil + } + } + return false, nil +} + +// r2ObjectKeyFromValue 接受当前对象 key 或旧 CDN URL。 +// 外部 URL 不属于当前 CDN，安全返回 owned=false；畸形 URL 返回错误以避免误判无引用。 +func r2ObjectKeyFromValue(value string) (key string, owned bool, err error) { + if value == "" { + return "", false, nil + } + if !strings.Contains(value, "://") && !strings.HasPrefix(value, "//") { + return value, true, nil + } + key, owned, err = storage.ObjectKeyFromCDNURL(value, conf.Upload.ImageBucketCDNHost) + if err != nil { + return "", false, errors.Wrap(err, "parse legacy CDN URL") + } + if !owned || key == "" { + return "", false, nil + } + return key, true, nil +} + +func enqueueR2DeleteTasksForValues(tx *gorm.DB, bucket string, values []string, excludedKeys ...string) error { + if bucket == "" || len(values) == 0 { + return nil + } + excluded := make(map[string]struct{}, len(excludedKeys)) + for _, key := range excludedKeys { + excluded[key] = struct{}{} + } + keys := make([]string, 0, len(values)) + for _, value := range values { + key, owned, err := r2ObjectKeyFromValue(value) + if err != nil { + return err + } + if !owned || isDefaultProfileObject(key) { + continue + } + if _, skip := excluded[key]; skip { + continue + } + keys = append(keys, key) + } + return EnqueueR2DeleteTasks(tx, bucket, keys) +} diff --git a/internal/db/r2_delete_task_test.go b/internal/db/r2_delete_task_test.go new file mode 100644 index 0000000..187591d --- /dev/null +++ b/internal/db/r2_delete_task_test.go @@ -0,0 +1,141 @@ +package db + +import ( + "context" + "crypto/sha256" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/tamakyi/TamaBox/internal/conf" +) + +func setR2TestConfig(t *testing.T) { + t.Helper() + oldBucket := conf.Upload.ImageBucket + oldCDNHost := conf.Upload.ImageBucketCDNHost + conf.Upload.ImageBucket = "test-r2-bucket" + conf.Upload.ImageBucketCDNHost = "cdn.example.test/assets" + oldDefaults := []string{conf.Upload.DefaultAvatarURL, conf.Upload.DefaultBackground, conf.Upload.DefaultBackgroundImageURL} + conf.Upload.DefaultAvatarURL = "https://cdn.example.test/assets/default-avatar.webp" + conf.Upload.DefaultBackground = "https://cdn.example.test/assets/default-background.webp" + conf.Upload.DefaultBackgroundImageURL = "https://cdn.example.test/assets/default-qr.webp" + t.Cleanup(func() { + conf.Upload.DefaultAvatarURL, conf.Upload.DefaultBackground, conf.Upload.DefaultBackgroundImageURL = oldDefaults[0], oldDefaults[1], oldDefaults[2] + conf.Upload.ImageBucket = oldBucket + conf.Upload.ImageBucketCDNHost = oldCDNHost + }) +} + +func TestR2DeleteTaskStore(t *testing.T) { + setR2TestConfig(t) + conn, _ := newTestDB(t) + ctx := context.Background() + store := NewR2DeleteTaskStore(conn) + + require.NoError(t, store.Enqueue(ctx, "", []string{"ignored.webp"})) + require.NoError(t, store.Enqueue(ctx, conf.Upload.ImageBucket, nil)) + require.NoError(t, store.Enqueue(ctx, conf.Upload.ImageBucket, []string{ + "2026/11/shared.webp", "2026/11/shared.webp", "2026/11/other.webp", "", + })) + + tasks, err := store.ListDue(ctx, 10) + require.NoError(t, err) + require.Len(t, tasks, 2) + byKey := make(map[string]*R2DeleteTask, len(tasks)) + for _, task := range tasks { + byKey[task.Key] = task + require.Equal(t, conf.Upload.ImageBucket, task.Bucket) + require.Len(t, task.KeyHash, 64) + } + shared := byKey["2026/11/shared.webp"] + other := byKey["2026/11/other.webp"] + require.NotNil(t, shared) + require.NotNil(t, other) + require.Equal(t, fmt.Sprintf("%x", sha256.Sum256([]byte(shared.Key))), shared.KeyHash) + + future := time.Now().Add(time.Hour).Truncate(time.Microsecond) + require.NoError(t, store.MarkFailed(shared.ID, future, "temporary failure")) + due, err := store.ListDue(ctx, 10) + require.NoError(t, err) + require.Len(t, due, 1) + require.Equal(t, other.ID, due[0].ID) + + require.NoError(t, store.Enqueue(ctx, conf.Upload.ImageBucket, []string{shared.Key, shared.Key})) + var requeued R2DeleteTask + require.NoError(t, conn.WithContext(ctx).First(&requeued, shared.ID).Error) + require.Equal(t, shared.ID, requeued.ID) + require.Zero(t, requeued.Attempts) + require.Empty(t, requeued.LastError) + require.WithinDuration(t, time.Now(), requeued.NextRetryAt, 5*time.Second) + + // 同一个 key 可在不同 bucket 各有一条任务，但同一 bucket 内仍保持唯一。 + require.NoError(t, store.Enqueue(ctx, "another-bucket", []string{shared.Key})) + var count int64 + require.NoError(t, conn.WithContext(ctx).Model(&R2DeleteTask{}).Count(&count).Error) + require.Equal(t, int64(3), count) + + require.NoError(t, store.Delete(other.ID)) + require.NoError(t, conn.WithContext(ctx).Model(&R2DeleteTask{}).Count(&count).Error) + require.Equal(t, int64(2), count) + empty, err := store.ListDue(ctx, 0) + require.NoError(t, err) + require.NotNil(t, empty) + require.Empty(t, empty) +} + +func TestR2DeleteTaskHasReferences(t *testing.T) { + setR2TestConfig(t) + conn, _ := newTestDB(t) + ctx := context.Background() + store := NewR2DeleteTaskStore(conn) + bucket := conf.Upload.ImageBucket + + upload := &UploadImage{Key: "upload.webp"} + require.NoError(t, conn.WithContext(ctx).Create(upload).Error) + require.NoError(t, conn.WithContext(ctx).Delete(upload).Error) + referenced, err := store.HasReferences(ctx, bucket, upload.Key) + require.NoError(t, err) + require.True(t, referenced, "软删除的 upload_images 仍必须通过 Unscoped 阻止对象删除") + + message := &MessageImage{MessageID: 1, ObjectKey: "message.webp", URL: "https://cdn.example.test/assets/message.webp"} + require.NoError(t, conn.WithContext(ctx).Create(message).Error) + require.NoError(t, conn.WithContext(ctx).Delete(message).Error) + referenced, err = store.HasReferences(ctx, bucket, message.ObjectKey) + require.NoError(t, err) + require.True(t, referenced, "软删除的 message_images 仍必须通过 Unscoped 阻止对象删除") + + legacyMessage := &MessageImage{MessageID: 2, URL: "https://cdn.example.test/assets/legacy-message.webp"} + legacyUpload := &UploadImage{Key: "https://cdn.example.test/assets/legacy-upload.webp"} + external := &MessageImage{MessageID: 3, URL: "https://external.example.net/external.webp"} + require.NoError(t, conn.WithContext(ctx).Create(legacyMessage).Error) + require.NoError(t, conn.WithContext(ctx).Create(legacyUpload).Error) + require.NoError(t, conn.WithContext(ctx).Create(external).Error) + + for _, key := range []string{"legacy-message.webp", "legacy-upload.webp"} { + referenced, err = store.HasReferences(ctx, bucket, key) + require.NoError(t, err) + require.True(t, referenced, key) + } + referenced, err = store.HasReferences(ctx, bucket, "external.webp") + require.NoError(t, err) + require.False(t, referenced) + referenced, err = store.HasReferences(ctx, bucket, "missing.webp") + require.NoError(t, err) + require.False(t, referenced) + + profileUser := &User{Avatar: "https://cdn.example.test/assets/profile-avatar.webp", Background: "profile-background.webp", BackgroundImage: "https://external.example.net/external-profile.webp"} + require.NoError(t, conn.WithContext(ctx).Create(profileUser).Error) + for _, key := range []string{"profile-avatar.webp", "profile-background.webp", "default-avatar.webp", "default-background.webp", "default-qr.webp"} { + referenced, err = store.HasReferences(ctx, bucket, key) + require.NoError(t, err) + require.True(t, referenced, key) + } + + // 第二个引用查询失败时必须显式返回错误，不能误报为无引用。 + require.NoError(t, conn.Migrator().DropTable(&MessageImage{})) + _, err = store.HasReferences(ctx, bucket, "missing-again.webp") + require.Error(t, err) +} diff --git a/internal/db/upload_images.go b/internal/db/upload_images.go index 6067201..80d1b74 100644 --- a/internal/db/upload_images.go +++ b/internal/db/upload_images.go @@ -10,6 +10,8 @@ import ( "github.com/pkg/errors" "github.com/samber/lo" "gorm.io/gorm" + + "github.com/tamakyi/TamaBox/internal/conf" ) var UploadImages UploadImagesStore @@ -23,12 +25,8 @@ type UploadImagesStore interface { // 用于账号导出等场景，将 N 次单题查询合并为少量 IN 查询，避免 N+1。 GetByQuestionIDs(ctx context.Context, questionIDs []uint) (map[uint][]*UploadImage, error) GetByTypeQuestionID(ctx context.Context, typ UploadImageQuestionType, questionID uint) ([]*UploadImage, error) - // DeleteByQuestionID 删除某条提问时级联清理其全部上传图片记录与关联行。 - // - // 用 Unscoped 真删：图片元数据（name/md5/key）随提问一并真删，与 questions.DeleteByID - // 对主表的真删口径一致。软删会留下带 deleted_at 的图片行，且关联行的 question_id - // 指向已不存在的提问。注意这里只清数据库记录，对象存储里的实际文件不在本函数删除 - // （图片可能被多条记录复用，避免误删共享资源）。 + // DeleteByQuestionID 先删除当前题目的关联，再物理删除不再被任何题目引用的主体。 + // 实际删除的主体的对象 key 与元数据删除在同一事务中登记 R2 outbox。 DeleteByQuestionID(ctx context.Context, questionID uint) error } @@ -87,9 +85,24 @@ func (db *uploadImages) Create(ctx context.Context, opts CreateUploadImageOption } if opts.IsDeletingPrevious { - if err := tx.WithContext(ctx).Where("type = ? AND question_id = ?", opts.Type, opts.QuestionID).Delete(&UploadImageQuestion{}).Error; err != nil { + var previousLinks []*UploadImageQuestion + if err := tx.WithContext(ctx). + Where("type = ? AND question_id = ?", opts.Type, opts.QuestionID). + Find(&previousLinks).Error; err != nil { + return errors.Wrap(err, "get previous image question links") + } + previousImageIDs := lo.Map(previousLinks, func(item *UploadImageQuestion, _ int) uint { + return item.UploadImageID + }) + if err := tx.WithContext(ctx). + Where("type = ? AND question_id = ?", opts.Type, opts.QuestionID). + Delete(&UploadImageQuestion{}).Error; err != nil { return errors.Wrap(err, "delete previous image question link") } + // 新图已经创建但尚未关联；即使它与旧图使用相同 key，也绝不加入删除队列。 + if err := deleteUnreferencedUploadImages(ctx, tx, previousImageIDs, opts.Key); err != nil { + return err + } } if err := tx.WithContext(ctx).Create(&UploadImageQuestion{ Type: opts.Type, @@ -167,10 +180,44 @@ func (db *uploadImages) GetByTypeQuestionID(ctx context.Context, typ UploadImage return db.getBy(ctx, "type = ? AND question_id = ?", typ, questionID) } -// DeleteByQuestionID 删除某条提问时级联清理其全部上传图片记录（软删除，与提问保持一致）。 -// 注意：question_id 列在关联表 upload_image_questions 上，不在 upload_images 表上， -// 需先查关联行拿到 imageID 列表，再软删图片主体并删除关联行（同一事务内完成）。 -// 注：仅清理数据库记录，对象存储中的实际文件不在此处删除（避免误删复用资源）。 +// unreferencedUploadImagePredicate 是所有 UploadImage 删除入口共用的共享保护谓词。 +// 调用方必须先解除当前作用域的关联，再使用该谓词收集并删除真正的孤儿主体。 +const unreferencedUploadImagePredicate = `NOT EXISTS ( + SELECT 1 FROM upload_image_questions AS shared_links + WHERE shared_links.upload_image_id = upload_images.id +)` + +func deleteUnreferencedUploadImages(ctx context.Context, tx *gorm.DB, imageIDs []uint, excludedKeys ...string) error { + if len(imageIDs) == 0 { + return nil + } + + var images []*UploadImage + if err := tx.WithContext(ctx).Unscoped(). + Where("id IN ?", imageIDs). + Where(unreferencedUploadImagePredicate). + Find(&images).Error; err != nil { + return errors.Wrap(err, "get unreferenced upload images") + } + if len(images) == 0 { + return nil + } + + ids := lo.Map(images, func(image *UploadImage, _ int) uint { return image.ID }) + if err := tx.WithContext(ctx).Unscoped(). + Where("id IN ?", ids). + Where(unreferencedUploadImagePredicate). + Delete(&UploadImage{}).Error; err != nil { + return errors.Wrap(err, "delete unreferenced upload images") + } + + values := lo.Map(images, func(image *UploadImage, _ int) string { return image.Key }) + if err := enqueueR2DeleteTasksForValues(tx.WithContext(ctx), conf.Upload.ImageBucket, values, excludedKeys...); err != nil { + return errors.Wrap(err, "enqueue deleted upload images") + } + return nil +} + func (db *uploadImages) DeleteByQuestionID(ctx context.Context, questionID uint) error { return db.Transaction(func(tx *gorm.DB) error { var links []*UploadImageQuestion @@ -184,13 +231,10 @@ func (db *uploadImages) DeleteByQuestionID(ctx context.Context, questionID uint) imageIDs := lo.Map(links, func(item *UploadImageQuestion, _ int) uint { return item.UploadImageID }) - if err := tx.WithContext(ctx).Where("id IN ?", imageIDs).Unscoped().Delete(&UploadImage{}).Error; err != nil { - return errors.Wrap(err, "delete upload images by question id") - } - // 关联表 UploadImageQuestion 无 DeletedAt，Delete 本就是真删，无需 Unscoped。 + // 关联表无 DeletedAt，Delete 本就是真删；先解除当前题目关联。 if err := tx.WithContext(ctx).Where("question_id = ?", questionID).Delete(&UploadImageQuestion{}).Error; err != nil { return errors.Wrap(err, "delete image question links") } - return nil + return deleteUnreferencedUploadImages(ctx, tx, imageIDs) }) } diff --git a/internal/db/upload_images_outbox_test.go b/internal/db/upload_images_outbox_test.go new file mode 100644 index 0000000..2a50a15 --- /dev/null +++ b/internal/db/upload_images_outbox_test.go @@ -0,0 +1,95 @@ +package db + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func uploadImageCount(t *testing.T, conn *gorm.DB, id uint) int64 { + t.Helper() + var count int64 + require.NoError(t, conn.Unscoped().Model(&UploadImage{}).Where("id = ?", id).Count(&count).Error) + return count +} + +func TestUploadImagesDeleteByQuestionIDOutbox(t *testing.T) { + setR2TestConfig(t) + conn, _ := newTestDB(t) + ctx := context.Background() + images := NewUploadImagesStore(conn) + tasks := NewR2DeleteTaskStore(conn) + + first, err := images.Create(ctx, CreateUploadImageOptions{QuestionID: 1, Key: "unique.webp"}) + require.NoError(t, err) + duplicate, err := images.Create(ctx, CreateUploadImageOptions{QuestionID: 1, Key: "unique.webp"}) + require.NoError(t, err) + shared, err := images.Create(ctx, CreateUploadImageOptions{QuestionID: 1, Key: "shared.webp"}) + require.NoError(t, err) + require.NoError(t, conn.WithContext(ctx).Create(&UploadImageQuestion{ + UploadImageID: shared.ID, + QuestionID: 2, + }).Error) + + require.NoError(t, images.DeleteByQuestionID(ctx, 1)) + require.Zero(t, uploadImageCount(t, conn, first.ID)) + require.Zero(t, uploadImageCount(t, conn, duplicate.ID)) + require.Equal(t, int64(1), uploadImageCount(t, conn, shared.ID)) + + due, err := tasks.ListDue(ctx, 10) + require.NoError(t, err) + require.Len(t, due, 1, "重复 key 必须去重，共享主体不得入队") + require.Equal(t, "unique.webp", due[0].Key) + + require.NoError(t, images.DeleteByQuestionID(ctx, 2)) + require.Zero(t, uploadImageCount(t, conn, shared.ID)) + due, err = tasks.ListDue(ctx, 10) + require.NoError(t, err) + require.Len(t, due, 2) + require.ElementsMatch(t, []string{"unique.webp", "shared.webp"}, []string{due[0].Key, due[1].Key}) +} + +func TestUploadImagesReplaceDeletesOnlyOrphanedPrevious(t *testing.T) { + setR2TestConfig(t) + conn, _ := newTestDB(t) + ctx := context.Background() + images := NewUploadImagesStore(conn) + tasks := NewR2DeleteTaskStore(conn) + + old, err := images.Create(ctx, CreateUploadImageOptions{ + Type: UploadImageQuestionTypeAnswer, + QuestionID: 10, + Key: "old.webp", + }) + require.NoError(t, err) + shared, err := images.Create(ctx, CreateUploadImageOptions{ + Type: UploadImageQuestionTypeAnswer, + QuestionID: 10, + Key: "shared-old.webp", + }) + require.NoError(t, err) + require.NoError(t, conn.WithContext(ctx).Create(&UploadImageQuestion{ + Type: UploadImageQuestionTypeAnswer, + UploadImageID: shared.ID, + QuestionID: 11, + }).Error) + + newImage, err := images.Create(ctx, CreateUploadImageOptions{ + Type: UploadImageQuestionTypeAnswer, + QuestionID: 10, + Key: "new.webp", + IsDeletingPrevious: true, + }) + require.NoError(t, err) + require.Zero(t, uploadImageCount(t, conn, old.ID)) + require.Equal(t, int64(1), uploadImageCount(t, conn, shared.ID)) + require.Equal(t, int64(1), uploadImageCount(t, conn, newImage.ID)) + + due, err := tasks.ListDue(ctx, 10) + require.NoError(t, err) + require.Len(t, due, 1) + require.Equal(t, "old.webp", due[0].Key) + require.NotEqual(t, newImage.Key, due[0].Key) +} diff --git a/internal/db/users.go b/internal/db/users.go index f400896..ff43300 100644 --- a/internal/db/users.go +++ b/internal/db/users.go @@ -386,10 +386,25 @@ func (db *users) Update(ctx context.Context, id uint, opts UpdateUserOptions) er "Language": opts.Language, } - if err := db.WithContext(ctx).Model(&User{}).Where("id = ?", id).Updates(updates).Error; err != nil { - return errors.Wrap(err, "update user") - } - return nil + return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var current User + if err := tx.WithContext(ctx).Select("id", "avatar", "background", "background_image").First(¤t, id).Error; err != nil { + return errors.Wrap(err, "get current profile images") + } + if err := tx.WithContext(ctx).Model(&User{}).Where("id = ?", id).Updates(updates).Error; err != nil { + return errors.Wrap(err, "update user") + } + oldValues := make([]string, 0, 3) + for _, pair := range [][2]string{{current.Avatar, opts.Avatar}, {current.Background, opts.Background}, {current.BackgroundImage, opts.BackgroundImage}} { + if pair[0] != pair[1] { + oldValues = append(oldValues, pair[0]) + } + } + if err := enqueueR2DeleteTasksForValues(tx.WithContext(ctx), conf.Upload.ImageBucket, oldValues, opts.Avatar, opts.Background, opts.BackgroundImage); err != nil { + return errors.Wrap(err, "enqueue replaced profile images") + } + return nil + }) } // UpdateHarassmentType 仅更新骚扰防护类型（none / register_only）。 @@ -1021,6 +1036,19 @@ func (db *users) BanUserAndCleanup(ctx context.Context, userID uint) error { // 直接复用会把上一条语句的 Where 条件串到下一条上（曾把 messages 的条件串进 follow_ups， // 报「字段 conversation_id 不存在」）。 func purgeUserContent(ctx context.Context, tx *gorm.DB, userID uint) error { + var profile User + if err := tx.WithContext(ctx).Select("id", "avatar", "background", "background_image").First(&profile, userID).Error; err != nil { + return errors.Wrap(err, "get profile images of user") + } + if err := enqueueR2DeleteTasksForValues(tx.WithContext(ctx), conf.Upload.ImageBucket, []string{profile.Avatar, profile.Background, profile.BackgroundImage}); err != nil { + return errors.Wrap(err, "enqueue profile images of user") + } + if err := tx.WithContext(ctx).Model(&User{}).Where("id = ?", userID).Updates(map[string]interface{}{ + "avatar": "", "background": "", "background_image": "", + }).Error; err != nil { + return errors.Wrap(err, "clear profile images of user") + } + // 追问：属于该用户提问下的追问 if err := tx.WithContext(ctx).Unscoped(). Where("question_id IN (SELECT id FROM questions WHERE user_id = ?)", userID). @@ -1029,11 +1057,31 @@ func purgeUserContent(ctx context.Context, tx *gorm.DB, userID uint) error { } // 消息图片必须先于消息删除，否则消息子查询将无法再匹配待删图片。 + // 新数据优先使用物化 ObjectKey，旧数据回退到 URL；外部 URL 不会进入 R2 outbox。 + messageImageScope := "message_id IN (SELECT id FROM messages WHERE conversation_id IN (SELECT id FROM conversations WHERE owner_id = ? OR asker_id = ?))" + var messageImages []*MessageImage if err := tx.WithContext(ctx).Unscoped(). - Where("message_id IN (SELECT id FROM messages WHERE conversation_id IN (SELECT id FROM conversations WHERE owner_id = ? OR asker_id = ?))", userID, userID). + Where(messageImageScope, userID, userID). + Find(&messageImages).Error; err != nil { + return errors.Wrap(err, "get message images of user") + } + messageImageValues := make([]string, 0, len(messageImages)) + for _, image := range messageImages { + value := image.ObjectKey + if value == "" { + value = image.URL + } + messageImageValues = append(messageImageValues, value) + } + if err := tx.WithContext(ctx).Unscoped(). + Where(messageImageScope, userID, userID). Delete(&MessageImage{}).Error; err != nil { return errors.Wrap(err, "delete message images of user") } + if err := enqueueR2DeleteTasksForValues(tx.WithContext(ctx), conf.Upload.ImageBucket, messageImageValues); err != nil { + return errors.Wrap(err, "enqueue deleted message images") + } + // 消息：属于该用户作为 owner 或 asker 的会话 if err := tx.WithContext(ctx).Unscoped(). Where("conversation_id IN (SELECT id FROM conversations WHERE owner_id = ? OR asker_id = ?)", userID, userID). @@ -1048,30 +1096,25 @@ func purgeUserContent(ctx context.Context, tx *gorm.DB, userID uint) error { return errors.Wrap(err, "delete conversations of user") } - // 提问图片元数据：删除本用户上传的孤立图片，以及仅链接到本用户提问的图片。 - // 仍被其他提问引用的图片必须保留；这里只清数据库元数据，不操作对象存储。 - if err := tx.WithContext(ctx).Unscoped(). - Where(`( - uploader_user_id = ? - OR upload_images.id IN ( - SELECT upload_image_id FROM upload_image_questions - WHERE question_id IN (SELECT id FROM questions WHERE user_id = ?) - ) - ) AND NOT EXISTS ( - SELECT 1 FROM upload_image_questions AS shared_links - WHERE shared_links.upload_image_id = upload_images.id - AND shared_links.question_id NOT IN (SELECT id FROM questions WHERE user_id = ?) - )`, userID, userID, userID). - Delete(&UploadImage{}).Error; err != nil { - return errors.Wrap(err, "delete upload images of user") + // 先确定本用户上传或仅与当前作用域有关的候选主体，再解除本用户题目的关联。 + // 最终是否可删除由与 UploadImages.DeleteByQuestionID 完全相同的共享保护谓词决定。 + var uploadImageIDs []uint + if err := tx.WithContext(ctx).Unscoped().Model(&UploadImage{}). + Where(`uploader_user_id = ? OR upload_images.id IN ( + SELECT upload_image_id FROM upload_image_questions + WHERE question_id IN (SELECT id FROM questions WHERE user_id = ?) + )`, userID, userID). + Pluck("upload_images.id", &uploadImageIDs).Error; err != nil { + return errors.Wrap(err, "get upload images of user") } - - // 提问-图片关联 if err := tx.WithContext(ctx).Unscoped(). Where("question_id IN (SELECT id FROM questions WHERE user_id = ?)", userID). Delete(&UploadImageQuestion{}).Error; err != nil { return errors.Wrap(err, "delete upload image links of user") } + if err := deleteUnreferencedUploadImages(ctx, tx, uploadImageIDs); err != nil { + return errors.Wrap(err, "delete upload images of user") + } // 提问 if err := tx.WithContext(ctx).Unscoped(). diff --git a/internal/db/users_outbox_test.go b/internal/db/users_outbox_test.go new file mode 100644 index 0000000..1d58a5e --- /dev/null +++ b/internal/db/users_outbox_test.go @@ -0,0 +1,142 @@ +package db + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "github.com/tamakyi/TamaBox/internal/conf" +) + +func createOutboxTestUser(t *testing.T, ctx context.Context, store UsersStore, email, domain string) *User { + t.Helper() + require.NoError(t, store.Create(ctx, CreateUserOptions{ + Name: domain, + Password: "test-only-password", + Email: email, + Domain: domain, + })) + user, err := store.GetByEmail(ctx, email) + require.NoError(t, err) + return user +} + +func TestPurgeUserContentEnqueuesOwnedR2Objects(t *testing.T) { + setR2TestConfig(t) + conn, _ := newTestDB(t) + ctx := context.Background() + users := NewUsersStore(conn) + owner := createOutboxTestUser(t, ctx, users, "outbox-owner@example.test", "outbox-owner") + other := createOutboxTestUser(t, ctx, users, "outbox-other@example.test", "outbox-other") + require.NoError(t, conn.Model(owner).Updates(map[string]interface{}{ + "avatar": "https://cdn.example.test/assets/profile-owned.webp", "background": "https://cdn.example.test/assets/default-background.webp", "background_image": "https://cdn.example.test/assets/profile-shared.webp", + }).Error) + require.NoError(t, conn.Model(other).Update("background_image", "https://cdn.example.test/assets/profile-shared.webp").Error) + questions := NewQuestionsStore(conn) + question, err := questions.Create(ctx, CreateQuestionOptions{UserID: owner.ID, Content: "owned"}) + require.NoError(t, err) + otherQuestion, err := questions.Create(ctx, CreateQuestionOptions{UserID: other.ID, Content: "other"}) + require.NoError(t, err) + + images := NewUploadImagesStore(conn) + owned, err := images.Create(ctx, CreateUploadImageOptions{QuestionID: question.ID, UploaderUserID: owner.ID, Key: "owned.webp"}) + require.NoError(t, err) + _, err = images.Create(ctx, CreateUploadImageOptions{QuestionID: question.ID, UploaderUserID: owner.ID, Key: "owned.webp"}) + require.NoError(t, err) + legacy, err := images.Create(ctx, CreateUploadImageOptions{ + QuestionID: question.ID, UploaderUserID: owner.ID, Key: "https://cdn.example.test/assets/legacy.webp", + }) + require.NoError(t, err) + external, err := images.Create(ctx, CreateUploadImageOptions{ + QuestionID: question.ID, UploaderUserID: owner.ID, Key: "https://external.example.net/external.webp", + }) + require.NoError(t, err) + shared, err := images.Create(ctx, CreateUploadImageOptions{QuestionID: question.ID, UploaderUserID: owner.ID, Key: "shared.webp"}) + require.NoError(t, err) + require.NoError(t, conn.WithContext(ctx).Create(&UploadImageQuestion{ + UploadImageID: shared.ID, QuestionID: otherQuestion.ID, + }).Error) + + conversation, message, err := NewConversationStore(conn).CreateWithFirstMessage(ctx, CreateConversationOptions{ + OwnerID: owner.ID, AskerID: other.ID, MaxMessages: 20, + }, other.ID, "private") + require.NoError(t, err) + require.NotNil(t, conversation) + messageImages := NewMessageImageStore(conn) + for _, opts := range []CreateMessageImageOptions{ + {MessageID: message.ID, ObjectKey: "message-key.webp", URL: "https://cdn.example.test/assets/message-key.webp"}, + {MessageID: message.ID, ObjectKey: "owned.webp", URL: "https://cdn.example.test/assets/owned.webp"}, + {MessageID: message.ID, URL: "https://cdn.example.test/assets/message-legacy.webp"}, + {MessageID: message.ID, URL: "https://external.example.net/message-external.webp"}, + } { + _, err = messageImages.Create(ctx, opts) + require.NoError(t, err) + } + + require.NoError(t, users.Deactivate(ctx, owner.ID)) + require.Zero(t, uploadImageCount(t, conn, owned.ID)) + require.Zero(t, uploadImageCount(t, conn, legacy.ID)) + require.Zero(t, uploadImageCount(t, conn, external.ID)) + require.Equal(t, int64(1), uploadImageCount(t, conn, shared.ID)) + var sharedLinks int64 + require.NoError(t, conn.WithContext(ctx).Model(&UploadImageQuestion{}). + Where("upload_image_id = ? AND question_id = ?", shared.ID, otherQuestion.ID).Count(&sharedLinks).Error) + require.Equal(t, int64(1), sharedLinks) + + due, err := NewR2DeleteTaskStore(conn).ListDue(ctx, 20) + require.NoError(t, err) + keys := make([]string, 0, len(due)) + for _, task := range due { + keys = append(keys, task.Key) + } + require.ElementsMatch(t, []string{ + "owned.webp", "legacy.webp", "message-key.webp", "message-legacy.webp", "profile-owned.webp", "profile-shared.webp", + }, keys) + require.NotContains(t, keys, "shared.webp") + require.NotContains(t, keys, "external.webp") + require.NotContains(t, keys, "message-external.webp") + require.NotContains(t, keys, "default-background.webp") + var deactivated User + require.NoError(t, conn.Unscoped().First(&deactivated, owner.ID).Error) + require.Empty(t, deactivated.Avatar) + require.Empty(t, deactivated.Background) + require.Empty(t, deactivated.BackgroundImage) + other, err = users.GetByID(ctx, other.ID) + require.NoError(t, err) + require.Equal(t, "https://cdn.example.test/assets/profile-shared.webp", other.BackgroundImage) + referenced, err := NewR2DeleteTaskStore(conn).HasReferences(ctx, conf.Upload.ImageBucket, "profile-shared.webp") + require.NoError(t, err) + require.True(t, referenced) + referenced, err = NewR2DeleteTaskStore(conn).HasReferences(ctx, conf.Upload.ImageBucket, "default-background.webp") + require.NoError(t, err) + require.True(t, referenced) +} + +func TestPurgeUserContentOutboxFailureRollsBack(t *testing.T) { + setR2TestConfig(t) + conn, _ := newTestDB(t) + ctx := context.Background() + users := NewUsersStore(conn) + owner := createOutboxTestUser(t, ctx, users, "outbox-rollback@example.test", "outbox-rollback") + question, err := NewQuestionsStore(conn).Create(ctx, CreateQuestionOptions{UserID: owner.ID, Content: "rollback"}) + require.NoError(t, err) + image, err := NewUploadImagesStore(conn).Create(ctx, CreateUploadImageOptions{ + QuestionID: question.ID, UploaderUserID: owner.ID, Key: "rollback.webp", + }) + require.NoError(t, err) + + require.NoError(t, conn.Migrator().DropTable(&R2DeleteTask{})) + require.Error(t, users.Deactivate(ctx, owner.ID)) + + require.Equal(t, int64(1), uploadImageCount(t, conn, image.ID)) + var links int64 + require.NoError(t, conn.WithContext(ctx).Model(&UploadImageQuestion{}). + Where("upload_image_id = ? AND question_id = ?", image.ID, question.ID).Count(&links).Error) + require.Equal(t, int64(1), links) + var questions int64 + require.NoError(t, conn.WithContext(ctx).Unscoped().Model(&Question{}).Where("id = ?", question.ID).Count(&questions).Error) + require.Equal(t, int64(1), questions) + var user User + require.NoError(t, conn.WithContext(ctx).First(&user, owner.ID).Error) + require.False(t, user.DeletedAt.Valid) +} diff --git a/internal/route/chat/chat_handler.go b/internal/route/chat/chat_handler.go index 40cb03a..aa23954 100644 --- a/internal/route/chat/chat_handler.go +++ b/internal/route/chat/chat_handler.go @@ -277,7 +277,8 @@ func SendMessage(ctx context.Context, cache cache.Cache) { } // 上传图片（如有） - var imageURL string + var imageURL, imageKey string + imageAttached := false if imageFile != nil { uploadedURL, err := storage.UploadPictureToS3(imageFile, imageHeader) if err != nil { @@ -291,6 +292,31 @@ func SendMessage(ctx context.Context, cache cache.Cache) { return } imageURL = uploadedURL + defer func() { + if imageAttached || imageKey != "" { + return + } + if err := db.R2DeleteTasks.EnqueueValues(stdctx.Background(), conf.Upload.ImageBucket, []string{imageURL}); err != nil { + logrus.WithError(err).Error("Failed to enqueue unresolved uploaded message image compensation") + } + }() + parsedKey, owned, err := storage.ObjectKeyFromCDNURL(imageURL, conf.Upload.ImageBucketCDNHost) + if err != nil || !owned || parsedKey == "" { + logrus.WithError(err).WithField("owned", owned).Error("Failed to resolve uploaded message image key") + ctx.SetInternalErrorFlash() + ctx.Redirect(fmt.Sprintf("/_/%s/chat/%s", domain, token)) + return + } + imageKey = parsedKey + // 只有消息图片元数据成功落库才标记 attached；此前任意失败出口都补偿新对象。 + defer func() { + if imageAttached { + return + } + if err := db.R2DeleteTasks.Enqueue(stdctx.Background(), conf.Upload.ImageBucket, []string{imageKey}); err != nil { + logrus.WithError(err).WithField("key", imageKey).Error("Failed to enqueue uploaded message image compensation") + } + }() } // 内容安全检测（仅文本） @@ -341,11 +367,16 @@ func SendMessage(ctx context.Context, cache cache.Cache) { if _, err := db.MessageImages.Create(ctx.Request().Context(), db.CreateMessageImageOptions{ MessageID: msg.ID, URL: imageURL, + ObjectKey: imageKey, FileName: imageHeader.Filename, FileSize: imageHeader.Size, }); err != nil { logrus.WithError(err).Error("Failed to save message image") + ctx.SetInternalErrorFlash() + ctx.Redirect(fmt.Sprintf("/_/%s/chat/%s", domain, token)) + return } + imageAttached = true } // 达到上限时关闭对话（幂等，多实例同时执行也安全） diff --git a/internal/storage/cdn_url.go b/internal/storage/cdn_url.go new file mode 100644 index 0000000..97de578 --- /dev/null +++ b/internal/storage/cdn_url.go @@ -0,0 +1,128 @@ +package storage + +import ( + "fmt" + "net/url" + "strings" + "unicode" +) + +// ObjectKeyFromCDNURL 将属于 cdnBase 的 HTTPS URL 转成原始 S3 对象 key。 +// cdnBase 可省略 scheme，也可带 base path。外站或不属于 base path 的 URL +// 返回 owned=false 且不报错；格式非法或试图伪装路径的 URL 返回错误。 +func ObjectKeyFromCDNURL(rawURL, cdnBase string) (key string, owned bool, err error) { + base, err := parseCDNBase(cdnBase) + if err != nil { + return "", false, err + } + parsed, err := parseCDNURL(rawURL, false) + if err != nil { + return "", false, err + } + if !strings.EqualFold(parsed.Host, base.Host) { + return "", false, nil + } + + basePath := base.Path + if basePath != "" { + if parsed.Path == basePath { + return "", false, fmt.Errorf("CDN URL does not contain an object key") + } + prefix := basePath + "/" + if !strings.HasPrefix(parsed.Path, prefix) { + return "", false, nil + } + key = strings.TrimPrefix(parsed.Path, prefix) + } else { + key = strings.TrimPrefix(parsed.Path, "/") + } + if key == "" { + return "", false, fmt.Errorf("CDN URL does not contain an object key") + } + + // 只去掉 cdnBase 的路径和前导分隔符；不添加或移除 bucket，也不 Clean， + // 因而返回的解码路径仍是实际 S3 key。 + return key, true, nil +} + +func parseCDNBase(cdnBase string) (*url.URL, error) { + if cdnBase == "" { + return nil, fmt.Errorf("CDN base URL is empty") + } + rawBase := cdnBase + if strings.HasPrefix(rawBase, "//") { + rawBase = "https:" + rawBase + } else if !strings.Contains(rawBase, "://") { + rawBase = "https://" + rawBase + } + return parseCDNURL(rawBase, true) +} + +func parseCDNURL(rawURL string, isBase bool) (*url.URL, error) { + if rawURL == "" { + return nil, fmt.Errorf("URL is empty") + } + if strings.Contains(rawURL, `\`) || containsControl(rawURL) { + return nil, fmt.Errorf("URL contains a backslash or control character") + } + parsed, err := url.Parse(rawURL) + if err != nil { + return nil, fmt.Errorf("parse URL: %w", err) + } + if !parsed.IsAbs() || !strings.EqualFold(parsed.Scheme, "https") || parsed.Host == "" || + parsed.Hostname() == "" || parsed.Opaque != "" { + return nil, fmt.Errorf("URL must be an absolute HTTPS URL") + } + if parsed.User != nil { + return nil, fmt.Errorf("URL must not contain userinfo") + } + if parsed.RawQuery != "" || parsed.ForceQuery { + return nil, fmt.Errorf("URL must not contain a query") + } + if parsed.Fragment != "" || strings.Contains(rawURL, "#") { + return nil, fmt.Errorf("URL must not contain a fragment") + } + + // 必须在 url.Parse 的 Path 解码之外检查 EscapedPath：编码斜杠在 Path 中 + // 与普通分隔符无法区分，若放行会导致 CDN 与 S3 对 key 的解释不一致。 + escapedPath := strings.ToLower(parsed.EscapedPath()) + if strings.Contains(escapedPath, "%2f") || strings.Contains(escapedPath, "%5c") { + return nil, fmt.Errorf("URL path contains an encoded slash or backslash") + } + if strings.Contains(parsed.Path, `\`) || containsControl(parsed.Path) { + return nil, fmt.Errorf("URL path contains a backslash or control character") + } + + if isBase && strings.HasSuffix(parsed.Path, "/") { + // 配置中的 CDN base 允许一个末尾斜杠；其余空路径段仍会被拒绝。 + parsed.Path = strings.TrimSuffix(parsed.Path, "/") + parsed.RawPath = "" + } + if parsed.Path == "" { + if isBase { + return parsed, nil + } + return nil, fmt.Errorf("URL path is empty") + } + if !strings.HasPrefix(parsed.Path, "/") { + return nil, fmt.Errorf("URL path must be absolute") + } + for _, segment := range strings.Split(strings.TrimPrefix(parsed.Path, "/"), "/") { + if segment == "" { + return nil, fmt.Errorf("URL path contains an empty segment") + } + if segment == "." || segment == ".." { + return nil, fmt.Errorf("URL path contains a dot segment") + } + } + return parsed, nil +} + +func containsControl(value string) bool { + for _, r := range value { + if unicode.IsControl(r) { + return true + } + } + return false +} diff --git a/internal/storage/cdn_url_test.go b/internal/storage/cdn_url_test.go new file mode 100644 index 0000000..959dee0 --- /dev/null +++ b/internal/storage/cdn_url_test.go @@ -0,0 +1,112 @@ +package storage + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestObjectKeyFromCDNURL(t *testing.T) { + tests := []struct { + name string + rawURL string + cdnBase string + wantKey string + owned bool + wantErr bool + }{ + { + name: "scheme less base", rawURL: "https://cdn.example.com/2026/11/pic.webp", + cdnBase: "cdn.example.com", wantKey: "2026/11/pic.webp", owned: true, + }, + { + name: "base path", rawURL: "https://cdn.example.com/assets/2026/pic.webp", + cdnBase: "https://cdn.example.com/assets", wantKey: "2026/pic.webp", owned: true, + }, + { + name: "base trailing slash and escaped space", rawURL: "https://cdn.example.com/tamabox/a%20b.webp", + cdnBase: "cdn.example.com/tamabox/", wantKey: "a b.webp", owned: true, + }, + { + name: "host case insensitive", rawURL: "https://CDN.EXAMPLE.COM/a.webp", + cdnBase: "cdn.example.com", wantKey: "a.webp", owned: true, + }, + { + name: "preserve bucket prefix", rawURL: "https://cdn.example.com/tamabox/2026/a.webp", + cdnBase: "cdn.example.com", wantKey: "tamabox/2026/a.webp", owned: true, + }, + { + name: "external host not owned", rawURL: "https://outside.example.com/a.webp", + cdnBase: "cdn.example.com", owned: false, + }, + { + name: "host suffix attack not owned", rawURL: "https://cdn.example.com.evil.test/a.webp", + cdnBase: "cdn.example.com", owned: false, + }, + { + name: "base path complete segment", rawURL: "https://cdn.example.com/assets2/a.webp", + cdnBase: "cdn.example.com/assets", owned: false, + }, + { + name: "base path case sensitive", rawURL: "https://cdn.example.com/Assets/a.webp", + cdnBase: "cdn.example.com/assets", owned: false, + }, + {name: "HTTP rejected", rawURL: "http://cdn.example.com/a", cdnBase: "cdn.example.com", wantErr: true}, + {name: "relative rejected", rawURL: "/a", cdnBase: "cdn.example.com", wantErr: true}, + {name: "scheme relative rejected", rawURL: "//cdn.example.com/a", cdnBase: "cdn.example.com", wantErr: true}, + {name: "userinfo rejected", rawURL: "https://user@cdn.example.com/a", cdnBase: "cdn.example.com", wantErr: true}, + {name: "query rejected", rawURL: "https://cdn.example.com/a?x=1", cdnBase: "cdn.example.com", wantErr: true}, + {name: "empty query rejected", rawURL: "https://cdn.example.com/a?", cdnBase: "cdn.example.com", wantErr: true}, + {name: "fragment rejected", rawURL: "https://cdn.example.com/a#fragment", cdnBase: "cdn.example.com", wantErr: true}, + {name: "empty fragment rejected", rawURL: "https://cdn.example.com/a#", cdnBase: "cdn.example.com", wantErr: true}, + {name: "backslash rejected", rawURL: `https://cdn.example.com/a\b`, cdnBase: "cdn.example.com", wantErr: true}, + {name: "encoded backslash rejected", rawURL: "https://cdn.example.com/a%5Cb", cdnBase: "cdn.example.com", wantErr: true}, + {name: "encoded slash rejected", rawURL: "https://cdn.example.com/a%2Fb", cdnBase: "cdn.example.com", wantErr: true}, + {name: "control rejected", rawURL: "https://cdn.example.com/a\x7fb", cdnBase: "cdn.example.com", wantErr: true}, + {name: "encoded control rejected", rawURL: "https://cdn.example.com/a%0Ab", cdnBase: "cdn.example.com", wantErr: true}, + {name: "dot rejected", rawURL: "https://cdn.example.com/a/./b", cdnBase: "cdn.example.com", wantErr: true}, + {name: "dot dot rejected", rawURL: "https://cdn.example.com/a/../b", cdnBase: "cdn.example.com", wantErr: true}, + {name: "encoded dot dot rejected", rawURL: "https://cdn.example.com/a/%2E%2E/b", cdnBase: "cdn.example.com", wantErr: true}, + {name: "empty path rejected", rawURL: "https://cdn.example.com", cdnBase: "cdn.example.com", wantErr: true}, + {name: "root path rejected", rawURL: "https://cdn.example.com/", cdnBase: "cdn.example.com", wantErr: true}, + {name: "empty middle segment rejected", rawURL: "https://cdn.example.com/a//b", cdnBase: "cdn.example.com", wantErr: true}, + {name: "empty last segment rejected", rawURL: "https://cdn.example.com/a/", cdnBase: "cdn.example.com", wantErr: true}, + {name: "missing key below base rejected", rawURL: "https://cdn.example.com/assets", cdnBase: "cdn.example.com/assets", wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + key, owned, err := ObjectKeyFromCDNURL(tt.rawURL, tt.cdnBase) + if tt.wantErr { + require.Error(t, err) + assert.False(t, owned) + assert.Empty(t, key) + return + } + require.NoError(t, err) + assert.Equal(t, tt.owned, owned) + assert.Equal(t, tt.wantKey, key) + }) + } +} + +func TestObjectKeyFromCDNURLRejectsInvalidBase(t *testing.T) { + for _, base := range []string{ + "", + "http://cdn.example.com", + "cdn.example.com?x=1", + "cdn.example.com#fragment", + "cdn.example.com/assets/../other", + "cdn.example.com/assets//other", + "cdn.example.com/assets%2Fother", + "user@cdn.example.com", + } { + t.Run(base, func(t *testing.T) { + key, owned, err := ObjectKeyFromCDNURL("https://cdn.example.com/assets/a", base) + require.Error(t, err) + assert.False(t, owned) + assert.Empty(t, key) + }) + } +} diff --git a/internal/storage/client.go b/internal/storage/client.go index ab15ac3..90be5e2 100644 --- a/internal/storage/client.go +++ b/internal/storage/client.go @@ -5,9 +5,6 @@ import ( "sync" "github.com/aliyun/aliyun-oss-go-sdk/oss" - "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" @@ -31,29 +28,13 @@ func GetS3Client() (*s3.Client, error) { s3ClientErr = errors.New("S3 endpoint not configured") return } - - r2Resolver := aws.EndpointResolverWithOptionsFunc(func(service, region string, options ...interface{}) (aws.Endpoint, error) { - return aws.Endpoint{ - URL: conf.Upload.ImageEndpoint, - HostnameImmutable: true, - Source: aws.EndpointSourceCustom, - }, nil - }) - - cfg, err := config.LoadDefaultConfig(context.Background(), - config.WithEndpointResolverWithOptions(r2Resolver), - config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider( - conf.Upload.ImageAccessID, - conf.Upload.ImageAccessSecret, - "", - )), - config.WithRegion("auto"), + s3Client, s3ClientErr = NewS3Client( + context.Background(), + conf.Upload.ImageEndpoint, + conf.Upload.ImageBucket, + conf.Upload.ImageAccessID, + conf.Upload.ImageAccessSecret, ) - if err != nil { - s3ClientErr = errors.Wrap(err, "load S3 config") - return - } - s3Client = s3.NewFromConfig(cfg) }) return s3Client, s3ClientErr } diff --git a/internal/storage/s3.go b/internal/storage/s3.go index 45ab2e6..a47594d 100644 --- a/internal/storage/s3.go +++ b/internal/storage/s3.go @@ -23,7 +23,6 @@ import ( "github.com/disintegration/imaging" "github.com/pkg/errors" "github.com/sirupsen/logrus" - "github.com/wuhan005/gadget" "github.com/tamakyi/TamaBox/internal/conf" "github.com/tamakyi/TamaBox/internal/imagecheck" @@ -189,19 +188,17 @@ func uploadPictureToS3(file multipart.File, fileHeader *multipart.FileHeader, as return "", errors.Wrap(err, "get S3 client") } - // 7. 上传到 R2 - if err := gadget.Retry(5, func() error { - _, err := client.PutObject(context.Background(), &s3.PutObjectInput{ - Bucket: aws.String(conf.Upload.ImageBucket), - Key: aws.String(key), - Body: bytes.NewReader(uploadData), - ContentLength: aws.Int64(int64(len(uploadData))), - ContentType: aws.String(contentType), - // key 含纳秒时间戳且内容不可变 → 可安全启用 immutable 强缓存， - // 避免每次页面加载都回源校验（曾导致头像/二维码 logo 每次重新下载）。 - CacheControl: aws.String("public, max-age=31536000, immutable"), - }) - return err + // 7. 上传到 R2。现有上传 API 不接收调用方 context，故显式标记为 TODO； + // PutObject 内部会为每次尝试设置独立超时，并对可取消 context 做退避。 + if err := PutObject(context.TODO(), client, &s3.PutObjectInput{ + Bucket: aws.String(conf.Upload.ImageBucket), + Key: aws.String(key), + Body: bytes.NewReader(uploadData), + ContentLength: aws.Int64(int64(len(uploadData))), + ContentType: aws.String(contentType), + // key 含纳秒时间戳且内容不可变 → 可安全启用 immutable 强缓存， + // 避免每次页面加载都回源校验（曾导致头像/二维码 logo 每次重新下载）。 + CacheControl: aws.String("public, max-age=31536000, immutable"), }); err != nil { return "", errors.Wrap(err, "put object after retry") } diff --git a/internal/storage/s3_client.go b/internal/storage/s3_client.go new file mode 100644 index 0000000..ec7b257 --- /dev/null +++ b/internal/storage/s3_client.go @@ -0,0 +1,112 @@ +package storage + +import ( + "context" + "fmt" + "net" + "net/http" + "net/url" + "strings" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/credentials" + "github.com/aws/aws-sdk-go-v2/service/s3" +) + +const ( + s3DialTimeout = 10 * time.Second + s3TLSHandshakeTimeout = 10 * time.Second + s3ResponseHeaderTimeout = 30 * time.Second +) + +// NewS3Client 构造用于 Cloudflare R2 的 S3 客户端。 +// endpoint 可以不带路径，也可以带且仅带 /；后者会被剥离， +// 再由 path-style 请求恰好添加一次 bucket，避免出现 /bucket/bucket/key。 +func NewS3Client(ctx context.Context, endpoint, bucket, accessID, accessSecret string) (*s3.Client, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + normalizedEndpoint, err := normalizeS3Endpoint(endpoint, bucket) + if err != nil { + return nil, err + } + + // 显式构造配置，不加载本机 AWS 配置或默认凭据链，避免离线环境意外访问 + // 元数据服务，也保证这里只使用调用方提供的静态凭据。 + cfg := aws.Config{ + Region: "auto", + BaseEndpoint: aws.String(normalizedEndpoint), + Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider( + accessID, + accessSecret, + "", + )), + HTTPClient: newS3HTTPClient(), + Retryer: func() aws.Retryer { + return aws.NopRetryer{} + }, + } + return s3.NewFromConfig(cfg, func(options *s3.Options) { + options.UsePathStyle = true + }), nil +} + +func normalizeS3Endpoint(endpoint, bucket string) (string, error) { + if bucket == "" || strings.ContainsAny(bucket, `/\`) { + return "", fmt.Errorf("invalid S3 bucket %q", bucket) + } + parsed, err := url.Parse(endpoint) + if err != nil { + return "", fmt.Errorf("parse S3 endpoint: %w", err) + } + if !parsed.IsAbs() || parsed.Host == "" || parsed.Hostname() == "" || parsed.Opaque != "" || + (!strings.EqualFold(parsed.Scheme, "https") && !strings.EqualFold(parsed.Scheme, "http")) { + return "", fmt.Errorf("S3 endpoint must be an absolute HTTP(S) URL: %q", endpoint) + } + if parsed.User != nil { + return "", fmt.Errorf("S3 endpoint must not contain userinfo") + } + if parsed.RawQuery != "" || parsed.ForceQuery { + return "", fmt.Errorf("S3 endpoint must not contain a query") + } + if parsed.Fragment != "" || strings.Contains(endpoint, "#") { + return "", fmt.Errorf("S3 endpoint must not contain a fragment") + } + + // EscapedPath 用于精确比较，拒绝用 %62ucket 等编码形式绕过路径约束。 + // 空路径保持原样；唯一允许的非空路径是 /，包括 "/" 在内的 + // 其他路径均拒绝，避免静默改变 endpoint 的语义。 + switch parsed.EscapedPath() { + case "": + return parsed.String(), nil + case "/" + bucket: + normalized := *parsed + normalized.Path = "" + normalized.RawPath = "" + return normalized.String(), nil + default: + return "", fmt.Errorf("S3 endpoint path must be empty or exactly /%s", bucket) + } +} + +func newS3HTTPClient() *http.Client { + dialer := &net.Dialer{ + Timeout: s3DialTimeout, + KeepAlive: 30 * time.Second, + } + transport := &http.Transport{ + Proxy: http.ProxyFromEnvironment, + DialContext: dialer.DialContext, + ForceAttemptHTTP2: true, + MaxIdleConns: 100, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: s3TLSHandshakeTimeout, + ResponseHeaderTimeout: s3ResponseHeaderTimeout, + ExpectContinueTimeout: 1 * time.Second, + } + + // 不设置 http.Client.Timeout：该超时覆盖读取/写入响应的整个生命周期， + // 会截断大对象传输。每次 PutObject/DeleteObject 尝试另有独立 60 秒上限。 + return &http.Client{Transport: transport} +} diff --git a/internal/storage/s3_client_test.go b/internal/storage/s3_client_test.go new file mode 100644 index 0000000..cd92b15 --- /dev/null +++ b/internal/storage/s3_client_test.go @@ -0,0 +1,126 @@ +package storage + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/s3" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNormalizeS3Endpoint(t *testing.T) { + tests := []struct { + name string + endpoint string + bucket string + want string + wantErr bool + }{ + {name: "empty path", endpoint: "https://account.example.com", bucket: "images", want: "https://account.example.com"}, + {name: "strip exact bucket path", endpoint: "https://account.example.com/images", bucket: "images", want: "https://account.example.com"}, + {name: "httptest HTTP supported", endpoint: "http://127.0.0.1:12345/images", bucket: "images", want: "http://127.0.0.1:12345"}, + {name: "root path rejected", endpoint: "https://account.example.com/", bucket: "images", wantErr: true}, + {name: "different path rejected", endpoint: "https://account.example.com/other", bucket: "images", wantErr: true}, + {name: "bucket trailing slash rejected", endpoint: "https://account.example.com/images/", bucket: "images", wantErr: true}, + {name: "encoded bucket rejected", endpoint: "https://account.example.com/%69mages", bucket: "images", wantErr: true}, + {name: "relative rejected", endpoint: "account.example.com", bucket: "images", wantErr: true}, + {name: "empty hostname rejected", endpoint: "https://:443", bucket: "images", wantErr: true}, + {name: "unsupported scheme rejected", endpoint: "ftp://account.example.com", bucket: "images", wantErr: true}, + {name: "query rejected", endpoint: "https://account.example.com?x=1", bucket: "images", wantErr: true}, + {name: "empty query rejected", endpoint: "https://account.example.com?", bucket: "images", wantErr: true}, + {name: "fragment rejected", endpoint: "https://account.example.com#fragment", bucket: "images", wantErr: true}, + {name: "userinfo rejected", endpoint: "https://user@account.example.com", bucket: "images", wantErr: true}, + {name: "empty bucket rejected", endpoint: "https://account.example.com", bucket: "", wantErr: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := normalizeS3Endpoint(tt.endpoint, tt.bucket) + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestNewS3ClientConfiguration(t *testing.T) { + client, err := NewS3Client(context.Background(), "https://account.example.com/images", "images", "test-id", "test-secret") + require.NoError(t, err) + + options := client.Options() + assert.Equal(t, "auto", options.Region) + assert.True(t, options.UsePathStyle) + credentials, err := options.Credentials.Retrieve(context.Background()) + require.NoError(t, err) + assert.Equal(t, "test-id", credentials.AccessKeyID) + assert.Equal(t, "test-secret", credentials.SecretAccessKey) + + httpClient, ok := options.HTTPClient.(*http.Client) + require.True(t, ok) + assert.Zero(t, httpClient.Timeout, "全局 Client.Timeout 会截断大对象传输") + transport, ok := httpClient.Transport.(*http.Transport) + require.True(t, ok) + assert.NotNil(t, transport.DialContext) + assert.Positive(t, transport.TLSHandshakeTimeout) + assert.Positive(t, transport.ResponseHeaderTimeout) + assert.Positive(t, transport.ExpectContinueTimeout) +} + +func TestS3PutObjectRequestShape(t *testing.T) { + type capturedRequest struct { + method string + path string + body string + contentType string + cacheControl string + authorization string + } + captured := make(chan capturedRequest, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + captured <- capturedRequest{ + method: r.Method, + path: r.URL.Path, + body: string(body), + contentType: r.Header.Get("Content-Type"), + cacheControl: r.Header.Get("Cache-Control"), + authorization: r.Header.Get("Authorization"), + } + w.Header().Set("ETag", `"test-etag"`) + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + client, err := NewS3Client(context.Background(), server.URL+"/images", "images", "test-id", "test-secret") + require.NoError(t, err) + err = PutObject(context.Background(), client, &s3.PutObjectInput{ + Bucket: aws.String("images"), + Key: aws.String("2026/11/picture.webp"), + Body: bytes.NewReader([]byte("fake-image")), + ContentLength: aws.Int64(10), + ContentType: aws.String("image/webp"), + CacheControl: aws.String("public, max-age=31536000, immutable"), + }) + require.NoError(t, err) + + request := <-captured + assert.Equal(t, http.MethodPut, request.method) + assert.Equal(t, "/images/2026/11/picture.webp", request.path, "bucket 必须且只能出现一次") + assert.Equal(t, "fake-image", request.body) + assert.Equal(t, "image/webp", request.contentType) + assert.Equal(t, "public, max-age=31536000, immutable", request.cacheControl) + assert.Contains(t, request.authorization, "Credential=test-id/") + assert.Contains(t, request.authorization, "/auto/s3/aws4_request") +} diff --git a/internal/storage/s3_retry.go b/internal/storage/s3_retry.go new file mode 100644 index 0000000..8276791 --- /dev/null +++ b/internal/storage/s3_retry.go @@ -0,0 +1,167 @@ +package storage + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/s3" + "github.com/aws/smithy-go" +) + +const ( + s3AttemptTimeout = 60 * time.Second + s3MaxAttempts = 3 + s3InitialRetryBackoff = 100 * time.Millisecond +) + +// PutObject 上传对象。每次尝试独立限时 60 秒，最多执行 3 次； +// SDK 内部重试被关闭，确保重试次数、Body 回绕和错误分类只由本层控制。 +func PutObject(ctx context.Context, client *s3.Client, input *s3.PutObjectInput) error { + if client == nil { + return errors.New("nil S3 client") + } + if input == nil { + return errors.New("nil S3 PutObject input") + } + return doS3RequestWithRetry(ctx, input.Body, func(attemptCtx context.Context) error { + _, err := client.PutObject(attemptCtx, input, func(options *s3.Options) { + options.Retryer = aws.NopRetryer{} + }) + return err + }) +} + +// DeleteObject 删除对象。S3 Delete 本身幂等；404 也表示目标已经不存在， +// 因而不重试并返回成功。其他错误采用与 PutObject 相同的超时和重试规则。 +func DeleteObject(ctx context.Context, client *s3.Client, bucket, key string) error { + if client == nil { + return errors.New("nil S3 client") + } + err := doS3RequestWithRetry(ctx, nil, func(attemptCtx context.Context) error { + _, err := client.DeleteObject(attemptCtx, &s3.DeleteObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }, func(options *s3.Options) { + options.Retryer = aws.NopRetryer{} + }) + return err + }) + if err != nil { + if statusCode, ok := s3ErrorHTTPStatus(err); ok && statusCode == 404 { + return nil + } + } + return err +} + +func doS3RequestWithRetry(ctx context.Context, body io.Reader, operation func(context.Context) error) error { + var ( + seeker io.Seeker + bodyStart int64 + bodyStartErr error + ) + if body != nil { + if candidate, ok := body.(io.Seeker); ok { + seeker = candidate + // 记录调用时的位置，而不是武断跳到 0，兼容从切片区间上传的 Reader。 + bodyStart, bodyStartErr = seeker.Seek(0, io.SeekCurrent) + } + } + + for attempt := 0; attempt < s3MaxAttempts; attempt++ { + if err := ctx.Err(); err != nil { + return err + } + if attempt > 0 && seeker != nil { + if bodyStartErr != nil { + return fmt.Errorf("determine S3 request body start: %w", bodyStartErr) + } + if _, err := seeker.Seek(bodyStart, io.SeekStart); err != nil { + return fmt.Errorf("rewind S3 request body: %w", err) + } + } + + attemptCtx, cancel := context.WithTimeout(ctx, s3AttemptTimeout) + err := operation(attemptCtx) + cancel() + if err == nil { + return nil + } + if ctx.Err() != nil { + return ctx.Err() + } + if attempt == s3MaxAttempts-1 || !retryableS3Error(err) { + return err + } + + backoff := s3InitialRetryBackoff << attempt + if err := waitForS3Retry(ctx, backoff); err != nil { + return err + } + } + return nil +} + +func waitForS3Retry(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 + } +} + +type httpStatusCoder interface { + HTTPStatusCode() int +} + +func s3ErrorHTTPStatus(err error) (int, bool) { + var statusErr httpStatusCoder + if errors.As(err, &statusErr) { + return statusErr.HTTPStatusCode(), true + } + return 0, false +} + +func retryableS3Error(err error) bool { + if err == nil { + return false + } + if statusCode, ok := s3ErrorHTTPStatus(err); ok { + return statusCode == 408 || statusCode == 429 || statusCode >= 500 && statusCode <= 599 + } + if errors.Is(err, context.Canceled) { + return false + } + // 父 context 已结束时会在调用方提前返回；这里的 DeadlineExceeded 是 + // 单次尝试或传输层超时，可以在父 context 仍有效时重试。 + if errors.Is(err, context.DeadlineExceeded) { + return true + } + var networkErr net.Error + if errors.As(err, &networkErr) { + return true + } + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + return true + } + + // 某些兼容 S3 的实现可能只返回错误码而没有可提取的 HTTP 状态。 + var apiErr smithy.APIError + if errors.As(err, &apiErr) { + switch apiErr.ErrorCode() { + case "RequestTimeout", "RequestTimeoutException", + "Throttling", "ThrottlingException", "TooManyRequestsException", "SlowDown", + "InternalError", "InternalServerError", "ServiceUnavailable", "ServiceUnavailableException": + return true + } + } + return false +} diff --git a/internal/storage/s3_retry_test.go b/internal/storage/s3_retry_test.go new file mode 100644 index 0000000..c7272ac --- /dev/null +++ b/internal/storage/s3_retry_test.go @@ -0,0 +1,171 @@ +package storage + +import ( + "bytes" + "context" + "io" + "net" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/s3" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPutObjectRetriesAndRewindsBody(t *testing.T) { + var attempts atomic.Int32 + bodies := make(chan string, 3) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + bodies <- string(body) + if attempts.Add(1) < 3 { + http.Error(w, "temporary", http.StatusInternalServerError) + return + } + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + client, err := NewS3Client(context.Background(), server.URL, "images", "id", "secret") + require.NoError(t, err) + err = PutObject(context.Background(), client, &s3.PutObjectInput{ + Bucket: aws.String("images"), + Key: aws.String("retry.txt"), + Body: bytes.NewReader([]byte("complete-body")), + }) + require.NoError(t, err) + assert.Equal(t, int32(3), attempts.Load()) + close(bodies) + for body := range bodies { + assert.Equal(t, "complete-body", body) + } +} + +func TestPutObjectDoesNotRetryPermanentErrors(t *testing.T) { + for _, statusCode := range []int{http.StatusForbidden, http.StatusNotFound} { + t.Run(http.StatusText(statusCode), func(t *testing.T) { + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempts.Add(1) + http.Error(w, "permanent", statusCode) + })) + defer server.Close() + + client, err := NewS3Client(context.Background(), server.URL, "images", "id", "secret") + require.NoError(t, err) + err = PutObject(context.Background(), client, &s3.PutObjectInput{ + Bucket: aws.String("images"), + Key: aws.String("permanent.txt"), + Body: bytes.NewReader([]byte("body")), + }) + require.Error(t, err) + assert.Equal(t, int32(1), attempts.Load()) + }) + } +} + +func TestPutObjectRetriesTemporaryStatusCodes(t *testing.T) { + for _, statusCode := range []int{http.StatusRequestTimeout, http.StatusTooManyRequests, http.StatusBadGateway} { + t.Run(http.StatusText(statusCode), func(t *testing.T) { + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if attempts.Add(1) == 1 { + http.Error(w, "temporary", statusCode) + return + } + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + client, err := NewS3Client(context.Background(), server.URL, "images", "id", "secret") + require.NoError(t, err) + err = PutObject(context.Background(), client, &s3.PutObjectInput{ + Bucket: aws.String("images"), + Key: aws.String("temporary.txt"), + Body: bytes.NewReader([]byte("body")), + }) + require.NoError(t, err) + assert.Equal(t, int32(2), attempts.Load()) + }) + } +} + +func TestPutObjectStopsAfterThreeAttempts(t *testing.T) { + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempts.Add(1) + http.Error(w, "still unavailable", http.StatusServiceUnavailable) + })) + defer server.Close() + + client, err := NewS3Client(context.Background(), server.URL, "images", "id", "secret") + require.NoError(t, err) + err = PutObject(context.Background(), client, &s3.PutObjectInput{ + Bucket: aws.String("images"), + Key: aws.String("unavailable.txt"), + Body: bytes.NewReader([]byte("body")), + }) + require.Error(t, err) + assert.Equal(t, int32(3), attempts.Load()) +} + +func TestPutObjectHonorsCanceledContext(t *testing.T) { + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempts.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + client, err := NewS3Client(context.Background(), server.URL, "images", "id", "secret") + require.NoError(t, err) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + err = PutObject(ctx, client, &s3.PutObjectInput{ + Bucket: aws.String("images"), + Key: aws.String("canceled.txt"), + Body: bytes.NewReader([]byte("body")), + }) + require.ErrorIs(t, err, context.Canceled) + assert.Zero(t, attempts.Load()) +} + +func TestRetryBackoffCanBeCanceled(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + done <- waitForS3Retry(ctx, time.Hour) + }() + cancel() + + select { + case err := <-done: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("退避未响应 context 取消") + } +} + +func TestNetworkErrorsAreRetryable(t *testing.T) { + assert.True(t, retryableS3Error(&net.DNSError{Err: "temporary resolver failure", IsTimeout: true})) +} + +func TestDeleteObjectNotFoundIsIdempotent(t *testing.T) { + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempts.Add(1) + assert.Equal(t, http.MethodDelete, r.Method) + http.Error(w, "not found", http.StatusNotFound) + })) + defer server.Close() + + client, err := NewS3Client(context.Background(), server.URL, "images", "id", "secret") + require.NoError(t, err) + require.NoError(t, DeleteObject(context.Background(), client, "images", "already-gone.txt")) + assert.Equal(t, int32(1), attempts.Load()) +} diff --git a/route/question/page.go b/route/question/page.go index 36776b6..200a8c3 100644 --- a/route/question/page.go +++ b/route/question/page.go @@ -871,7 +871,7 @@ func uploadImage(ctx context.Context, opts uploadImageOptions) error { if err != nil { return errors.Wrap(err, "get S3 client") } - if _, err := client.PutObject(ctx.Request().Context(), &s3.PutObjectInput{ + if err := storage.PutObject(ctx.Request().Context(), client, &s3.PutObjectInput{ Bucket: aws.String(conf.Upload.ImageBucket), Key: aws.String(uploadKey), Body: bytes.NewReader(uploadData), @@ -894,6 +894,9 @@ func uploadImage(ctx context.Context, opts uploadImageOptions) error { IsDeletingPrevious: opts.IsDeletingPrevious, }) if err != nil { + if enqueueErr := db.R2DeleteTasks.Enqueue(ctx.Request().Context(), conf.Upload.ImageBucket, []string{uploadKey}); enqueueErr != nil { + return errors.Wrapf(enqueueErr, "compensate upload image after metadata failure: %v", err) + } return errors.Wrap(err, "create upload image") } return nil diff --git a/route/user/profile.go b/route/user/profile.go index ad83afd..d90c3fa 100644 --- a/route/user/profile.go +++ b/route/user/profile.go @@ -5,6 +5,7 @@ package user import ( + stdctx "context" "encoding/json" "fmt" "net/url" @@ -121,6 +122,17 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) { return } + var uploadedProfileImages []string + profileSaved := false + defer func() { + if profileSaved || len(uploadedProfileImages) == 0 { + return + } + if err := db.R2DeleteTasks.EnqueueValues(stdctx.Background(), conf.Upload.ImageBucket, uploadedProfileImages); err != nil { + logrus.WithError(err).Error("Failed to enqueue uploaded profile image compensation") + } + }() + // 处理头像上传（强制 1:1，见 storage.UploadAvatarToS3） avatarURL := currentUser.Avatar // 默认使用原值 avatarFile, avatarFileHeader, err := ctx.Request().FormFile("avatar") @@ -146,6 +158,7 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) { logrus.WithContext(ctx.Request().Context()).WithError(err).Error("Failed to upload avatar") } else { avatarURL = uploadedURL + uploadedProfileImages = append(uploadedProfileImages, uploadedURL) } } @@ -177,6 +190,7 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) { logrus.WithContext(ctx.Request().Context()).WithError(err).Error("Failed to upload background") } else { backgroundURL = uploadedURL + uploadedProfileImages = append(uploadedProfileImages, uploadedURL) } } @@ -200,6 +214,7 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) { logrus.WithContext(ctx.Request().Context()).WithError(err).Error("Failed to upload QRbackgroundImage") } else { backgroundImageURL = uploadedURL + uploadedProfileImages = append(uploadedProfileImages, uploadedURL) } } @@ -268,6 +283,7 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) { } ctx.SetInternalErrorFlash() } else if passwordChanged { + profileSaved = true // 密码修改成功且其余资料也保存成功：强制登出重登。 // 先 Flush 清会话再写入 flash，保证提示能带到登录页显示。 ctx.Session.Flush() @@ -275,6 +291,7 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) { ctx.Redirect("/login") return } else { + profileSaved = true ctx.SetSuccessFlash(i18n.T(ctx.Lang(), "admin.flash.profileUpdated")) } ctx.Redirect("/user/profile") diff --git a/tamabox.go b/tamabox.go index 7e55ae2..3880010 100644 --- a/tamabox.go +++ b/tamabox.go @@ -30,6 +30,8 @@ func main() { cmd.HidePublic, cmd.Honeypot, cmd.Backup, + cmd.RestoreImages, + cmd.R2Outbox, } if err := app.Run(os.Args); err != nil { logrus.WithError(err).Fatal("Failed to start application")