可靠性: 完善 R2 备份恢复与对象生命周期
This commit is contained in:
@@ -30,6 +30,8 @@ func main() {
|
|||||||
cmd.HidePublic,
|
cmd.HidePublic,
|
||||||
cmd.Honeypot,
|
cmd.Honeypot,
|
||||||
cmd.Backup,
|
cmd.Backup,
|
||||||
|
cmd.RestoreImages,
|
||||||
|
cmd.R2Outbox,
|
||||||
}
|
}
|
||||||
if err := app.Run(os.Args); err != nil {
|
if err := app.Run(os.Args); err != nil {
|
||||||
logrus.WithError(err).Fatal("Failed to start application")
|
logrus.WithError(err).Fatal("Failed to start application")
|
||||||
|
|||||||
+56
-5
@@ -32,7 +32,7 @@ ps -eo pid,lstart,cmd | grep '[t]amabox' # 看进程启动时间,与二进
|
|||||||
|
|
||||||
```bash
|
```bash
|
||||||
./tamabox --help | sed -n '/COMMANDS/,/GLOBAL/p'
|
./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,
|
> 想确认某个字符串有没有编进二进制,**不要用 `strings`** —— 它默认只输出可打印 ASCII,
|
||||||
@@ -77,10 +77,13 @@ pg_dump -Fc <dbname> > before-$(date +%F-%H%M).dump
|
|||||||
| `censor` | 重新审核历史内容 | 是 | 审核规则变更后复检 |
|
| `censor` | 重新审核历史内容 | 是 | 审核规则变更后复检 |
|
||||||
| `uid` | 为用户生成/补齐 uid | 是 | 数据修复 |
|
| `uid` | 为用户生成/补齐 uid | 是 | 数据修复 |
|
||||||
| `honeypot` | 扫描器诱捕取证(假登录页) | 部分是(`prune` / `purge-creds`) | 看谁在扫、拿到了什么字典 |
|
| `honeypot` | 扫描器诱捕取证(假登录页) | 部分是(`prune` / `purge-creds`) | 看谁在扫、拿到了什么字典 |
|
||||||
|
| `backup` | 备份 PostgreSQL 与配置 bucket 的全部 R2 对象 | 写本地备份目录 | 灾难恢复前留档 |
|
||||||
|
| `restore-images` | 校验并恢复 R2 图片及元数据 | 是(除 `--dry-run`) | 图片灾难恢复、隔离恢复演练 |
|
||||||
| `web` | 启动 Web 服务 | — | 生产入口 |
|
| `web` | 启动 Web 服务 | — | 生产入口 |
|
||||||
|
|
||||||
除 `reports`、`honeypot hits/stats/paths/size` 外,所有写操作都要求交互式输入大写 `YES` 确认;
|
除 `reports`、`honeypot hits/stats/paths/size` 外,既有运营写操作要求交互式输入大写 `YES` 确认;
|
||||||
脚本化时加 `--yes`。注意 `honeypot prune` / `purge-creds` 的确认词是 **`yes`**(小写)。
|
脚本化时加 `--yes`。注意 `honeypot prune` / `purge-creds` 的确认词是 **`yes`**(小写)。
|
||||||
|
`restore-images` 不做交互确认:实际写入必须显式提供 `--yes`,`--dry-run` 除外。
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
@@ -329,7 +332,55 @@ pg_dump -Fc <dbname> > 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 保存到 `<output>/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 次」。
|
1. **发现**:收到告警邮件(同一箱主 24h ≥ 8 条 或 7d ≥ 20 条)→ 或箱主后台出现「被举报 N 次」。
|
||||||
2. **取证**:`reports --domain <domain>` 看条目、占比、历史处置。
|
2. **取证**:`reports --domain <domain>` 看条目、占比、历史处置。
|
||||||
@@ -351,7 +402,7 @@ pg_dump -Fc <dbname> > before-$(date +%F-%H%M).dump
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 9. 排错
|
## 10. 排错
|
||||||
|
|
||||||
| 现象 | 原因 | 处理 |
|
| 现象 | 原因 | 处理 |
|
||||||
|---|---|---|
|
|---|---|---|
|
||||||
@@ -374,7 +425,7 @@ pg_dump -Fc <dbname> > before-$(date +%F-%H%M).dump
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 10. 明确不做的事
|
## 11. 明确不做的事
|
||||||
|
|
||||||
- **不做自动处置**:举报数达标只发邮件,不自动封禁。这个场景里举报本身就是攻击武器。
|
- **不做自动处置**:举报数达标只发邮件,不自动封禁。这个场景里举报本身就是攻击武器。
|
||||||
- **不做申诉后台**:申诉走邮件,处置结果由 CLI 写入。
|
- **不做申诉后台**:申诉走邮件,处置结果由 CLI 写入。
|
||||||
|
|||||||
+412
-143
@@ -6,40 +6,46 @@ package cmd
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/hex"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"os"
|
"os"
|
||||||
"os/exec"
|
"os/exec"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"runtime"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/aws/aws-sdk-go-v2/aws"
|
"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/aws/aws-sdk-go-v2/service/s3"
|
||||||
"github.com/pkg/errors"
|
"github.com/pkg/errors"
|
||||||
"github.com/sirupsen/logrus"
|
"github.com/sirupsen/logrus"
|
||||||
"github.com/urfave/cli/v2"
|
"github.com/urfave/cli/v2"
|
||||||
|
|
||||||
"github.com/tamakyi/TamaBox/internal/conf"
|
"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)图片到本地目录。
|
// Backup 全量备份数据库与对象存储(Cloudflare R2)图片到本地目录。
|
||||||
//
|
//
|
||||||
// 为什么必须同时备份图片:站点的所有图片(头像 / 背景 / 二维码 logo / 提问图 / 回答图)
|
// 为什么必须同时备份图片:站点的所有图片(头像 / 背景 / 二维码 logo / 提问图 / 回答图)
|
||||||
// 都以 picture/YYYY/MM/DD/<hex>(新)或 tamabox/YYYY/MM/<hex>.<ext>(历史)的形式存在
|
// 都以根日期 key、picture/ 或 tamabox/ 等形式存在 R2 上,数据库里只存 key 或 CDN URL。
|
||||||
// R2 上,数据库里只存 key(upload_images.key)或 URL(users.avatar/background 里的完整
|
// 只备份数据库的话,恢复后这些引用指向的图片若已不在 R2,页面就全是裂图。
|
||||||
// CDN 地址)。只备份数据库的话,恢复后这些引用指向的图片若已不在 R2(bucket 误删 /
|
|
||||||
// 账号丢失 / 图片被清理),页面就全是裂图。因此备份必须把 R2 上的图片一并落盘。
|
|
||||||
//
|
//
|
||||||
// 输出目录结构:
|
// 输出目录结构:
|
||||||
//
|
//
|
||||||
// <output>/
|
// <output>/
|
||||||
// db-<20060102-150405>.sql # pg_dump 全量导出(schema + data,可用 psql 恢复)
|
// db-<20060102-150405>.sql # pg_dump 全量导出(schema + data,可用 psql 恢复)
|
||||||
// images/<key> # R2 对象,保持原 key 路径
|
// images/<key> # R2 对象,保持原 key 路径
|
||||||
// images-manifest.json # key → md5 / size 清单,用于恢复后完整性校验
|
// images-manifest.json # key → 实际大小 / SHA-256 / GET 元数据
|
||||||
//
|
//
|
||||||
// 用法:TamaBox backup --output /var/backups/tamabox
|
// 用法:TamaBox backup --output /var/backups/tamabox
|
||||||
//
|
//
|
||||||
@@ -47,7 +53,7 @@ import (
|
|||||||
// TamaBox backup --output . --images-only # 只备份图片
|
// TamaBox backup --output . --images-only # 只备份图片
|
||||||
var Backup = &cli.Command{
|
var Backup = &cli.Command{
|
||||||
Name: "backup",
|
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,
|
Action: runBackup,
|
||||||
Flags: []cli.Flag{
|
Flags: []cli.Flag{
|
||||||
&cli.StringFlag{
|
&cli.StringFlag{
|
||||||
@@ -67,6 +73,9 @@ var Backup = &cli.Command{
|
|||||||
}
|
}
|
||||||
|
|
||||||
func runBackup(ctx *cli.Context) error {
|
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 {
|
if err := conf.Init(); err != nil {
|
||||||
return errors.Wrap(err, "load configuration")
|
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)
|
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
|
var dumpErr strings.Builder
|
||||||
cmd.Stdout = f
|
if err := writeFileAtomically(outFile, func(w io.Writer) error {
|
||||||
cmd.Stderr = &dumpErr
|
cmd.Stdout = w
|
||||||
if err := cmd.Run(); err != nil {
|
cmd.Stderr = &dumpErr
|
||||||
return errors.Wrapf(err, "pg_dump failed: %s", strings.TrimSpace(dumpErr.String()))
|
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("数据库已备份")
|
logrus.WithField("file", outFile).Info("数据库已备份")
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// imageManifestEntry 图片备份清单里的一条:key → 大小 → md5(用于恢复后校验完整性)。
|
// writeFileAtomically 保留数据库、manifest 与对象下载共用的原子写修复:
|
||||||
type imageManifestEntry struct {
|
// 只有完整写入、Sync 并关闭成功后才替换目标,失败时清理临时 .part 文件。
|
||||||
Key string `json:"key"`
|
func writeFileAtomically(outFile string, write func(io.Writer) error) error {
|
||||||
Size int64 `json:"size"`
|
f, err := os.CreateTemp(filepath.Dir(outFile), "."+filepath.Base(outFile)+"-*.part")
|
||||||
MD5 string `json:"md5,omitempty"`
|
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。
|
// imageManifestEntry 记录实际下载内容的大小和 SHA-256,以及 GET 返回的原始元数据。
|
||||||
//
|
// ETag 只是服务商元数据,不能当作内容 MD5 或完整性校验的替代品。
|
||||||
// R2 里 tamabox 站图片的 key 都带 `tamabox/` 前缀(如 tamabox/2025/12/<ts>.jpg),
|
type imageManifestEntry struct {
|
||||||
// 这是「bucket 用 path-style 访问时 List 返回的 key 带 bucket 名作前缀」这一 R2 行为。
|
Key string `json:"key"`
|
||||||
// 因此只需列出 `tamabox/` 前缀即可覆盖全部图片。blog/ 是同一 R2 account 下另一个站点
|
Size int64 `json:"size"`
|
||||||
// 的内容,不在 tamabox 的备份范围。
|
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 {
|
func backupImages(ctx context.Context, outDir string) error {
|
||||||
client, err := backupS3Client()
|
client, err := backupS3Client(ctx)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return errors.Wrap(err, "get R2 client for backup")
|
return errors.Wrap(err, "get R2 client for backup")
|
||||||
}
|
}
|
||||||
|
return backupImagesWithClient(ctx, client, conf.Upload.ImageBucket, outDir, "")
|
||||||
|
}
|
||||||
|
|
||||||
imageDir := filepath.Join(outDir, "images")
|
// backupImagesWithClient 允许测试和隔离演练显式指定 prefix。
|
||||||
if err := os.MkdirAll(imageDir, 0o700); err != nil {
|
// 所有分页 List、key 校验和本地路径预检都在第一次 GetObject 之前完成。
|
||||||
return errors.Wrap(err, "create images directory")
|
func backupImagesWithClient(ctx context.Context, client backupS3API, bucket, outDir, prefix string) error {
|
||||||
}
|
keys, err := listAllBackupKeys(ctx, client, bucket, prefix)
|
||||||
|
if err != nil {
|
||||||
var manifest []imageManifestEntry
|
|
||||||
var count int64
|
|
||||||
var totalBytes int64
|
|
||||||
|
|
||||||
// 见函数注释:tamabox/ 前缀即覆盖全部图片(含新旧 key 格式)。
|
|
||||||
if err := backupPrefix(ctx, client, imageDir, "tamabox/", &manifest, &count, &totalBytes); err != nil {
|
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
// 写 manifest:记录 key → 大小 → md5,恢复后可用于校验完整性。
|
imageDir := filepath.Join(outDir, "images")
|
||||||
mf, err := os.Create(filepath.Join(outDir, "images-manifest.json"))
|
plans, err := planBackupDownloads(imageDir, keys)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return errors.Wrap(err, "create manifest file")
|
return err
|
||||||
}
|
}
|
||||||
defer mf.Close()
|
|
||||||
enc := json.NewEncoder(mf)
|
manifest := make([]imageManifestEntry, 0, len(plans))
|
||||||
enc.SetIndent("", " ")
|
var totalBytes int64
|
||||||
if err := enc.Encode(manifest); err != nil {
|
for _, plan := range plans {
|
||||||
return errors.Wrap(err, "write manifest")
|
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{
|
logrus.WithFields(logrus.Fields{
|
||||||
"objects": count,
|
"objects": len(manifest),
|
||||||
"bytes": totalBytes,
|
"bytes": totalBytes,
|
||||||
"dir": imageDir,
|
"dir": imageDir,
|
||||||
|
"prefix": prefix,
|
||||||
}).Info("图片已备份")
|
}).Info("图片已备份")
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// backupPrefix 分页列出并下载指定前缀下的全部对象。
|
func listAllBackupKeys(ctx context.Context, client backupS3API, bucket, prefix string) ([]string, error) {
|
||||||
func backupPrefix(ctx context.Context, client *s3.Client, imageDir, prefix string, manifest *[]imageManifestEntry, count, totalBytes *int64) error {
|
keys := make([]string, 0)
|
||||||
paginator := s3.NewListObjectsV2Paginator(client, &s3.ListObjectsV2Input{
|
var continuationToken *string
|
||||||
Bucket: aws.String(conf.Upload.ImageBucket),
|
seenTokens := make(map[string]struct{})
|
||||||
Prefix: aws.String(prefix),
|
|
||||||
})
|
for {
|
||||||
for paginator.HasMorePages() {
|
page, err := client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
|
||||||
page, err := paginator.NextPage(ctx)
|
Bucket: aws.String(bucket),
|
||||||
|
Prefix: aws.String(prefix),
|
||||||
|
ContinuationToken: continuationToken,
|
||||||
|
})
|
||||||
if err != nil {
|
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/ 前缀),
|
if !aws.ToBool(page.IsTruncated) {
|
||||||
// 恢复时可直接用同一 key 回传 R2。
|
return keys, nil
|
||||||
dst := filepath.Join(imageDir, filepath.FromSlash(key))
|
}
|
||||||
if err := os.MkdirAll(filepath.Dir(dst), 0o700); err != nil {
|
next := aws.ToString(page.NextContinuationToken)
|
||||||
return errors.Wrapf(err, "create dir for %q", key)
|
if next == "" {
|
||||||
}
|
return nil, errors.New("truncated R2 list response has no continuation token")
|
||||||
if err := downloadObject(ctx, client, conf.Upload.ImageBucket, key, dst); err != nil {
|
}
|
||||||
return err
|
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)
|
func planBackupDownloads(imageDir string, keys []string) ([]backupDownloadPlan, error) {
|
||||||
if obj.Size != nil {
|
plans := make([]backupDownloadPlan, 0, len(keys))
|
||||||
size = *obj.Size
|
files := make(map[string]string, len(keys))
|
||||||
}
|
directories := make(map[string]string)
|
||||||
md5 := ""
|
|
||||||
if obj.ETag != nil {
|
|
||||||
md5 = strings.Trim(*obj.ETag, `"`)
|
|
||||||
}
|
|
||||||
*manifest = append(*manifest, imageManifestEntry{Key: key, Size: size, MD5: md5})
|
|
||||||
|
|
||||||
*count++
|
root, err := filepath.Abs(imageDir)
|
||||||
*totalBytes += size
|
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
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// backupS3Client 构造用于备份的 R2 客户端。
|
func localImagePath(imageDir, key string) (string, error) {
|
||||||
//
|
if err := validateObjectKey(key); err != nil {
|
||||||
// 不能直接复用 storage.GetS3Client 做 List:那个客户端把 bucket 路径拼进了 endpoint
|
return "", err
|
||||||
// (image_endpoint 形如 https://<account>.r2.cloudflarestorage.com/<bucket>),
|
|
||||||
// 它对 GetObject/PutObject 恰好可用(R2 容错),但 ListObjectsV2 会把 URL 拼成
|
|
||||||
// /<bucket>/<bucket>?list-type=2,R2 找不到这个 key 前缀而返回 404 NoSuchKey。
|
|
||||||
//
|
|
||||||
// 这里剥离 endpoint 的 bucket 后缀,得到纯 account endpoint,并启用 path-style:
|
|
||||||
// List/Get 都用 Bucket 参数 + 完整 key(List 返回的 key 自带 bucket 名作前缀,
|
|
||||||
// 如 tamabox/2025/12/<ts>.jpg),这是 R2 上唯一稳定可用的姿势。
|
|
||||||
func backupS3Client() (*s3.Client, error) {
|
|
||||||
endpoint := strings.TrimRight(conf.Upload.ImageEndpoint, "/")
|
|
||||||
// 剥离最后一个路径段(bucket 名)。endpoint 形如
|
|
||||||
// https://<account>.r2.cloudflarestorage.com/<bucket>,去掉 /<bucket>。
|
|
||||||
if i := strings.LastIndex(endpoint, "/"); i > len("https://") {
|
|
||||||
endpoint = endpoint[:i]
|
|
||||||
}
|
}
|
||||||
|
root, err := filepath.Abs(imageDir)
|
||||||
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"),
|
|
||||||
)
|
|
||||||
if err != nil {
|
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) {
|
destination := filepath.Join(root, filepath.FromSlash(key))
|
||||||
o.UsePathStyle = true
|
if !pathWithinDirectory(root, destination) || destination == root {
|
||||||
}), nil
|
return "", errors.Errorf("R2 object key %q escapes the images directory", key)
|
||||||
|
}
|
||||||
|
return destination, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// downloadObject 下载单个对象到本地文件,失败时删除半成品避免留下损坏文件。
|
func pathWithinDirectory(root, candidate string) bool {
|
||||||
func downloadObject(ctx context.Context, client *s3.Client, bucket, key, dst string) error {
|
rel, err := filepath.Rel(root, candidate)
|
||||||
|
if err != nil || filepath.IsAbs(rel) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return rel == "." || rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))
|
||||||
|
}
|
||||||
|
|
||||||
|
func canonicalLocalPath(path string) string {
|
||||||
|
cleaned := filepath.Clean(path)
|
||||||
|
if runtime.GOOS == "windows" {
|
||||||
|
return strings.ToLower(cleaned)
|
||||||
|
}
|
||||||
|
return cleaned
|
||||||
|
}
|
||||||
|
|
||||||
|
func prepareBackupImageRoot(imageDir string) error {
|
||||||
|
info, err := os.Lstat(imageDir)
|
||||||
|
if os.IsNotExist(err) {
|
||||||
|
if err := os.MkdirAll(imageDir, 0o700); err != nil {
|
||||||
|
return errors.Wrap(err, "create images directory")
|
||||||
|
}
|
||||||
|
info, err = os.Lstat(imageDir)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return errors.Wrap(err, "inspect images directory")
|
||||||
|
}
|
||||||
|
if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
|
||||||
|
return errors.Errorf("images path %q must be a real directory, not a symlink or file", imageDir)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func preflightBackupLocalPath(imageDir, destination string) error {
|
||||||
|
info, err := os.Lstat(destination)
|
||||||
|
if err == nil && (info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular()) {
|
||||||
|
return errors.Errorf("backup destination %q is not a regular file", destination)
|
||||||
|
}
|
||||||
|
if err != nil && !os.IsNotExist(err) {
|
||||||
|
return errors.Wrapf(err, "inspect backup destination %q", destination)
|
||||||
|
}
|
||||||
|
|
||||||
|
for parent := filepath.Dir(destination); parent != imageDir; parent = filepath.Dir(parent) {
|
||||||
|
if !pathWithinDirectory(imageDir, parent) {
|
||||||
|
return errors.Errorf("backup destination %q escapes images directory", destination)
|
||||||
|
}
|
||||||
|
info, err := os.Lstat(parent)
|
||||||
|
if os.IsNotExist(err) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return errors.Wrapf(err, "inspect backup parent %q", parent)
|
||||||
|
}
|
||||||
|
if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
|
||||||
|
return errors.Errorf("backup parent %q is not a real directory", parent)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func ensureBackupDestinationDirectory(imageDir, destination string) error {
|
||||||
|
root, err := filepath.Abs(imageDir)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
parent := filepath.Dir(destination)
|
||||||
|
rel, err := filepath.Rel(root, parent)
|
||||||
|
if err != nil || !pathWithinDirectory(root, parent) {
|
||||||
|
return errors.Errorf("invalid backup destination %q", destination)
|
||||||
|
}
|
||||||
|
if rel == "." {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
current := root
|
||||||
|
for _, segment := range strings.Split(rel, string(filepath.Separator)) {
|
||||||
|
current = filepath.Join(current, segment)
|
||||||
|
info, err := os.Lstat(current)
|
||||||
|
if os.IsNotExist(err) {
|
||||||
|
if err := os.Mkdir(current, 0o700); err != nil && !os.IsExist(err) {
|
||||||
|
return errors.Wrapf(err, "create backup directory %q", current)
|
||||||
|
}
|
||||||
|
info, err = os.Lstat(current)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return errors.Wrapf(err, "inspect backup directory %q", current)
|
||||||
|
}
|
||||||
|
if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
|
||||||
|
return errors.Errorf("backup parent %q is not a real directory", current)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// backupS3Client 与上传、恢复共用 storage.NewS3Client 的 endpoint 语义:
|
||||||
|
// endpoint 可不带路径或只带 /<bucket>,请求始终以 Bucket 参数和原始完整 key 寻址。
|
||||||
|
// List 返回什么 key 就下载什么 key,不额外添加或剥离 bucket 前缀。
|
||||||
|
func backupS3Client(ctx context.Context) (*s3.Client, error) {
|
||||||
|
return storage.NewS3Client(
|
||||||
|
ctx,
|
||||||
|
conf.Upload.ImageEndpoint,
|
||||||
|
conf.Upload.ImageBucket,
|
||||||
|
conf.Upload.ImageAccessID,
|
||||||
|
conf.Upload.ImageAccessSecret,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// downloadObject 原子下载单个对象,并返回实际落盘内容的 SHA-256 和 GET 元数据。
|
||||||
|
// 每次尝试都有独立 60 秒 context;流中断会从头重试整个对象,最多 3 次。
|
||||||
|
func downloadObject(ctx context.Context, client backupS3API, bucket, key, destination string) (downloadedImageObject, error) {
|
||||||
|
var lastErr error
|
||||||
|
for attempt := 0; attempt < imageObjectMaxAttempts; attempt++ {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return downloadedImageObject{}, err
|
||||||
|
}
|
||||||
|
attemptCtx, cancel := context.WithTimeout(ctx, imageObjectAttemptTimeout)
|
||||||
|
downloaded, err := downloadObjectOnce(attemptCtx, client, bucket, key, destination)
|
||||||
|
cancel()
|
||||||
|
if err == nil {
|
||||||
|
return downloaded, nil
|
||||||
|
}
|
||||||
|
lastErr = err
|
||||||
|
if ctx.Err() != nil {
|
||||||
|
return downloadedImageObject{}, ctx.Err()
|
||||||
|
}
|
||||||
|
if attempt == imageObjectMaxAttempts-1 {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if err := waitForImageObjectRetry(ctx, imageObjectRetryBackoff<<attempt); err != nil {
|
||||||
|
return downloadedImageObject{}, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return downloadedImageObject{}, errors.Wrapf(lastErr, "download %q after %d attempts", key, imageObjectMaxAttempts)
|
||||||
|
}
|
||||||
|
|
||||||
|
func downloadObjectOnce(ctx context.Context, client backupS3API, bucket, key, destination string) (downloadedImageObject, error) {
|
||||||
out, err := client.GetObject(ctx, &s3.GetObjectInput{
|
out, err := client.GetObject(ctx, &s3.GetObjectInput{
|
||||||
Bucket: aws.String(bucket),
|
Bucket: aws.String(bucket),
|
||||||
Key: aws.String(key),
|
Key: aws.String(key),
|
||||||
})
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return errors.Wrapf(err, "get R2 object %q", key)
|
return downloadedImageObject{}, errors.Wrapf(err, "get R2 object %q", key)
|
||||||
|
}
|
||||||
|
if out == nil || out.Body == nil {
|
||||||
|
return downloadedImageObject{}, errors.Errorf("get R2 object %q returned an empty response body", key)
|
||||||
}
|
}
|
||||||
defer out.Body.Close()
|
|
||||||
|
|
||||||
tmp := dst + ".part"
|
hash := sha256.New()
|
||||||
f, err := os.OpenFile(tmp, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
|
var size int64
|
||||||
if err != nil {
|
writeErr := writeFileAtomically(destination, func(w io.Writer) error {
|
||||||
return errors.Wrapf(err, "create %q", tmp)
|
var err error
|
||||||
|
size, err = io.Copy(io.MultiWriter(w, hash), out.Body)
|
||||||
|
if err != nil {
|
||||||
|
return errors.Wrapf(err, "download %q", key)
|
||||||
|
}
|
||||||
|
if out.ContentLength != nil && size != aws.ToInt64(out.ContentLength) {
|
||||||
|
return errors.Wrapf(io.ErrUnexpectedEOF, "download %q: GET content length is %d, received %d", key, aws.ToInt64(out.ContentLength), size)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
_ = out.Body.Close()
|
||||||
|
if writeErr != nil {
|
||||||
|
return downloadedImageObject{}, writeErr
|
||||||
|
}
|
||||||
|
return downloadedImageObject{
|
||||||
|
size: size,
|
||||||
|
sha256: hex.EncodeToString(hash.Sum(nil)),
|
||||||
|
etag: aws.ToString(out.ETag),
|
||||||
|
contentType: aws.ToString(out.ContentType),
|
||||||
|
cacheControl: aws.ToString(out.CacheControl),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func waitForImageObjectRetry(ctx context.Context, delay time.Duration) error {
|
||||||
|
timer := time.NewTimer(delay)
|
||||||
|
defer timer.Stop()
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return ctx.Err()
|
||||||
|
case <-timer.C:
|
||||||
|
return nil
|
||||||
}
|
}
|
||||||
if _, err := io.Copy(f, out.Body); err != nil {
|
|
||||||
f.Close()
|
|
||||||
os.Remove(tmp)
|
|
||||||
return errors.Wrapf(err, "download %q", key)
|
|
||||||
}
|
|
||||||
if err := f.Close(); err != nil {
|
|
||||||
os.Remove(tmp)
|
|
||||||
return errors.Wrapf(err, "close %q", tmp)
|
|
||||||
}
|
|
||||||
if err := os.Rename(tmp, dst); err != nil {
|
|
||||||
os.Remove(tmp)
|
|
||||||
return errors.Wrapf(err, "finalize %q", dst)
|
|
||||||
}
|
|
||||||
return nil
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,437 @@
|
|||||||
|
// Copyright 2026 tamakyi. All rights reserved.
|
||||||
|
// Use of this source code is governed by a MIT-style
|
||||||
|
// license that can be found in the LICENSE file.
|
||||||
|
|
||||||
|
package cmd
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/hex"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/aws/aws-sdk-go-v2/aws"
|
||||||
|
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||||
|
"github.com/aws/smithy-go"
|
||||||
|
"github.com/pkg/errors"
|
||||||
|
"github.com/sirupsen/logrus"
|
||||||
|
"github.com/urfave/cli/v2"
|
||||||
|
|
||||||
|
"github.com/tamakyi/TamaBox/internal/conf"
|
||||||
|
"github.com/tamakyi/TamaBox/internal/storage"
|
||||||
|
)
|
||||||
|
|
||||||
|
const maxImageManifestBytes = 64 << 20 // 64 MiB,足以容纳大量对象且避免恶意清单耗尽内存。
|
||||||
|
|
||||||
|
// RestoreImages 从 backup 生成的 images-manifest.json 和 images/ 恢复 R2 对象。
|
||||||
|
var RestoreImages = &cli.Command{
|
||||||
|
Name: "restore-images",
|
||||||
|
Usage: "Validate and restore backed-up images to R2",
|
||||||
|
Action: runRestoreImages,
|
||||||
|
Flags: []cli.Flag{
|
||||||
|
&cli.StringFlag{
|
||||||
|
Name: "input",
|
||||||
|
Required: true,
|
||||||
|
Usage: "Backup directory or images-manifest.json path",
|
||||||
|
},
|
||||||
|
&cli.BoolFlag{
|
||||||
|
Name: "dry-run",
|
||||||
|
Usage: "Validate local files and remote targets without writing objects",
|
||||||
|
},
|
||||||
|
&cli.BoolFlag{
|
||||||
|
Name: "yes",
|
||||||
|
Usage: "Confirm actual remote writes",
|
||||||
|
},
|
||||||
|
&cli.BoolFlag{
|
||||||
|
Name: "overwrite",
|
||||||
|
Usage: "Allow replacing objects that already exist",
|
||||||
|
},
|
||||||
|
&cli.StringFlag{
|
||||||
|
Name: "target-prefix",
|
||||||
|
Value: "",
|
||||||
|
Usage: "Prefix prepended to every restored key (for isolated drills)",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
type restoreImagesOptions struct {
|
||||||
|
Input string
|
||||||
|
DryRun bool
|
||||||
|
Yes bool
|
||||||
|
Overwrite bool
|
||||||
|
TargetPrefix string
|
||||||
|
}
|
||||||
|
|
||||||
|
// restoreObjectStore 的 PutObject 是带重试的高层操作,便于离线 fake 精确断言零写入。
|
||||||
|
type restoreObjectStore interface {
|
||||||
|
HeadObject(context.Context, *s3.HeadObjectInput, ...func(*s3.Options)) (*s3.HeadObjectOutput, error)
|
||||||
|
PutObject(context.Context, *s3.PutObjectInput) error
|
||||||
|
}
|
||||||
|
|
||||||
|
type r2RestoreObjectStore struct {
|
||||||
|
client *s3.Client
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *r2RestoreObjectStore) HeadObject(ctx context.Context, input *s3.HeadObjectInput, options ...func(*s3.Options)) (*s3.HeadObjectOutput, error) {
|
||||||
|
return s.client.HeadObject(ctx, input, options...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *r2RestoreObjectStore) PutObject(ctx context.Context, input *s3.PutObjectInput) error {
|
||||||
|
// storage.PutObject 提供每次独立 60 秒超时、最多 3 次重试、Body 回绕和可取消退避。
|
||||||
|
return storage.PutObject(ctx, s.client, input)
|
||||||
|
}
|
||||||
|
|
||||||
|
type restorePlannedObject struct {
|
||||||
|
entry imageManifestEntry
|
||||||
|
sourcePath string
|
||||||
|
targetKey string
|
||||||
|
}
|
||||||
|
|
||||||
|
func runRestoreImages(ctx *cli.Context) error {
|
||||||
|
opts := restoreImagesOptions{
|
||||||
|
Input: ctx.String("input"),
|
||||||
|
DryRun: ctx.Bool("dry-run"),
|
||||||
|
Yes: ctx.Bool("yes"),
|
||||||
|
Overwrite: ctx.Bool("overwrite"),
|
||||||
|
TargetPrefix: ctx.String("target-prefix"),
|
||||||
|
}
|
||||||
|
if opts.Input == "" {
|
||||||
|
return errors.New("--input must not be empty")
|
||||||
|
}
|
||||||
|
if !opts.DryRun && !opts.Yes {
|
||||||
|
return errors.New("actual restore writes require --yes (or use --dry-run)")
|
||||||
|
}
|
||||||
|
if err := conf.Init(); err != nil {
|
||||||
|
return errors.Wrap(err, "load configuration")
|
||||||
|
}
|
||||||
|
|
||||||
|
client, err := backupS3Client(ctx.Context)
|
||||||
|
if err != nil {
|
||||||
|
return errors.Wrap(err, "get R2 client for image restore")
|
||||||
|
}
|
||||||
|
return restoreImagesWithStore(ctx.Context, &r2RestoreObjectStore{client: client}, conf.Upload.ImageBucket, opts)
|
||||||
|
}
|
||||||
|
|
||||||
|
func restoreImagesWithStore(ctx context.Context, store restoreObjectStore, bucket string, opts restoreImagesOptions) error {
|
||||||
|
if opts.Input == "" {
|
||||||
|
return errors.New("restore input must not be empty")
|
||||||
|
}
|
||||||
|
if !opts.DryRun && !opts.Yes {
|
||||||
|
return errors.New("actual restore writes require --yes (or use --dry-run)")
|
||||||
|
}
|
||||||
|
|
||||||
|
manifestPath, imageDir, err := resolveRestoreInput(opts.Input)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
entries, err := readImageManifest(manifestPath)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// 关键安全顺序:先验证所有 manifest 条目和全部本地文件,再做任何 Head/Put。
|
||||||
|
// 因此重复 key、非法 hash/size、缺文件、篡改或 symlink 都不会造成部分上传。
|
||||||
|
objects := make([]restorePlannedObject, 0, len(entries))
|
||||||
|
for _, entry := range entries {
|
||||||
|
targetKey, err := restoreTargetKey(opts.TargetPrefix, entry.Key)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
sourcePath, err := localImagePath(imageDir, entry.Key)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := validateRestoreLocalFile(imageDir, sourcePath, entry); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
objects = append(objects, restorePlannedObject{
|
||||||
|
entry: entry,
|
||||||
|
sourcePath: sourcePath,
|
||||||
|
targetKey: targetKey,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// 默认先 Head 所有目标。任何目标已存在都会在 Put 循环开始前失败,保证零 Put。
|
||||||
|
if !opts.Overwrite {
|
||||||
|
for _, object := range objects {
|
||||||
|
headCtx, cancel := context.WithTimeout(ctx, imageObjectAttemptTimeout)
|
||||||
|
_, err := store.HeadObject(headCtx, &s3.HeadObjectInput{
|
||||||
|
Bucket: aws.String(bucket),
|
||||||
|
Key: aws.String(object.targetKey),
|
||||||
|
})
|
||||||
|
cancel()
|
||||||
|
if err == nil {
|
||||||
|
return errors.Errorf("restore target %q already exists; no objects were uploaded (use --overwrite to replace)", object.targetKey)
|
||||||
|
}
|
||||||
|
if !isS3ObjectNotFound(err) {
|
||||||
|
return errors.Wrapf(err, "check restore target %q; no objects were uploaded", object.targetKey)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if opts.DryRun {
|
||||||
|
logrus.WithFields(logrus.Fields{
|
||||||
|
"objects": len(objects),
|
||||||
|
"prefix": opts.TargetPrefix,
|
||||||
|
}).Info("图片恢复演练完成,未写入对象")
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
succeeded := 0
|
||||||
|
for _, object := range objects {
|
||||||
|
file, err := openRestoreFileForPut(object.sourcePath, object.entry.Size)
|
||||||
|
if err != nil {
|
||||||
|
return errors.Wrapf(err, "open validated restore file for %q after %d successful uploads; restore is not transactional", object.entry.Key, succeeded)
|
||||||
|
}
|
||||||
|
input := &s3.PutObjectInput{
|
||||||
|
Bucket: aws.String(bucket),
|
||||||
|
Key: aws.String(object.targetKey),
|
||||||
|
Body: file,
|
||||||
|
ContentLength: aws.Int64(object.entry.Size),
|
||||||
|
ContentType: nonEmptyStringPointer(object.entry.ContentType),
|
||||||
|
CacheControl: nonEmptyStringPointer(object.entry.CacheControl),
|
||||||
|
}
|
||||||
|
putErr := store.PutObject(ctx, input)
|
||||||
|
closeErr := file.Close()
|
||||||
|
if putErr != nil {
|
||||||
|
return errors.Wrapf(putErr, "restore object %q to %q failed after %d successful uploads; restore is not transactional", object.entry.Key, object.targetKey, succeeded)
|
||||||
|
}
|
||||||
|
if closeErr != nil {
|
||||||
|
return errors.Wrapf(closeErr, "close restored file %q (remote object was uploaded); restore is not transactional", object.sourcePath)
|
||||||
|
}
|
||||||
|
succeeded++
|
||||||
|
}
|
||||||
|
|
||||||
|
logrus.WithFields(logrus.Fields{
|
||||||
|
"objects": succeeded,
|
||||||
|
"prefix": opts.TargetPrefix,
|
||||||
|
}).Info("图片恢复完成")
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func resolveRestoreInput(input string) (manifestPath, imageDir string, err error) {
|
||||||
|
info, err := os.Lstat(input)
|
||||||
|
if err != nil {
|
||||||
|
return "", "", errors.Wrapf(err, "inspect restore input %q", input)
|
||||||
|
}
|
||||||
|
if info.Mode()&os.ModeSymlink != 0 {
|
||||||
|
return "", "", errors.Errorf("restore input %q must not be a symlink", input)
|
||||||
|
}
|
||||||
|
|
||||||
|
manifestPath = input
|
||||||
|
if info.IsDir() {
|
||||||
|
manifestPath = filepath.Join(input, "images-manifest.json")
|
||||||
|
} else if !info.Mode().IsRegular() {
|
||||||
|
return "", "", errors.Errorf("restore input %q is neither a directory nor a regular manifest file", input)
|
||||||
|
}
|
||||||
|
return manifestPath, filepath.Join(filepath.Dir(manifestPath), "images"), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func readImageManifest(manifestPath string) ([]imageManifestEntry, error) {
|
||||||
|
info, err := os.Lstat(manifestPath)
|
||||||
|
if err != nil {
|
||||||
|
return nil, errors.Wrapf(err, "inspect image manifest %q", manifestPath)
|
||||||
|
}
|
||||||
|
if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
|
||||||
|
return nil, errors.Errorf("image manifest %q must be a regular non-symlink file", manifestPath)
|
||||||
|
}
|
||||||
|
if info.Size() > maxImageManifestBytes {
|
||||||
|
return nil, errors.Errorf("image manifest %q exceeds the %d-byte limit", manifestPath, maxImageManifestBytes)
|
||||||
|
}
|
||||||
|
|
||||||
|
file, err := os.Open(manifestPath)
|
||||||
|
if err != nil {
|
||||||
|
return nil, errors.Wrapf(err, "open image manifest %q", manifestPath)
|
||||||
|
}
|
||||||
|
data, readErr := io.ReadAll(io.LimitReader(file, maxImageManifestBytes+1))
|
||||||
|
closeErr := file.Close()
|
||||||
|
if readErr != nil {
|
||||||
|
return nil, errors.Wrap(readErr, "read image manifest")
|
||||||
|
}
|
||||||
|
if closeErr != nil {
|
||||||
|
return nil, errors.Wrap(closeErr, "close image manifest")
|
||||||
|
}
|
||||||
|
if len(data) > maxImageManifestBytes {
|
||||||
|
return nil, errors.Errorf("image manifest %q exceeds the %d-byte limit", manifestPath, maxImageManifestBytes)
|
||||||
|
}
|
||||||
|
|
||||||
|
decoder := json.NewDecoder(bytes.NewReader(data))
|
||||||
|
var entries []imageManifestEntry
|
||||||
|
if err := decoder.Decode(&entries); err != nil {
|
||||||
|
return nil, errors.Wrap(err, "decode image manifest")
|
||||||
|
}
|
||||||
|
if entries == nil {
|
||||||
|
return nil, errors.New("image manifest must be a JSON array, not null")
|
||||||
|
}
|
||||||
|
var trailing json.RawMessage
|
||||||
|
if err := decoder.Decode(&trailing); err != io.EOF {
|
||||||
|
if err == nil {
|
||||||
|
return nil, errors.New("image manifest contains trailing JSON")
|
||||||
|
}
|
||||||
|
return nil, errors.Wrap(err, "image manifest contains trailing data")
|
||||||
|
}
|
||||||
|
|
||||||
|
seen := make(map[string]struct{}, len(entries))
|
||||||
|
for index := range entries {
|
||||||
|
entry := &entries[index]
|
||||||
|
if err := validateObjectKey(entry.Key); err != nil {
|
||||||
|
return nil, errors.Wrapf(err, "invalid manifest entry %d", index)
|
||||||
|
}
|
||||||
|
if _, exists := seen[entry.Key]; exists {
|
||||||
|
return nil, errors.Errorf("image manifest contains duplicate key %q", entry.Key)
|
||||||
|
}
|
||||||
|
seen[entry.Key] = struct{}{}
|
||||||
|
if entry.Size < 0 {
|
||||||
|
return nil, errors.Errorf("manifest entry %q has negative size %d", entry.Key, entry.Size)
|
||||||
|
}
|
||||||
|
decodedHash, err := hex.DecodeString(entry.SHA256)
|
||||||
|
if err != nil || len(decodedHash) != sha256.Size {
|
||||||
|
return nil, errors.Errorf("manifest entry %q has invalid SHA-256 %q", entry.Key, entry.SHA256)
|
||||||
|
}
|
||||||
|
entry.SHA256 = strings.ToLower(entry.SHA256)
|
||||||
|
}
|
||||||
|
return entries, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func restoreTargetKey(prefix, key string) (string, error) {
|
||||||
|
if prefix != "" && !strings.HasSuffix(prefix, "/") {
|
||||||
|
prefix += "/"
|
||||||
|
}
|
||||||
|
target := prefix + key
|
||||||
|
if err := validateObjectKey(target); err != nil {
|
||||||
|
return "", errors.Wrapf(err, "invalid restore target for source key %q", key)
|
||||||
|
}
|
||||||
|
return target, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func validateRestoreLocalFile(imageDir, sourcePath string, entry imageManifestEntry) error {
|
||||||
|
root, err := filepath.Abs(imageDir)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := validateRestoreDirectoryChain(root, filepath.Dir(sourcePath)); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
info, err := os.Lstat(sourcePath)
|
||||||
|
if err != nil {
|
||||||
|
return errors.Wrapf(err, "inspect local image for %q", entry.Key)
|
||||||
|
}
|
||||||
|
if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
|
||||||
|
return errors.Errorf("local image for %q must be a regular non-symlink file", entry.Key)
|
||||||
|
}
|
||||||
|
if info.Size() != entry.Size {
|
||||||
|
return errors.Errorf("local image size mismatch for %q: manifest=%d actual=%d", entry.Key, entry.Size, info.Size())
|
||||||
|
}
|
||||||
|
|
||||||
|
file, err := os.Open(sourcePath)
|
||||||
|
if err != nil {
|
||||||
|
return errors.Wrapf(err, "open local image for %q", entry.Key)
|
||||||
|
}
|
||||||
|
openedInfo, statErr := file.Stat()
|
||||||
|
if statErr != nil {
|
||||||
|
_ = file.Close()
|
||||||
|
return errors.Wrapf(statErr, "stat opened local image for %q", entry.Key)
|
||||||
|
}
|
||||||
|
if !openedInfo.Mode().IsRegular() || openedInfo.Size() != entry.Size {
|
||||||
|
_ = file.Close()
|
||||||
|
return errors.Errorf("opened local image for %q is not the expected regular file", entry.Key)
|
||||||
|
}
|
||||||
|
hash := sha256.New()
|
||||||
|
_, hashErr := io.Copy(hash, file)
|
||||||
|
closeErr := file.Close()
|
||||||
|
if hashErr != nil {
|
||||||
|
return errors.Wrapf(hashErr, "hash local image for %q", entry.Key)
|
||||||
|
}
|
||||||
|
if closeErr != nil {
|
||||||
|
return errors.Wrapf(closeErr, "close local image for %q", entry.Key)
|
||||||
|
}
|
||||||
|
actualHash := hex.EncodeToString(hash.Sum(nil))
|
||||||
|
if actualHash != entry.SHA256 {
|
||||||
|
return errors.Errorf("local image SHA-256 mismatch for %q: manifest=%s actual=%s", entry.Key, entry.SHA256, actualHash)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func validateRestoreDirectoryChain(root, parent string) error {
|
||||||
|
rootInfo, err := os.Lstat(root)
|
||||||
|
if err != nil {
|
||||||
|
return errors.Wrapf(err, "inspect restore images directory %q", root)
|
||||||
|
}
|
||||||
|
if rootInfo.Mode()&os.ModeSymlink != 0 || !rootInfo.IsDir() {
|
||||||
|
return errors.Errorf("restore images path %q must be a real directory", root)
|
||||||
|
}
|
||||||
|
rel, err := filepath.Rel(root, parent)
|
||||||
|
if err != nil || !pathWithinDirectory(root, parent) {
|
||||||
|
return errors.Errorf("restore file parent %q escapes images directory", parent)
|
||||||
|
}
|
||||||
|
if rel == "." {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
current := root
|
||||||
|
for _, segment := range strings.Split(rel, string(filepath.Separator)) {
|
||||||
|
current = filepath.Join(current, segment)
|
||||||
|
info, err := os.Lstat(current)
|
||||||
|
if err != nil {
|
||||||
|
return errors.Wrapf(err, "inspect restore parent directory %q", current)
|
||||||
|
}
|
||||||
|
if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
|
||||||
|
return errors.Errorf("restore parent %q must be a real directory", current)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func openRestoreFileForPut(sourcePath string, expectedSize int64) (*os.File, error) {
|
||||||
|
info, err := os.Lstat(sourcePath)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() || info.Size() != expectedSize {
|
||||||
|
return nil, fmt.Errorf("local file %q changed after validation", sourcePath)
|
||||||
|
}
|
||||||
|
file, err := os.Open(sourcePath)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
openedInfo, err := file.Stat()
|
||||||
|
if err != nil {
|
||||||
|
_ = file.Close()
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if !openedInfo.Mode().IsRegular() || openedInfo.Size() != expectedSize {
|
||||||
|
_ = file.Close()
|
||||||
|
return nil, fmt.Errorf("local file %q changed after validation", sourcePath)
|
||||||
|
}
|
||||||
|
return file, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func isS3ObjectNotFound(err error) bool {
|
||||||
|
if err == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
var statusErr interface{ HTTPStatusCode() int }
|
||||||
|
if errors.As(err, &statusErr) && statusErr.HTTPStatusCode() == 404 {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
var apiErr smithy.APIError
|
||||||
|
if errors.As(err, &apiErr) {
|
||||||
|
switch apiErr.ErrorCode() {
|
||||||
|
case "NotFound", "NoSuchKey", "404":
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func nonEmptyStringPointer(value string) *string {
|
||||||
|
if value == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return aws.String(value)
|
||||||
|
}
|
||||||
@@ -0,0 +1,184 @@
|
|||||||
|
package cmd
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"crypto/rand"
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/hex"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"os"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/aws/aws-sdk-go-v2/aws"
|
||||||
|
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
|
||||||
|
"github.com/tamakyi/TamaBox/internal/conf"
|
||||||
|
"github.com/tamakyi/TamaBox/internal/storage"
|
||||||
|
)
|
||||||
|
|
||||||
|
// TestR2BackupRestoreIntegration 默认跳过。只有显式设置 TB_R2_INTEGRATION=1,
|
||||||
|
// 并通过 NEKOBOX_CONFIG_PATH 提供真实 R2 配置时才运行。
|
||||||
|
// 测试只访问 tamabox-backup-test/<随机>/ 下的 source/ 与 target/ 两个专属前缀,
|
||||||
|
// 绝不执行全桶 List,也不会读取、覆盖或删除其他 key。
|
||||||
|
func TestR2BackupRestoreIntegration(t *testing.T) {
|
||||||
|
if os.Getenv("TB_R2_INTEGRATION") != "1" {
|
||||||
|
t.Skip("set TB_R2_INTEGRATION=1 to run the real R2 integration test")
|
||||||
|
}
|
||||||
|
if os.Getenv("NEKOBOX_CONFIG_PATH") == "" {
|
||||||
|
t.Fatal("NEKOBOX_CONFIG_PATH must point to the R2 integration configuration")
|
||||||
|
}
|
||||||
|
require.NoError(t, conf.Init())
|
||||||
|
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
|
||||||
|
defer cancel()
|
||||||
|
client, err := storage.NewS3Client(
|
||||||
|
ctx,
|
||||||
|
conf.Upload.ImageEndpoint,
|
||||||
|
conf.Upload.ImageBucket,
|
||||||
|
conf.Upload.ImageAccessID,
|
||||||
|
conf.Upload.ImageAccessSecret,
|
||||||
|
)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
random := make([]byte, 12)
|
||||||
|
_, err = rand.Read(random)
|
||||||
|
require.NoError(t, err)
|
||||||
|
basePrefix := fmt.Sprintf("tamabox-backup-test/%d-%s/", time.Now().UTC().UnixNano(), hex.EncodeToString(random))
|
||||||
|
sourcePrefix := basePrefix + "source/"
|
||||||
|
targetPrefix := basePrefix + "target/"
|
||||||
|
|
||||||
|
objects := []struct {
|
||||||
|
relativeKey string
|
||||||
|
body []byte
|
||||||
|
contentType string
|
||||||
|
cacheControl string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
relativeKey: "2026/11/04/root.webp",
|
||||||
|
body: []byte("tamabox-r2-integration-root"),
|
||||||
|
contentType: "image/webp",
|
||||||
|
cacheControl: "public, max-age=60",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
relativeKey: "picture/2026/11/04/picture.webp",
|
||||||
|
body: []byte("tamabox-r2-integration-picture"),
|
||||||
|
contentType: "image/webp",
|
||||||
|
cacheControl: "public, max-age=120",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
// created 只记录本测试成功创建的精确 key。任何失败清理也只删除这些 key。
|
||||||
|
created := make(map[string]bool)
|
||||||
|
t.Cleanup(func() {
|
||||||
|
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||||||
|
defer cleanupCancel()
|
||||||
|
for key := range created {
|
||||||
|
if err := storage.DeleteObject(cleanupCtx, client, conf.Upload.ImageBucket, key); err != nil {
|
||||||
|
t.Errorf("clean up own R2 key %q: %v", key, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
for _, object := range objects {
|
||||||
|
key := sourcePrefix + object.relativeKey
|
||||||
|
err := storage.PutObject(ctx, client, &s3.PutObjectInput{
|
||||||
|
Bucket: aws.String(conf.Upload.ImageBucket),
|
||||||
|
Key: aws.String(key),
|
||||||
|
Body: bytes.NewReader(object.body),
|
||||||
|
ContentLength: aws.Int64(int64(len(object.body))),
|
||||||
|
ContentType: aws.String(object.contentType),
|
||||||
|
CacheControl: aws.String(object.cacheControl),
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
created[key] = true
|
||||||
|
}
|
||||||
|
|
||||||
|
backupDir := t.TempDir()
|
||||||
|
require.NoError(t, backupImagesWithClient(ctx, client, conf.Upload.ImageBucket, backupDir, sourcePrefix))
|
||||||
|
manifest, err := readImageManifest(backupDir + string(os.PathSeparator) + "images-manifest.json")
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Len(t, manifest, len(objects))
|
||||||
|
manifestKeys := make(map[string]bool, len(manifest))
|
||||||
|
for _, entry := range manifest {
|
||||||
|
manifestKeys[entry.Key] = true
|
||||||
|
}
|
||||||
|
for _, object := range objects {
|
||||||
|
require.True(t, manifestKeys[sourcePrefix+object.relativeKey])
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, object := range objects {
|
||||||
|
sourceKey := sourcePrefix + object.relativeKey
|
||||||
|
require.NoError(t, storage.DeleteObject(ctx, client, conf.Upload.ImageBucket, sourceKey))
|
||||||
|
delete(created, sourceKey)
|
||||||
|
}
|
||||||
|
|
||||||
|
store := &integrationTrackingRestoreStore{
|
||||||
|
base: &r2RestoreObjectStore{client: client},
|
||||||
|
onUpload: func(key string) {
|
||||||
|
created[key] = true
|
||||||
|
},
|
||||||
|
}
|
||||||
|
require.NoError(t, restoreImagesWithStore(ctx, store, conf.Upload.ImageBucket, restoreImagesOptions{
|
||||||
|
Input: backupDir,
|
||||||
|
Yes: true,
|
||||||
|
TargetPrefix: targetPrefix,
|
||||||
|
}))
|
||||||
|
|
||||||
|
for _, object := range objects {
|
||||||
|
targetKey := targetPrefix + sourcePrefix + object.relativeKey
|
||||||
|
require.True(t, created[targetKey], "restore must report its own successful Put")
|
||||||
|
output, err := client.GetObject(ctx, &s3.GetObjectInput{
|
||||||
|
Bucket: aws.String(conf.Upload.ImageBucket),
|
||||||
|
Key: aws.String(targetKey),
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
body, readErr := io.ReadAll(output.Body)
|
||||||
|
closeErr := output.Body.Close()
|
||||||
|
require.NoError(t, readErr)
|
||||||
|
require.NoError(t, closeErr)
|
||||||
|
digest := sha256.Sum256(object.body)
|
||||||
|
require.Equal(t, fmt.Sprintf("%x", digest), fmt.Sprintf("%x", sha256.Sum256(body)))
|
||||||
|
require.Equal(t, object.contentType, aws.ToString(output.ContentType))
|
||||||
|
require.Equal(t, object.cacheControl, aws.ToString(output.CacheControl))
|
||||||
|
}
|
||||||
|
|
||||||
|
for key := range created {
|
||||||
|
require.NoError(t, storage.DeleteObject(ctx, client, conf.Upload.ImageBucket, key))
|
||||||
|
delete(created, key)
|
||||||
|
}
|
||||||
|
requireOwnPrefixEmpty(ctx, t, client, conf.Upload.ImageBucket, sourcePrefix)
|
||||||
|
requireOwnPrefixEmpty(ctx, t, client, conf.Upload.ImageBucket, targetPrefix)
|
||||||
|
}
|
||||||
|
|
||||||
|
type integrationTrackingRestoreStore struct {
|
||||||
|
base restoreObjectStore
|
||||||
|
onUpload func(string)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *integrationTrackingRestoreStore) HeadObject(ctx context.Context, input *s3.HeadObjectInput, options ...func(*s3.Options)) (*s3.HeadObjectOutput, error) {
|
||||||
|
return s.base.HeadObject(ctx, input, options...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *integrationTrackingRestoreStore) PutObject(ctx context.Context, input *s3.PutObjectInput) error {
|
||||||
|
if err := s.base.PutObject(ctx, input); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
s.onUpload(aws.ToString(input.Key))
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func requireOwnPrefixEmpty(ctx context.Context, t *testing.T, client *s3.Client, bucket, prefix string) {
|
||||||
|
t.Helper()
|
||||||
|
require.NotEmpty(t, prefix)
|
||||||
|
output, err := client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
|
||||||
|
Bucket: aws.String(bucket),
|
||||||
|
Prefix: aws.String(prefix),
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Empty(t, output.Contents, "专属 prefix 清理后必须为空:%q", prefix)
|
||||||
|
require.False(t, aws.ToBool(output.IsTruncated), "随机专属 prefix 下对象数不应触发分页")
|
||||||
|
}
|
||||||
@@ -0,0 +1,603 @@
|
|||||||
|
package cmd
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/aws/aws-sdk-go-v2/aws"
|
||||||
|
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||||
|
"github.com/aws/aws-sdk-go-v2/service/s3/types"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
type fakeBackupObject struct {
|
||||||
|
body []byte
|
||||||
|
etag string
|
||||||
|
contentType string
|
||||||
|
cacheControl string
|
||||||
|
interruptedAttempts int
|
||||||
|
wrongLengthAttempts int
|
||||||
|
}
|
||||||
|
|
||||||
|
type fakeBackupS3 struct {
|
||||||
|
pages []*s3.ListObjectsV2Output
|
||||||
|
listInputs []*s3.ListObjectsV2Input
|
||||||
|
objects map[string]fakeBackupObject
|
||||||
|
getCalls map[string]int
|
||||||
|
getOrder []string
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeBackupS3) ListObjectsV2(_ context.Context, input *s3.ListObjectsV2Input, _ ...func(*s3.Options)) (*s3.ListObjectsV2Output, error) {
|
||||||
|
f.listInputs = append(f.listInputs, input)
|
||||||
|
index := len(f.listInputs) - 1
|
||||||
|
if index >= len(f.pages) {
|
||||||
|
return nil, fmt.Errorf("unexpected ListObjectsV2 call %d", index+1)
|
||||||
|
}
|
||||||
|
return f.pages[index], nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeBackupS3) GetObject(_ context.Context, input *s3.GetObjectInput, _ ...func(*s3.Options)) (*s3.GetObjectOutput, error) {
|
||||||
|
key := aws.ToString(input.Key)
|
||||||
|
if f.getCalls == nil {
|
||||||
|
f.getCalls = make(map[string]int)
|
||||||
|
}
|
||||||
|
f.getCalls[key]++
|
||||||
|
f.getOrder = append(f.getOrder, key)
|
||||||
|
object, exists := f.objects[key]
|
||||||
|
if !exists {
|
||||||
|
return nil, fmt.Errorf("unexpected GetObject key %q", key)
|
||||||
|
}
|
||||||
|
|
||||||
|
var body io.ReadCloser
|
||||||
|
if f.getCalls[key] <= object.interruptedAttempts {
|
||||||
|
body = io.NopCloser(&interruptedReader{data: append([]byte(nil), object.body...)})
|
||||||
|
} else {
|
||||||
|
body = io.NopCloser(strings.NewReader(string(object.body)))
|
||||||
|
}
|
||||||
|
contentLength := int64(len(object.body))
|
||||||
|
if f.getCalls[key] <= object.wrongLengthAttempts {
|
||||||
|
contentLength++
|
||||||
|
}
|
||||||
|
return &s3.GetObjectOutput{
|
||||||
|
Body: body,
|
||||||
|
ContentLength: aws.Int64(contentLength),
|
||||||
|
ETag: aws.String(object.etag),
|
||||||
|
ContentType: aws.String(object.contentType),
|
||||||
|
CacheControl: aws.String(object.cacheControl),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type interruptedReader struct {
|
||||||
|
data []byte
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *interruptedReader) Read(p []byte) (int, error) {
|
||||||
|
if len(r.data) == 0 {
|
||||||
|
return 0, io.ErrUnexpectedEOF
|
||||||
|
}
|
||||||
|
n := copy(p, r.data)
|
||||||
|
r.data = r.data[n:]
|
||||||
|
if len(r.data) == 0 {
|
||||||
|
return n, io.ErrUnexpectedEOF
|
||||||
|
}
|
||||||
|
return n, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func backupListPage(keys []string, truncated bool, nextToken string) *s3.ListObjectsV2Output {
|
||||||
|
objects := make([]types.Object, 0, len(keys))
|
||||||
|
for _, key := range keys {
|
||||||
|
objects = append(objects, types.Object{
|
||||||
|
Key: aws.String(key),
|
||||||
|
Size: aws.Int64(999999), // 故意不可信;manifest 必须记录实际下载字节数。
|
||||||
|
ETag: aws.String("list-etag-must-not-be-used"),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
output := &s3.ListObjectsV2Output{Contents: objects, IsTruncated: aws.Bool(truncated)}
|
||||||
|
if nextToken != "" {
|
||||||
|
output.NextContinuationToken = aws.String(nextToken)
|
||||||
|
}
|
||||||
|
return output
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBackupImagesWithClientCoversAllKeysAndPages(t *testing.T) {
|
||||||
|
keys := []string{
|
||||||
|
"2026/11/04/root.webp",
|
||||||
|
"picture/2026/11/04/picture.webp",
|
||||||
|
"tamabox/2025/12/historical.jpg",
|
||||||
|
"static/site.css",
|
||||||
|
}
|
||||||
|
objects := make(map[string]fakeBackupObject, len(keys))
|
||||||
|
for index, key := range keys {
|
||||||
|
objects[key] = fakeBackupObject{
|
||||||
|
body: []byte(fmt.Sprintf("actual-body-%d", index)),
|
||||||
|
etag: fmt.Sprintf("get-etag-%d", index),
|
||||||
|
contentType: []string{"image/webp", "image/webp", "image/jpeg", "text/css"}[index],
|
||||||
|
cacheControl: fmt.Sprintf("public, max-age=%d", index+60),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
client := &fakeBackupS3{
|
||||||
|
pages: []*s3.ListObjectsV2Output{
|
||||||
|
backupListPage(keys[:2], true, "page-2"),
|
||||||
|
backupListPage(keys[2:], false, ""),
|
||||||
|
},
|
||||||
|
objects: objects,
|
||||||
|
}
|
||||||
|
outDir := t.TempDir()
|
||||||
|
require.NoError(t, backupImagesWithClient(context.Background(), client, "configured-bucket", outDir, ""))
|
||||||
|
|
||||||
|
require.Len(t, client.listInputs, 2)
|
||||||
|
for index, input := range client.listInputs {
|
||||||
|
require.Equal(t, "configured-bucket", aws.ToString(input.Bucket))
|
||||||
|
require.Equal(t, "", aws.ToString(input.Prefix), "生产默认必须列举整个 bucket")
|
||||||
|
if index == 0 {
|
||||||
|
require.Nil(t, input.ContinuationToken)
|
||||||
|
} else {
|
||||||
|
require.Equal(t, "page-2", aws.ToString(input.ContinuationToken))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
require.ElementsMatch(t, keys, client.getOrder)
|
||||||
|
|
||||||
|
manifestBytes, err := os.ReadFile(filepath.Join(outDir, "images-manifest.json"))
|
||||||
|
require.NoError(t, err)
|
||||||
|
var manifest []imageManifestEntry
|
||||||
|
require.NoError(t, json.Unmarshal(manifestBytes, &manifest))
|
||||||
|
require.Len(t, manifest, len(keys))
|
||||||
|
for index, entry := range manifest {
|
||||||
|
require.Equal(t, keys[index], entry.Key)
|
||||||
|
object := objects[entry.Key]
|
||||||
|
require.Equal(t, int64(len(object.body)), entry.Size)
|
||||||
|
digest := sha256.Sum256(object.body)
|
||||||
|
require.Equal(t, fmt.Sprintf("%x", digest), entry.SHA256)
|
||||||
|
require.Equal(t, object.etag, entry.ETag)
|
||||||
|
require.Equal(t, object.contentType, entry.ContentType)
|
||||||
|
require.Equal(t, object.cacheControl, entry.CacheControl)
|
||||||
|
|
||||||
|
path, err := localImagePath(filepath.Join(outDir, "images"), entry.Key)
|
||||||
|
require.NoError(t, err)
|
||||||
|
content, err := os.ReadFile(path)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, object.body, content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBackupImagesWithClientSupportsIsolatedPrefixAndEmptyManifest(t *testing.T) {
|
||||||
|
t.Run("prefix", func(t *testing.T) {
|
||||||
|
key := "tamabox-backup-test/unit/source/root.webp"
|
||||||
|
client := &fakeBackupS3{
|
||||||
|
pages: []*s3.ListObjectsV2Output{backupListPage([]string{key}, false, "")},
|
||||||
|
objects: map[string]fakeBackupObject{
|
||||||
|
key: {body: []byte("isolated"), contentType: "image/webp"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
require.NoError(t, backupImagesWithClient(context.Background(), client, "bucket", t.TempDir(), "tamabox-backup-test/unit/source/"))
|
||||||
|
require.Len(t, client.listInputs, 1)
|
||||||
|
require.Equal(t, "tamabox-backup-test/unit/source/", aws.ToString(client.listInputs[0].Prefix))
|
||||||
|
require.Equal(t, []string{key}, client.getOrder)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("empty array", func(t *testing.T) {
|
||||||
|
client := &fakeBackupS3{pages: []*s3.ListObjectsV2Output{backupListPage(nil, false, "")}}
|
||||||
|
outDir := t.TempDir()
|
||||||
|
require.NoError(t, backupImagesWithClient(context.Background(), client, "bucket", outDir, ""))
|
||||||
|
content, err := os.ReadFile(filepath.Join(outDir, "images-manifest.json"))
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, "[]", strings.TrimSpace(string(content)))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBackupRejectsPathAttacksBeforeAnyDownload(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
key *string
|
||||||
|
}{
|
||||||
|
{name: "nil key", key: nil},
|
||||||
|
{name: "empty key", key: aws.String("")},
|
||||||
|
{name: "absolute POSIX", key: aws.String("/outside")},
|
||||||
|
{name: "absolute Windows", key: aws.String("C:/outside")},
|
||||||
|
{name: "volume relative", key: aws.String("C:outside")},
|
||||||
|
{name: "parent", key: aws.String("../outside")},
|
||||||
|
{name: "nested parent", key: aws.String("picture/../outside")},
|
||||||
|
{name: "backslash", key: aws.String(`picture\outside`)},
|
||||||
|
{name: "NUL", key: aws.String("picture/\x00outside")},
|
||||||
|
{name: "empty segment", key: aws.String("picture//outside")},
|
||||||
|
{name: "leading dot segment", key: aws.String("./outside")},
|
||||||
|
{name: "inner dot segment", key: aws.String("picture/./outside")},
|
||||||
|
{name: "directory", key: aws.String("picture/")},
|
||||||
|
{name: "dot only", key: aws.String(".")},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
page := backupListPage(nil, false, "")
|
||||||
|
page.Contents = []types.Object{{Key: tt.key}}
|
||||||
|
client := &fakeBackupS3{pages: []*s3.ListObjectsV2Output{page}}
|
||||||
|
err := backupImagesWithClient(context.Background(), client, "bucket", t.TempDir(), "")
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Empty(t, client.getOrder, "非法 key 必须在任何下载前被拒绝")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBackupPreflightsAllPagesAndLocalCollisionsBeforeDownload(t *testing.T) {
|
||||||
|
t.Run("invalid key on later page", func(t *testing.T) {
|
||||||
|
client := &fakeBackupS3{
|
||||||
|
pages: []*s3.ListObjectsV2Output{
|
||||||
|
backupListPage([]string{"picture/valid.webp"}, true, "next"),
|
||||||
|
backupListPage([]string{"../escape"}, false, ""),
|
||||||
|
},
|
||||||
|
objects: map[string]fakeBackupObject{"picture/valid.webp": {body: []byte("valid")}},
|
||||||
|
}
|
||||||
|
require.Error(t, backupImagesWithClient(context.Background(), client, "bucket", t.TempDir(), ""))
|
||||||
|
require.Len(t, client.listInputs, 2)
|
||||||
|
require.Empty(t, client.getOrder)
|
||||||
|
})
|
||||||
|
|
||||||
|
for _, keys := range [][]string{{"a", "a/b"}, {"a/b", "a"}, {"same", "same"}} {
|
||||||
|
t.Run(fmt.Sprintf("collision %q", keys), func(t *testing.T) {
|
||||||
|
client := &fakeBackupS3{pages: []*s3.ListObjectsV2Output{backupListPage(keys, false, "")}}
|
||||||
|
require.Error(t, backupImagesWithClient(context.Background(), client, "bucket", t.TempDir(), ""))
|
||||||
|
require.Empty(t, client.getOrder)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
t.Run("existing file blocks planned directory", func(t *testing.T) {
|
||||||
|
outDir := t.TempDir()
|
||||||
|
require.NoError(t, os.MkdirAll(filepath.Join(outDir, "images"), 0o700))
|
||||||
|
require.NoError(t, os.WriteFile(filepath.Join(outDir, "images", "a"), []byte("existing"), 0o600))
|
||||||
|
client := &fakeBackupS3{pages: []*s3.ListObjectsV2Output{backupListPage([]string{"a/b"}, false, "")}}
|
||||||
|
require.Error(t, backupImagesWithClient(context.Background(), client, "bucket", outDir, ""))
|
||||||
|
require.Empty(t, client.getOrder)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("existing directory blocks planned file", func(t *testing.T) {
|
||||||
|
outDir := t.TempDir()
|
||||||
|
require.NoError(t, os.MkdirAll(filepath.Join(outDir, "images", "a"), 0o700))
|
||||||
|
client := &fakeBackupS3{pages: []*s3.ListObjectsV2Output{backupListPage([]string{"a"}, false, "")}}
|
||||||
|
require.Error(t, backupImagesWithClient(context.Background(), client, "bucket", outDir, ""))
|
||||||
|
require.Empty(t, client.getOrder)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBackupRetriesInterruptedStreamsAndLeavesNoPart(t *testing.T) {
|
||||||
|
t.Run("eventual success", func(t *testing.T) {
|
||||||
|
key := "picture/retry.webp"
|
||||||
|
body := []byte("complete-object-after-retries")
|
||||||
|
client := &fakeBackupS3{
|
||||||
|
pages: []*s3.ListObjectsV2Output{backupListPage([]string{key}, false, "")},
|
||||||
|
objects: map[string]fakeBackupObject{
|
||||||
|
key: {body: body, interruptedAttempts: 2, etag: "successful-get-etag"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
outDir := t.TempDir()
|
||||||
|
require.NoError(t, backupImagesWithClient(context.Background(), client, "bucket", outDir, ""))
|
||||||
|
require.Equal(t, 3, client.getCalls[key])
|
||||||
|
content, err := os.ReadFile(filepath.Join(outDir, "images", "picture", "retry.webp"))
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, body, content)
|
||||||
|
requireNoPartFiles(t, filepath.Join(outDir, "images", "picture"))
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("short body relative to GET content length", func(t *testing.T) {
|
||||||
|
key := "short-body.webp"
|
||||||
|
client := &fakeBackupS3{
|
||||||
|
pages: []*s3.ListObjectsV2Output{backupListPage([]string{key}, false, "")},
|
||||||
|
objects: map[string]fakeBackupObject{
|
||||||
|
key: {body: []byte("complete-on-second-attempt"), wrongLengthAttempts: 1},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
outDir := t.TempDir()
|
||||||
|
require.NoError(t, backupImagesWithClient(context.Background(), client, "bucket", outDir, ""))
|
||||||
|
require.Equal(t, 2, client.getCalls[key])
|
||||||
|
requireNoPartFiles(t, filepath.Join(outDir, "images"))
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("final failure", func(t *testing.T) {
|
||||||
|
key := "retry.txt"
|
||||||
|
client := &fakeBackupS3{
|
||||||
|
pages: []*s3.ListObjectsV2Output{backupListPage([]string{key}, false, "")},
|
||||||
|
objects: map[string]fakeBackupObject{
|
||||||
|
key: {body: []byte("never-complete"), interruptedAttempts: 3},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
outDir := t.TempDir()
|
||||||
|
err := backupImagesWithClient(context.Background(), client, "bucket", outDir, "")
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Contains(t, err.Error(), "after 3 attempts")
|
||||||
|
require.Equal(t, 3, client.getCalls[key])
|
||||||
|
require.NoFileExists(t, filepath.Join(outDir, "images", key))
|
||||||
|
requireNoPartFiles(t, filepath.Join(outDir, "images"))
|
||||||
|
require.NoFileExists(t, filepath.Join(outDir, "images-manifest.json"))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
type capturedRestorePut struct {
|
||||||
|
key string
|
||||||
|
body []byte
|
||||||
|
contentType string
|
||||||
|
cacheControl string
|
||||||
|
size int64
|
||||||
|
}
|
||||||
|
|
||||||
|
type fakeRestoreStore struct {
|
||||||
|
existing map[string]bool
|
||||||
|
headErrors map[string]error
|
||||||
|
putFailures map[string]error
|
||||||
|
headCalls []string
|
||||||
|
puts []capturedRestorePut
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeRestoreStore) HeadObject(_ context.Context, input *s3.HeadObjectInput, _ ...func(*s3.Options)) (*s3.HeadObjectOutput, error) {
|
||||||
|
key := aws.ToString(input.Key)
|
||||||
|
f.headCalls = append(f.headCalls, key)
|
||||||
|
if err := f.headErrors[key]; err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if f.existing[key] {
|
||||||
|
return &s3.HeadObjectOutput{}, nil
|
||||||
|
}
|
||||||
|
return nil, &types.NotFound{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeRestoreStore) PutObject(_ context.Context, input *s3.PutObjectInput) error {
|
||||||
|
body, err := io.ReadAll(input.Body)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
key := aws.ToString(input.Key)
|
||||||
|
f.puts = append(f.puts, capturedRestorePut{
|
||||||
|
key: key,
|
||||||
|
body: body,
|
||||||
|
contentType: aws.ToString(input.ContentType),
|
||||||
|
cacheControl: aws.ToString(input.CacheControl),
|
||||||
|
size: aws.ToInt64(input.ContentLength),
|
||||||
|
})
|
||||||
|
return f.putFailures[key]
|
||||||
|
}
|
||||||
|
|
||||||
|
type restoreFixtureObject struct {
|
||||||
|
key string
|
||||||
|
body []byte
|
||||||
|
storedBody []byte
|
||||||
|
manifestSize *int64
|
||||||
|
manifestSHA string
|
||||||
|
contentType string
|
||||||
|
cacheControl string
|
||||||
|
}
|
||||||
|
|
||||||
|
func createRestoreFixture(t *testing.T, objects []restoreFixtureObject) string {
|
||||||
|
t.Helper()
|
||||||
|
dir := t.TempDir()
|
||||||
|
manifest := make([]imageManifestEntry, 0, len(objects))
|
||||||
|
for _, object := range objects {
|
||||||
|
storedBody := object.storedBody
|
||||||
|
if storedBody == nil {
|
||||||
|
storedBody = object.body
|
||||||
|
}
|
||||||
|
path := filepath.Join(dir, "images", filepath.FromSlash(object.key))
|
||||||
|
require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700))
|
||||||
|
require.NoError(t, os.WriteFile(path, storedBody, 0o600))
|
||||||
|
|
||||||
|
size := int64(len(object.body))
|
||||||
|
if object.manifestSize != nil {
|
||||||
|
size = *object.manifestSize
|
||||||
|
}
|
||||||
|
digest := sha256.Sum256(object.body)
|
||||||
|
hash := fmt.Sprintf("%x", digest)
|
||||||
|
if object.manifestSHA != "" {
|
||||||
|
hash = object.manifestSHA
|
||||||
|
}
|
||||||
|
manifest = append(manifest, imageManifestEntry{
|
||||||
|
Key: object.key,
|
||||||
|
Size: size,
|
||||||
|
SHA256: hash,
|
||||||
|
ETag: "source-get-etag",
|
||||||
|
ContentType: object.contentType,
|
||||||
|
CacheControl: object.cacheControl,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
data, err := json.MarshalIndent(manifest, "", " ")
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NoError(t, os.WriteFile(filepath.Join(dir, "images-manifest.json"), data, 0o600))
|
||||||
|
return dir
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRestoreImagesPreservesKeysContentAndMetadataWithTargetPrefix(t *testing.T) {
|
||||||
|
objects := []restoreFixtureObject{
|
||||||
|
{key: "2026/11/04/root.webp", body: []byte("root-body"), contentType: "image/webp", cacheControl: "public, max-age=60"},
|
||||||
|
{key: "picture/2026/11/04/pic.webp", body: []byte("picture-body"), contentType: "image/webp", cacheControl: "public, max-age=120"},
|
||||||
|
{key: "static/site.css", body: []byte("body{}"), contentType: "text/css", cacheControl: "no-cache"},
|
||||||
|
}
|
||||||
|
input := createRestoreFixture(t, objects)
|
||||||
|
store := &fakeRestoreStore{}
|
||||||
|
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||||
|
Input: input,
|
||||||
|
Yes: true,
|
||||||
|
TargetPrefix: "isolated-restore", // 不带斜杠也应规范化为恰好一个路径分隔符。
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Len(t, store.headCalls, len(objects))
|
||||||
|
require.Len(t, store.puts, len(objects))
|
||||||
|
|
||||||
|
byKey := make(map[string]capturedRestorePut, len(store.puts))
|
||||||
|
for _, put := range store.puts {
|
||||||
|
byKey[put.key] = put
|
||||||
|
}
|
||||||
|
for _, object := range objects {
|
||||||
|
target := "isolated-restore/" + object.key
|
||||||
|
put := byKey[target]
|
||||||
|
require.Equal(t, object.body, put.body)
|
||||||
|
require.Equal(t, object.contentType, put.contentType)
|
||||||
|
require.Equal(t, object.cacheControl, put.cacheControl)
|
||||||
|
require.Equal(t, int64(len(object.body)), put.size)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRestoreImagesDefaultRefusesOverwriteWithZeroPut(t *testing.T) {
|
||||||
|
objects := []restoreFixtureObject{
|
||||||
|
{key: "first.webp", body: []byte("first")},
|
||||||
|
{key: "second.webp", body: []byte("second")},
|
||||||
|
}
|
||||||
|
store := &fakeRestoreStore{existing: map[string]bool{"second.webp": true}}
|
||||||
|
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||||
|
Input: createRestoreFixture(t, objects),
|
||||||
|
Yes: true,
|
||||||
|
})
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Contains(t, err.Error(), "no objects were uploaded")
|
||||||
|
require.Empty(t, store.puts)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRestoreImagesOverwriteAndDryRun(t *testing.T) {
|
||||||
|
t.Run("overwrite", func(t *testing.T) {
|
||||||
|
store := &fakeRestoreStore{existing: map[string]bool{"picture.webp": true}}
|
||||||
|
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||||
|
Input: createRestoreFixture(t, []restoreFixtureObject{{key: "picture.webp", body: []byte("replacement")}}),
|
||||||
|
Yes: true,
|
||||||
|
Overwrite: true,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Empty(t, store.headCalls, "显式 overwrite 不需要默认防覆盖 Head")
|
||||||
|
require.Len(t, store.puts, 1)
|
||||||
|
require.Equal(t, []byte("replacement"), store.puts[0].body)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("dry run without yes", func(t *testing.T) {
|
||||||
|
store := &fakeRestoreStore{}
|
||||||
|
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||||
|
Input: createRestoreFixture(t, []restoreFixtureObject{{key: "picture.webp", body: []byte("dry-run")}}),
|
||||||
|
DryRun: true,
|
||||||
|
TargetPrefix: "drill/",
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, []string{"drill/picture.webp"}, store.headCalls)
|
||||||
|
require.Empty(t, store.puts)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("actual write requires yes", func(t *testing.T) {
|
||||||
|
store := &fakeRestoreStore{}
|
||||||
|
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||||
|
Input: createRestoreFixture(t, []restoreFixtureObject{{key: "picture.webp", body: []byte("write")}}),
|
||||||
|
})
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Contains(t, err.Error(), "--yes")
|
||||||
|
require.Empty(t, store.headCalls)
|
||||||
|
require.Empty(t, store.puts)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRestoreRejectsTamperedFilesBeforeAnyRemoteCall(t *testing.T) {
|
||||||
|
wrongSize := int64(999)
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
second restoreFixtureObject
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "SHA256 tamper with unchanged size",
|
||||||
|
second: restoreFixtureObject{
|
||||||
|
key: "second.webp", body: []byte("second"), storedBody: []byte("SECOND"),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "size tamper",
|
||||||
|
second: restoreFixtureObject{
|
||||||
|
key: "second.webp", body: []byte("second"), manifestSize: &wrongSize,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "manifest hash tamper",
|
||||||
|
second: restoreFixtureObject{
|
||||||
|
key: "second.webp", body: []byte("second"), manifestSHA: strings.Repeat("0", 64),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
store := &fakeRestoreStore{}
|
||||||
|
input := createRestoreFixture(t, []restoreFixtureObject{
|
||||||
|
{key: "first.webp", body: []byte("first")},
|
||||||
|
tt.second,
|
||||||
|
})
|
||||||
|
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{Input: input, Yes: true})
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Empty(t, store.headCalls, "全部本地文件应在任何远端操作前完成校验")
|
||||||
|
require.Empty(t, store.puts, "首个有效文件也不能在后续文件校验失败前上传")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReadImageManifestRejectsInvalidStructure(t *testing.T) {
|
||||||
|
validHash := strings.Repeat("a", 64)
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
json string
|
||||||
|
}{
|
||||||
|
{name: "null", json: `null`},
|
||||||
|
{name: "trailing JSON", json: `[] {}`},
|
||||||
|
{name: "trailing garbage", json: `[] garbage`},
|
||||||
|
{name: "duplicate key", json: fmt.Sprintf(`[{"key":"same","size":0,"sha256":%q},{"key":"same","size":0,"sha256":%q}]`, validHash, validHash)},
|
||||||
|
{name: "short hash", json: `[{"key":"image","size":1,"sha256":"abc"}]`},
|
||||||
|
{name: "non hex hash", json: fmt.Sprintf(`[{"key":"image","size":1,"sha256":%q}]`, strings.Repeat("z", 64))},
|
||||||
|
{name: "negative size", json: fmt.Sprintf(`[{"key":"image","size":-1,"sha256":%q}]`, validHash)},
|
||||||
|
{name: "path traversal", json: fmt.Sprintf(`[{"key":"../image","size":1,"sha256":%q}]`, validHash)},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
path := filepath.Join(t.TempDir(), "images-manifest.json")
|
||||||
|
// 非数组用例直接保留原 JSON;其余用例同样应被解析层拒绝。
|
||||||
|
require.NoError(t, os.WriteFile(path, []byte(tt.json), 0o600))
|
||||||
|
_, err := readImageManifest(path)
|
||||||
|
require.Error(t, err)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRestoreRejectsSymlinkBeforeRemoteCalls(t *testing.T) {
|
||||||
|
input := createRestoreFixture(t, []restoreFixtureObject{{key: "picture.webp", body: []byte("content")}})
|
||||||
|
source := filepath.Join(input, "images", "picture.webp")
|
||||||
|
outside := filepath.Join(t.TempDir(), "outside.webp")
|
||||||
|
require.NoError(t, os.WriteFile(outside, []byte("content"), 0o600))
|
||||||
|
require.NoError(t, os.Remove(source))
|
||||||
|
if err := os.Symlink(outside, source); err != nil {
|
||||||
|
t.Skipf("当前平台无法创建 symlink:%v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
store := &fakeRestoreStore{}
|
||||||
|
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{Input: input, DryRun: true})
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Empty(t, store.headCalls)
|
||||||
|
require.Empty(t, store.puts)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRestoreReportsPartialFailureWithoutClaimingTransaction(t *testing.T) {
|
||||||
|
store := &fakeRestoreStore{putFailures: map[string]error{"second.webp": errors.New("injected put failure")}}
|
||||||
|
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||||
|
Input: createRestoreFixture(t, []restoreFixtureObject{
|
||||||
|
{key: "first.webp", body: []byte("first")},
|
||||||
|
{key: "second.webp", body: []byte("second")},
|
||||||
|
}),
|
||||||
|
Yes: true,
|
||||||
|
})
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Contains(t, err.Error(), "1 successful uploads")
|
||||||
|
require.Contains(t, err.Error(), "not transactional")
|
||||||
|
require.Len(t, store.puts, 2)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRestoreRejectsUnsafeTargetPrefixBeforeRemoteCalls(t *testing.T) {
|
||||||
|
store := &fakeRestoreStore{}
|
||||||
|
err := restoreImagesWithStore(context.Background(), store, "bucket", restoreImagesOptions{
|
||||||
|
Input: createRestoreFixture(t, []restoreFixtureObject{{key: "picture.webp", body: []byte("content")}}),
|
||||||
|
Yes: true,
|
||||||
|
TargetPrefix: "../outside/",
|
||||||
|
})
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Empty(t, store.headCalls)
|
||||||
|
require.Empty(t, store.puts)
|
||||||
|
}
|
||||||
@@ -0,0 +1,96 @@
|
|||||||
|
package cmd
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"flag"
|
||||||
|
"io"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
"github.com/urfave/cli/v2"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestRunBackupRejectsMutuallyExclusiveModes(t *testing.T) {
|
||||||
|
set := flag.NewFlagSet("backup", flag.ContinueOnError)
|
||||||
|
set.Bool("db-only", true, "")
|
||||||
|
set.Bool("images-only", true, "")
|
||||||
|
|
||||||
|
err := runBackup(cli.NewContext(nil, set, nil))
|
||||||
|
require.EqualError(t, err, "--db-only and --images-only cannot be used together")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWriteFileAtomically(t *testing.T) {
|
||||||
|
t.Run("success", func(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
target := filepath.Join(dir, "backup.sql")
|
||||||
|
require.NoError(t, writeFileAtomically(target, func(w io.Writer) error {
|
||||||
|
_, err := io.WriteString(w, "complete backup")
|
||||||
|
return err
|
||||||
|
}))
|
||||||
|
|
||||||
|
content, err := os.ReadFile(target)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, "complete backup", string(content))
|
||||||
|
requireNoPartFiles(t, dir)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("failure preserves previous backup", func(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
target := filepath.Join(dir, "backup.sql")
|
||||||
|
require.NoError(t, os.WriteFile(target, []byte("known good"), 0o600))
|
||||||
|
writeErr := errors.New("dump failed")
|
||||||
|
|
||||||
|
err := writeFileAtomically(target, func(w io.Writer) error {
|
||||||
|
_, _ = io.WriteString(w, "partial")
|
||||||
|
return writeErr
|
||||||
|
})
|
||||||
|
require.ErrorIs(t, err, writeErr)
|
||||||
|
|
||||||
|
content, err := os.ReadFile(target)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, "known good", string(content))
|
||||||
|
requireNoPartFiles(t, dir)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("failure leaves no new backup", func(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
target := filepath.Join(dir, "backup.sql")
|
||||||
|
writeErr := errors.New("dump failed")
|
||||||
|
|
||||||
|
err := writeFileAtomically(target, func(w io.Writer) error {
|
||||||
|
_, _ = io.WriteString(w, "partial")
|
||||||
|
return writeErr
|
||||||
|
})
|
||||||
|
require.ErrorIs(t, err, writeErr)
|
||||||
|
require.NoFileExists(t, target)
|
||||||
|
requireNoPartFiles(t, dir)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func requireNoPartFiles(t *testing.T, dir string) {
|
||||||
|
t.Helper()
|
||||||
|
matches, err := filepath.Glob(filepath.Join(dir, "*.part"))
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Empty(t, matches)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestImageManifestEntryIntegrityFields(t *testing.T) {
|
||||||
|
entry := imageManifestEntry{
|
||||||
|
Key: "tamabox/image.webp",
|
||||||
|
Size: 3,
|
||||||
|
SHA256: "actual-content-sha256",
|
||||||
|
ETag: "provider-etag",
|
||||||
|
}
|
||||||
|
encoded, err := json.Marshal(entry)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.JSONEq(t, `{
|
||||||
|
"key":"tamabox/image.webp",
|
||||||
|
"size":3,
|
||||||
|
"sha256":"actual-content-sha256",
|
||||||
|
"etag":"provider-etag"
|
||||||
|
}`, string(encoded))
|
||||||
|
require.NotContains(t, string(encoded), "md5")
|
||||||
|
}
|
||||||
+3
-1
@@ -95,7 +95,9 @@ func runBan(ctx *cli.Context) error {
|
|||||||
return errors.Wrap(err, "ban user and cleanup")
|
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
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,365 @@
|
|||||||
|
// Copyright 2026 tamakyi. All rights reserved.
|
||||||
|
// Use of this source code is governed by a MIT-style
|
||||||
|
// license that can be found in the LICENSE file.
|
||||||
|
|
||||||
|
package cmd
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
"unicode/utf8"
|
||||||
|
|
||||||
|
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||||
|
"github.com/sirupsen/logrus"
|
||||||
|
"github.com/urfave/cli/v2"
|
||||||
|
|
||||||
|
"github.com/tamakyi/TamaBox/internal/conf"
|
||||||
|
"github.com/tamakyi/TamaBox/internal/db"
|
||||||
|
"github.com/tamakyi/TamaBox/internal/storage"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
r2OutboxDefaultLimit = 100
|
||||||
|
r2OutboxWorkerInterval = time.Minute
|
||||||
|
r2OutboxInitialBackoff = time.Minute
|
||||||
|
r2OutboxMaxBackoff = time.Hour
|
||||||
|
r2OutboxMaxLastErrorSize = 1024
|
||||||
|
)
|
||||||
|
|
||||||
|
var errR2OutboxNotConfigured = errors.New("R2 configuration is incomplete")
|
||||||
|
|
||||||
|
// R2Outbox 有界处理一批持久化的 R2 对象删除任务。
|
||||||
|
var R2Outbox = &cli.Command{
|
||||||
|
Name: "r2-outbox",
|
||||||
|
Usage: "Drain a bounded batch of pending R2 object deletion tasks",
|
||||||
|
Action: runR2Outbox,
|
||||||
|
Flags: []cli.Flag{
|
||||||
|
&cli.IntFlag{
|
||||||
|
Name: "limit",
|
||||||
|
Value: r2OutboxDefaultLimit,
|
||||||
|
Usage: "Maximum number of due tasks to process",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
type r2OutboxTaskStore interface {
|
||||||
|
ListDue(ctx context.Context, limit int) ([]*db.R2DeleteTask, error)
|
||||||
|
Delete(id uint) error
|
||||||
|
MarkFailed(id uint, nextRetry time.Time, lastError string) error
|
||||||
|
HasReferences(ctx context.Context, bucket, key string) (bool, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
type r2ObjectDeleter interface {
|
||||||
|
DeleteObject(ctx context.Context, bucket, key string) error
|
||||||
|
}
|
||||||
|
|
||||||
|
type r2S3ObjectDeleter struct {
|
||||||
|
client *s3.Client
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *r2S3ObjectDeleter) DeleteObject(ctx context.Context, bucket, key string) error {
|
||||||
|
return storage.DeleteObject(ctx, d.client, bucket, key)
|
||||||
|
}
|
||||||
|
|
||||||
|
type r2OutboxResult struct {
|
||||||
|
Listed int
|
||||||
|
Removed int
|
||||||
|
ObjectsDeleted int
|
||||||
|
ReferenceOnly int
|
||||||
|
Failed int
|
||||||
|
Skipped int
|
||||||
|
Remaining int
|
||||||
|
}
|
||||||
|
|
||||||
|
// processR2Outbox 串行处理一批到期任务。单个任务失败不会中断同批后续任务;
|
||||||
|
// context 取消则立即停止,且不会把关闭过程中的取消误记为一次任务失败。
|
||||||
|
func processR2Outbox(ctx context.Context, store r2OutboxTaskStore, deleter r2ObjectDeleter, limit int) (result r2OutboxResult, err error) {
|
||||||
|
defer func() {
|
||||||
|
result.Remaining = result.Listed - result.Removed
|
||||||
|
}()
|
||||||
|
|
||||||
|
tasks, err := store.ListDue(ctx, limit)
|
||||||
|
if err != nil {
|
||||||
|
if ctx.Err() != nil {
|
||||||
|
return result, ctx.Err()
|
||||||
|
}
|
||||||
|
return result, fmt.Errorf("list R2 outbox tasks: %s", sanitizeR2OutboxError(err))
|
||||||
|
}
|
||||||
|
result.Listed = len(tasks)
|
||||||
|
|
||||||
|
for _, task := range tasks {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return result, err
|
||||||
|
}
|
||||||
|
if task == nil {
|
||||||
|
result.Failed++
|
||||||
|
logrus.Warn("R2 outbox 返回了空任务,已跳过")
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
logger := logrus.WithField("task_id", task.ID)
|
||||||
|
if task.Bucket != conf.Upload.ImageBucket {
|
||||||
|
result.Skipped++
|
||||||
|
logger.Warn("R2 outbox 任务 bucket 与当前配置不匹配,任务已保留")
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
referenced, refErr := store.HasReferences(ctx, task.Bucket, task.Key)
|
||||||
|
if ctx.Err() != nil {
|
||||||
|
return result, ctx.Err()
|
||||||
|
}
|
||||||
|
if refErr != nil {
|
||||||
|
result.Failed++
|
||||||
|
markR2OutboxTaskFailed(store, task, refErr)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if referenced {
|
||||||
|
// 还有引用时只清理过期 outbox 记录,绝不删除对象。最后一个引用被
|
||||||
|
// 删除时,对应数据库事务会重新入队一条任务。
|
||||||
|
if err := store.Delete(task.ID); err != nil {
|
||||||
|
result.Failed++
|
||||||
|
logger.WithField("error", sanitizeR2OutboxError(err)).Warn("R2 outbox 引用任务清理失败,任务已保留")
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
result.Removed++
|
||||||
|
result.ReferenceOnly++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
deleteErr := deleter.DeleteObject(ctx, task.Bucket, task.Key)
|
||||||
|
if deleteErr != nil && !isS3ObjectNotFound(deleteErr) {
|
||||||
|
if ctx.Err() != nil {
|
||||||
|
return result, ctx.Err()
|
||||||
|
}
|
||||||
|
result.Failed++
|
||||||
|
markR2OutboxTaskFailed(store, task, deleteErr)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// storage.DeleteObject 已将 404 归一为成功;这里再次识别 404,保证其他
|
||||||
|
// deleter 实现也遵循相同的幂等语义。只有任务记录也删除后才算完成。
|
||||||
|
if err := store.Delete(task.ID); err != nil {
|
||||||
|
result.Failed++
|
||||||
|
logger.WithField("error", sanitizeR2OutboxError(err)).Warn("R2 对象已删除,但 outbox 任务清理失败;后续将幂等重试")
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
result.Removed++
|
||||||
|
result.ObjectsDeleted++
|
||||||
|
}
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func markR2OutboxTaskFailed(store r2OutboxTaskStore, task *db.R2DeleteTask, cause error) {
|
||||||
|
nextRetry := time.Now().Add(r2OutboxBackoff(task.Attempts))
|
||||||
|
lastError := sanitizeR2OutboxError(cause)
|
||||||
|
if err := store.MarkFailed(task.ID, nextRetry, lastError); err != nil {
|
||||||
|
logrus.WithFields(logrus.Fields{
|
||||||
|
"task_id": task.ID,
|
||||||
|
"processing_error": lastError,
|
||||||
|
"persist_error": sanitizeR2OutboxError(err),
|
||||||
|
}).Warn("R2 outbox 任务失败且重试状态写入失败,任务已保留")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
logrus.WithFields(logrus.Fields{
|
||||||
|
"task_id": task.ID,
|
||||||
|
"error": lastError,
|
||||||
|
"next_retry": nextRetry,
|
||||||
|
}).Warn("R2 outbox 任务处理失败,已安排重试")
|
||||||
|
}
|
||||||
|
|
||||||
|
func r2OutboxBackoff(attempts int) time.Duration {
|
||||||
|
if attempts < 0 {
|
||||||
|
attempts = 0
|
||||||
|
}
|
||||||
|
// 1、2、4、8、16、32 分钟后,下一次本应到 64 分钟,直接封顶一小时。
|
||||||
|
if attempts >= 6 {
|
||||||
|
return r2OutboxMaxBackoff
|
||||||
|
}
|
||||||
|
return r2OutboxInitialBackoff << attempts
|
||||||
|
}
|
||||||
|
|
||||||
|
// sanitizeR2OutboxError 在持久化或记录错误前移除已知配置与凭据,并限制为
|
||||||
|
// 1024 字节。截断保持 UTF-8 边界,避免向 text 列写入半个多字节字符。
|
||||||
|
func sanitizeR2OutboxError(err error) string {
|
||||||
|
if err == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
message := err.Error()
|
||||||
|
sensitiveValues := []string{
|
||||||
|
conf.Database.DSN,
|
||||||
|
conf.Database.Password,
|
||||||
|
conf.Upload.ImageEndpoint,
|
||||||
|
conf.Upload.ImageAccessID,
|
||||||
|
conf.Upload.ImageAccessSecret,
|
||||||
|
conf.Upload.ImageBucket,
|
||||||
|
conf.Upload.ImageBucketCDNHost,
|
||||||
|
}
|
||||||
|
for _, value := range sensitiveValues {
|
||||||
|
if value != "" {
|
||||||
|
message = strings.ReplaceAll(message, value, "[redacted]")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(message) <= r2OutboxMaxLastErrorSize {
|
||||||
|
return message
|
||||||
|
}
|
||||||
|
end := r2OutboxMaxLastErrorSize
|
||||||
|
for end > 0 && !utf8.ValidString(message[:end]) {
|
||||||
|
end--
|
||||||
|
}
|
||||||
|
return message[:end]
|
||||||
|
}
|
||||||
|
|
||||||
|
func r2OutboxConfigured() bool {
|
||||||
|
return strings.TrimSpace(conf.Upload.ImageEndpoint) != "" &&
|
||||||
|
strings.TrimSpace(conf.Upload.ImageAccessID) != "" &&
|
||||||
|
strings.TrimSpace(conf.Upload.ImageAccessSecret) != "" &&
|
||||||
|
strings.TrimSpace(conf.Upload.ImageBucket) != ""
|
||||||
|
}
|
||||||
|
|
||||||
|
func newR2OutboxObjectDeleter(ctx context.Context) (r2ObjectDeleter, error) {
|
||||||
|
if !r2OutboxConfigured() {
|
||||||
|
return nil, errR2OutboxNotConfigured
|
||||||
|
}
|
||||||
|
client, err := storage.NewS3Client(
|
||||||
|
ctx,
|
||||||
|
conf.Upload.ImageEndpoint,
|
||||||
|
conf.Upload.ImageBucket,
|
||||||
|
conf.Upload.ImageAccessID,
|
||||||
|
conf.Upload.ImageAccessSecret,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("initialize R2 client: %s", sanitizeR2OutboxError(err))
|
||||||
|
}
|
||||||
|
return &r2S3ObjectDeleter{client: client}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// startR2OutboxWorker 只能在 db.Init 成功后调用。goroutine 由调用方提供的
|
||||||
|
// WaitGroup 跟踪,Web 退出时先取消 ctx,再等待正在执行的一轮收尾。
|
||||||
|
func startR2OutboxWorker(ctx context.Context, wg *sync.WaitGroup) {
|
||||||
|
if !r2OutboxConfigured() {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
deleter, err := newR2OutboxObjectDeleter(ctx)
|
||||||
|
if err != nil {
|
||||||
|
logrus.WithField("error", sanitizeR2OutboxError(err)).Warn("R2 outbox worker 启动失败,Web 服务继续运行")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
runR2OutboxWorker(ctx, db.R2DeleteTasks, deleter)
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
func runR2OutboxWorker(ctx context.Context, store r2OutboxTaskStore, deleter r2ObjectDeleter) {
|
||||||
|
ticker := time.NewTicker(r2OutboxWorkerInterval)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
for {
|
||||||
|
// 启动先立即跑一轮;处理期间不启动并发轮次。若一轮超过一分钟,
|
||||||
|
// ticker 只会合并唤醒,不会形成多个 worker 同时消费。
|
||||||
|
_, err := processR2Outbox(ctx, store, deleter, r2OutboxDefaultLimit)
|
||||||
|
if err != nil {
|
||||||
|
if ctx.Err() != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
logrus.WithField("error", sanitizeR2OutboxError(err)).Warn("R2 outbox 本轮处理失败,将在下一分钟重试")
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return
|
||||||
|
case <-ticker.C:
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func drainR2Outbox(ctx context.Context, limit int) (r2OutboxResult, error) {
|
||||||
|
deleter, err := newR2OutboxObjectDeleter(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return r2OutboxResult{}, err
|
||||||
|
}
|
||||||
|
return processR2Outbox(ctx, db.R2DeleteTasks, deleter, limit)
|
||||||
|
}
|
||||||
|
|
||||||
|
func runR2Outbox(ctx *cli.Context) error {
|
||||||
|
limit := ctx.Int("limit")
|
||||||
|
if limit <= 0 {
|
||||||
|
return errors.New("--limit must be greater than zero")
|
||||||
|
}
|
||||||
|
if err := conf.Init(); err != nil {
|
||||||
|
return fmt.Errorf("load configuration: %w", err)
|
||||||
|
}
|
||||||
|
if !r2OutboxConfigured() {
|
||||||
|
return errR2OutboxNotConfigured
|
||||||
|
}
|
||||||
|
|
||||||
|
dbType := conf.Database.Type
|
||||||
|
var dsn string
|
||||||
|
switch dbType {
|
||||||
|
case "mysql", "":
|
||||||
|
dsn = conf.MySQLDsn()
|
||||||
|
case "postgres":
|
||||||
|
dsn = conf.PostgresDsn()
|
||||||
|
default:
|
||||||
|
return fmt.Errorf("unknown database type: %q", dbType)
|
||||||
|
}
|
||||||
|
conf.Database.DSN = dsn
|
||||||
|
gdb, err := db.Init(dbType, dsn)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("connect to database: %w", err)
|
||||||
|
}
|
||||||
|
if sqlDB, dbErr := gdb.DB(); dbErr == nil {
|
||||||
|
defer func() {
|
||||||
|
if err := sqlDB.Close(); err != nil {
|
||||||
|
logrus.WithField("error", sanitizeR2OutboxError(err)).Warn("关闭数据库连接池失败")
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
result, err := drainR2Outbox(ctx.Context, limit)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("drain R2 outbox: %s", sanitizeR2OutboxError(err))
|
||||||
|
}
|
||||||
|
logR2OutboxResult("R2 outbox 本轮处理完成", result)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func logR2OutboxResult(message string, result r2OutboxResult) {
|
||||||
|
entry := logrus.WithFields(logrus.Fields{
|
||||||
|
"listed": result.Listed,
|
||||||
|
"removed": result.Removed,
|
||||||
|
"object_deletes": result.ObjectsDeleted,
|
||||||
|
"reference_only": result.ReferenceOnly,
|
||||||
|
"failed": result.Failed,
|
||||||
|
"skipped": result.Skipped,
|
||||||
|
"remaining": result.Remaining,
|
||||||
|
})
|
||||||
|
if result.Remaining > 0 {
|
||||||
|
entry.Warn(message + ",仍有任务待重试或保留")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
entry.Info(message)
|
||||||
|
}
|
||||||
|
|
||||||
|
// drainR2OutboxAfterBan 在封禁事务提交后执行。R2 失败只影响对象清理进度,
|
||||||
|
// 不向 runBan 返回错误,避免运维误以为需要重新执行不可逆的数据库封禁。
|
||||||
|
func drainR2OutboxAfterBan(ctx context.Context) {
|
||||||
|
result, err := drainR2Outbox(ctx, r2OutboxDefaultLimit)
|
||||||
|
if err != nil {
|
||||||
|
logrus.WithField("error", sanitizeR2OutboxError(err)).Warn(
|
||||||
|
"数据库封禁与清理已提交,但 R2 outbox 本轮未能处理;任务仍保留,无需重跑 ban,可稍后运行 r2-outbox",
|
||||||
|
)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if result.Remaining > 0 {
|
||||||
|
logR2OutboxResult("数据库封禁与清理已提交;R2 对象清理将在后续重试,无需重跑 ban", result)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
logR2OutboxResult("数据库封禁与清理已提交;R2 outbox 本轮处理完成", result)
|
||||||
|
}
|
||||||
@@ -0,0 +1,270 @@
|
|||||||
|
package cmd
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
|
||||||
|
"github.com/tamakyi/TamaBox/internal/conf"
|
||||||
|
"github.com/tamakyi/TamaBox/internal/db"
|
||||||
|
)
|
||||||
|
|
||||||
|
type fakeR2OutboxFailure struct {
|
||||||
|
nextRetry time.Time
|
||||||
|
lastError string
|
||||||
|
}
|
||||||
|
|
||||||
|
type fakeR2OutboxStore struct {
|
||||||
|
tasks []*db.R2DeleteTask
|
||||||
|
references map[string]bool
|
||||||
|
referenceErrs map[string]error
|
||||||
|
referenceCalls []string
|
||||||
|
deleted map[uint]bool
|
||||||
|
deleteErr error
|
||||||
|
failures map[uint]fakeR2OutboxFailure
|
||||||
|
markFailedErr error
|
||||||
|
}
|
||||||
|
|
||||||
|
func newFakeR2OutboxStore(tasks ...*db.R2DeleteTask) *fakeR2OutboxStore {
|
||||||
|
return &fakeR2OutboxStore{
|
||||||
|
tasks: tasks,
|
||||||
|
references: make(map[string]bool),
|
||||||
|
referenceErrs: make(map[string]error),
|
||||||
|
deleted: make(map[uint]bool),
|
||||||
|
failures: make(map[uint]fakeR2OutboxFailure),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *fakeR2OutboxStore) ListDue(ctx context.Context, limit int) ([]*db.R2DeleteTask, error) {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
tasks := s.tasks
|
||||||
|
if limit >= 0 && len(tasks) > limit {
|
||||||
|
tasks = tasks[:limit]
|
||||||
|
}
|
||||||
|
return append([]*db.R2DeleteTask(nil), tasks...), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *fakeR2OutboxStore) Delete(id uint) error {
|
||||||
|
if s.deleteErr != nil {
|
||||||
|
return s.deleteErr
|
||||||
|
}
|
||||||
|
s.deleted[id] = true
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *fakeR2OutboxStore) MarkFailed(id uint, nextRetry time.Time, lastError string) error {
|
||||||
|
if s.markFailedErr != nil {
|
||||||
|
return s.markFailedErr
|
||||||
|
}
|
||||||
|
s.failures[id] = fakeR2OutboxFailure{nextRetry: nextRetry, lastError: lastError}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *fakeR2OutboxStore) HasReferences(ctx context.Context, _, key string) (bool, error) {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return false, err
|
||||||
|
}
|
||||||
|
s.referenceCalls = append(s.referenceCalls, key)
|
||||||
|
return s.references[key], s.referenceErrs[key]
|
||||||
|
}
|
||||||
|
|
||||||
|
type fakeR2DeleteCall struct {
|
||||||
|
bucket string
|
||||||
|
key string
|
||||||
|
}
|
||||||
|
|
||||||
|
type fakeR2ObjectDeleter struct {
|
||||||
|
calls []fakeR2DeleteCall
|
||||||
|
errs map[string]error
|
||||||
|
onDelete func(context.Context, string, string) error
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *fakeR2ObjectDeleter) DeleteObject(ctx context.Context, bucket, key string) error {
|
||||||
|
d.calls = append(d.calls, fakeR2DeleteCall{bucket: bucket, key: key})
|
||||||
|
if d.onDelete != nil {
|
||||||
|
return d.onDelete(ctx, bucket, key)
|
||||||
|
}
|
||||||
|
return d.errs[key]
|
||||||
|
}
|
||||||
|
|
||||||
|
type fakeR2HTTPStatusError int
|
||||||
|
|
||||||
|
func (e fakeR2HTTPStatusError) Error() string {
|
||||||
|
return fmt.Sprintf("R2 returned HTTP %d", int(e))
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e fakeR2HTTPStatusError) HTTPStatusCode() int {
|
||||||
|
return int(e)
|
||||||
|
}
|
||||||
|
|
||||||
|
func setR2OutboxTestConfig(t *testing.T) {
|
||||||
|
t.Helper()
|
||||||
|
old := conf.Upload
|
||||||
|
conf.Upload.ImageEndpoint = "https://r2.example.test"
|
||||||
|
conf.Upload.ImageAccessID = "test-access-id"
|
||||||
|
conf.Upload.ImageAccessSecret = "test-secret"
|
||||||
|
conf.Upload.ImageBucket = "configured-bucket"
|
||||||
|
conf.Upload.ImageBucketCDNHost = "cdn.example.test"
|
||||||
|
t.Cleanup(func() {
|
||||||
|
conf.Upload = old
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProcessR2OutboxSuccessAndNotFound(t *testing.T) {
|
||||||
|
setR2OutboxTestConfig(t)
|
||||||
|
for _, tt := range []struct {
|
||||||
|
name string
|
||||||
|
deleteErr error
|
||||||
|
}{
|
||||||
|
{name: "success"},
|
||||||
|
{name: "not found is success", deleteErr: fakeR2HTTPStatusError(404)},
|
||||||
|
} {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
first := &db.R2DeleteTask{ID: 1, Bucket: conf.Upload.ImageBucket, Key: "first.webp"}
|
||||||
|
second := &db.R2DeleteTask{ID: 2, Bucket: conf.Upload.ImageBucket, Key: "second.webp"}
|
||||||
|
store := newFakeR2OutboxStore(first, second)
|
||||||
|
deleter := &fakeR2ObjectDeleter{errs: map[string]error{"second.webp": tt.deleteErr}}
|
||||||
|
|
||||||
|
result, err := processR2Outbox(context.Background(), store, deleter, 100)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, 2, result.Listed)
|
||||||
|
require.Equal(t, 2, result.Removed)
|
||||||
|
require.Equal(t, 2, result.ObjectsDeleted)
|
||||||
|
require.Zero(t, result.Failed)
|
||||||
|
require.Zero(t, result.Remaining)
|
||||||
|
require.Equal(t, []fakeR2DeleteCall{
|
||||||
|
{bucket: conf.Upload.ImageBucket, key: "first.webp"},
|
||||||
|
{bucket: conf.Upload.ImageBucket, key: "second.webp"},
|
||||||
|
}, deleter.calls, "任务必须按 ListDue 顺序串行处理")
|
||||||
|
require.True(t, store.deleted[first.ID])
|
||||||
|
require.True(t, store.deleted[second.ID])
|
||||||
|
require.Empty(t, store.failures)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProcessR2OutboxReferencedTaskDeletesNoObject(t *testing.T) {
|
||||||
|
setR2OutboxTestConfig(t)
|
||||||
|
task := &db.R2DeleteTask{ID: 11, Bucket: conf.Upload.ImageBucket, Key: "shared.webp"}
|
||||||
|
store := newFakeR2OutboxStore(task)
|
||||||
|
store.references[task.Key] = true
|
||||||
|
deleter := &fakeR2ObjectDeleter{}
|
||||||
|
|
||||||
|
result, err := processR2Outbox(context.Background(), store, deleter, 100)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Empty(t, deleter.calls, "仍有引用时绝不能调用 R2 删除")
|
||||||
|
require.True(t, store.deleted[task.ID], "过期任务本身应删除,最后引用移除时会重新入队")
|
||||||
|
require.Equal(t, 1, result.ReferenceOnly)
|
||||||
|
require.Equal(t, 1, result.Removed)
|
||||||
|
require.Zero(t, result.ObjectsDeleted)
|
||||||
|
require.Zero(t, result.Remaining)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProcessR2OutboxReferenceErrorSchedulesRetry(t *testing.T) {
|
||||||
|
setR2OutboxTestConfig(t)
|
||||||
|
task := &db.R2DeleteTask{ID: 21, Bucket: conf.Upload.ImageBucket, Key: "unknown.webp"}
|
||||||
|
store := newFakeR2OutboxStore(task)
|
||||||
|
store.referenceErrs[task.Key] = errors.New("reference query failed")
|
||||||
|
deleter := &fakeR2ObjectDeleter{}
|
||||||
|
before := time.Now()
|
||||||
|
|
||||||
|
result, err := processR2Outbox(context.Background(), store, deleter, 100)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Empty(t, deleter.calls)
|
||||||
|
require.False(t, store.deleted[task.ID])
|
||||||
|
require.Contains(t, store.failures, task.ID)
|
||||||
|
failure := store.failures[task.ID]
|
||||||
|
require.Equal(t, "reference query failed", failure.lastError)
|
||||||
|
require.WithinDuration(t, before.Add(time.Minute), failure.nextRetry, 2*time.Second)
|
||||||
|
require.Equal(t, 1, result.Failed)
|
||||||
|
require.Equal(t, 1, result.Remaining)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProcessR2OutboxDeleteFailureUsesBackoffAndBoundsLastError(t *testing.T) {
|
||||||
|
setR2OutboxTestConfig(t)
|
||||||
|
task := &db.R2DeleteTask{ID: 31, Bucket: conf.Upload.ImageBucket, Key: "retry.webp", Attempts: 2}
|
||||||
|
store := newFakeR2OutboxStore(task)
|
||||||
|
longErr := fmt.Errorf("delete failed endpoint=%s credential=%s %s",
|
||||||
|
conf.Upload.ImageEndpoint,
|
||||||
|
conf.Upload.ImageAccessSecret,
|
||||||
|
strings.Repeat("x", 1100),
|
||||||
|
)
|
||||||
|
deleter := &fakeR2ObjectDeleter{errs: map[string]error{task.Key: longErr}}
|
||||||
|
before := time.Now()
|
||||||
|
|
||||||
|
result, err := processR2Outbox(context.Background(), store, deleter, 100)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Len(t, deleter.calls, 1)
|
||||||
|
require.False(t, store.deleted[task.ID])
|
||||||
|
require.Contains(t, store.failures, task.ID)
|
||||||
|
failure := store.failures[task.ID]
|
||||||
|
require.WithinDuration(t, before.Add(4*time.Minute), failure.nextRetry, 2*time.Second)
|
||||||
|
require.LessOrEqual(t, len(failure.lastError), r2OutboxMaxLastErrorSize)
|
||||||
|
require.NotContains(t, failure.lastError, conf.Upload.ImageEndpoint)
|
||||||
|
require.NotContains(t, failure.lastError, conf.Upload.ImageAccessSecret)
|
||||||
|
require.Equal(t, 1, result.Failed)
|
||||||
|
require.Equal(t, 1, result.Remaining)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProcessR2OutboxBucketMismatchIsRetainedWithoutDeletion(t *testing.T) {
|
||||||
|
setR2OutboxTestConfig(t)
|
||||||
|
task := &db.R2DeleteTask{ID: 41, Bucket: "different-bucket", Key: "foreign.webp"}
|
||||||
|
store := newFakeR2OutboxStore(task)
|
||||||
|
deleter := &fakeR2ObjectDeleter{}
|
||||||
|
|
||||||
|
result, err := processR2Outbox(context.Background(), store, deleter, 100)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Empty(t, store.referenceCalls, "bucket 不匹配时不应查询当前 bucket 的引用")
|
||||||
|
require.Empty(t, deleter.calls)
|
||||||
|
require.Empty(t, store.failures)
|
||||||
|
require.False(t, store.deleted[task.ID])
|
||||||
|
require.Equal(t, 1, result.Skipped)
|
||||||
|
require.Zero(t, result.Failed)
|
||||||
|
require.Equal(t, 1, result.Remaining)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProcessR2OutboxContextCancellationStopsWithoutMarkFailed(t *testing.T) {
|
||||||
|
setR2OutboxTestConfig(t)
|
||||||
|
first := &db.R2DeleteTask{ID: 51, Bucket: conf.Upload.ImageBucket, Key: "first.webp"}
|
||||||
|
second := &db.R2DeleteTask{ID: 52, Bucket: conf.Upload.ImageBucket, Key: "second.webp"}
|
||||||
|
store := newFakeR2OutboxStore(first, second)
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
deleter := &fakeR2ObjectDeleter{
|
||||||
|
onDelete: func(ctx context.Context, _, _ string) error {
|
||||||
|
cancel()
|
||||||
|
return ctx.Err()
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
result, err := processR2Outbox(ctx, store, deleter, 100)
|
||||||
|
require.ErrorIs(t, err, context.Canceled)
|
||||||
|
require.Len(t, deleter.calls, 1, "取消后不得继续处理第二个任务")
|
||||||
|
require.Empty(t, store.failures, "优雅关闭导致的取消不应增加任务失败次数")
|
||||||
|
require.Empty(t, store.deleted)
|
||||||
|
require.Zero(t, result.Failed)
|
||||||
|
require.Equal(t, 2, result.Remaining)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestR2OutboxBackoffCapsAtOneHour(t *testing.T) {
|
||||||
|
for _, tt := range []struct {
|
||||||
|
attempts int
|
||||||
|
expected time.Duration
|
||||||
|
}{
|
||||||
|
{attempts: -1, expected: time.Minute},
|
||||||
|
{attempts: 0, expected: time.Minute},
|
||||||
|
{attempts: 1, expected: 2 * time.Minute},
|
||||||
|
{attempts: 2, expected: 4 * time.Minute},
|
||||||
|
{attempts: 5, expected: 32 * time.Minute},
|
||||||
|
{attempts: 6, expected: time.Hour},
|
||||||
|
{attempts: 100, expected: time.Hour},
|
||||||
|
} {
|
||||||
|
require.Equal(t, tt.expected, r2OutboxBackoff(tt.attempts))
|
||||||
|
}
|
||||||
|
}
|
||||||
+10
-1
@@ -10,6 +10,7 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
"os/signal"
|
"os/signal"
|
||||||
|
"sync"
|
||||||
"syscall"
|
"syscall"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
@@ -209,7 +210,15 @@ func runWeb(ctx *cli.Context) error {
|
|||||||
// 「提问已 INSERT、图片还没传完」「回答已写入、配额还没扣」这类半成品。
|
// 「提问已 INSERT、图片还没传完」「回答已写入、配额还没扣」这类半成品。
|
||||||
// systemd Restart=always 与 docker restart 都救不了这个。
|
// systemd Restart=always 与 docker restart 都救不了这个。
|
||||||
runCtx, stop := signal.NotifyContext(ctx.Context, syscall.SIGTERM, os.Interrupt)
|
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)
|
serveErr := make(chan error, 1)
|
||||||
go func() {
|
go func() {
|
||||||
|
|||||||
+2
-1
@@ -36,7 +36,7 @@ func Ping(ctx context.Context) error {
|
|||||||
|
|
||||||
var AllTables = []interface{}{
|
var AllTables = []interface{}{
|
||||||
&User{}, &Question{}, &FollowUp{}, &CensorLog{}, &UploadImage{}, &UploadImageQuestion{},
|
&User{}, &Question{}, &FollowUp{}, &CensorLog{}, &UploadImage{}, &UploadImageQuestion{},
|
||||||
&Conversation{}, &Message{}, &MessageImage{}, &UserBlock{},
|
&Conversation{}, &Message{}, &MessageImage{}, &UserBlock{}, &R2DeleteTask{},
|
||||||
// 匿名投票
|
// 匿名投票
|
||||||
&Vote{}, &VoteOption{}, &VoteFormField{}, &VoteSubmission{},
|
&Vote{}, &VoteOption{}, &VoteFormField{}, &VoteSubmission{},
|
||||||
// 功能开关(提问 / 私聊 / 投票 的开关、定时、定量收集)
|
// 功能开关(提问 / 私聊 / 投票 的开关、定时、定量收集)
|
||||||
@@ -132,6 +132,7 @@ func Init(typ, dsn string) (*gorm.DB, error) {
|
|||||||
Conversations = NewConversationStore(db)
|
Conversations = NewConversationStore(db)
|
||||||
Messages = NewMessageStore(db)
|
Messages = NewMessageStore(db)
|
||||||
MessageImages = NewMessageImageStore(db)
|
MessageImages = NewMessageImageStore(db)
|
||||||
|
R2DeleteTasks = NewR2DeleteTaskStore(db)
|
||||||
Votes = NewVotesStore(db)
|
Votes = NewVotesStore(db)
|
||||||
VoteOptions = NewVoteOptionsStore(db)
|
VoteOptions = NewVoteOptionsStore(db)
|
||||||
VoteFormFields = NewVoteFormFieldsStore(db)
|
VoteFormFields = NewVoteFormFieldsStore(db)
|
||||||
|
|||||||
@@ -26,6 +26,7 @@ type MessageImage struct {
|
|||||||
gorm.Model
|
gorm.Model
|
||||||
MessageID uint `gorm:"index" json:"message_id"`
|
MessageID uint `gorm:"index" json:"message_id"`
|
||||||
URL string `gorm:"size:500" json:"url"`
|
URL string `gorm:"size:500" json:"url"`
|
||||||
|
ObjectKey string `gorm:"type:text" json:"-"`
|
||||||
FileName string `gorm:"size:200" json:"file_name"`
|
FileName string `gorm:"size:200" json:"file_name"`
|
||||||
FileSize int64 `json:"file_size"`
|
FileSize int64 `json:"file_size"`
|
||||||
}
|
}
|
||||||
@@ -33,6 +34,7 @@ type MessageImage struct {
|
|||||||
type CreateMessageImageOptions struct {
|
type CreateMessageImageOptions struct {
|
||||||
MessageID uint
|
MessageID uint
|
||||||
URL string
|
URL string
|
||||||
|
ObjectKey string
|
||||||
FileName string
|
FileName string
|
||||||
FileSize int64
|
FileSize int64
|
||||||
}
|
}
|
||||||
@@ -41,6 +43,7 @@ func (db *messageImages) Create(ctx context.Context, opts CreateMessageImageOpti
|
|||||||
img := &MessageImage{
|
img := &MessageImage{
|
||||||
MessageID: opts.MessageID,
|
MessageID: opts.MessageID,
|
||||||
URL: opts.URL,
|
URL: opts.URL,
|
||||||
|
ObjectKey: opts.ObjectKey,
|
||||||
FileName: opts.FileName,
|
FileName: opts.FileName,
|
||||||
FileSize: opts.FileSize,
|
FileSize: opts.FileSize,
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -10,6 +10,8 @@ import (
|
|||||||
"github.com/pkg/errors"
|
"github.com/pkg/errors"
|
||||||
"github.com/samber/lo"
|
"github.com/samber/lo"
|
||||||
"gorm.io/gorm"
|
"gorm.io/gorm"
|
||||||
|
|
||||||
|
"github.com/tamakyi/TamaBox/internal/conf"
|
||||||
)
|
)
|
||||||
|
|
||||||
var UploadImages UploadImagesStore
|
var UploadImages UploadImagesStore
|
||||||
@@ -23,12 +25,8 @@ type UploadImagesStore interface {
|
|||||||
// 用于账号导出等场景,将 N 次单题查询合并为少量 IN 查询,避免 N+1。
|
// 用于账号导出等场景,将 N 次单题查询合并为少量 IN 查询,避免 N+1。
|
||||||
GetByQuestionIDs(ctx context.Context, questionIDs []uint) (map[uint][]*UploadImage, error)
|
GetByQuestionIDs(ctx context.Context, questionIDs []uint) (map[uint][]*UploadImage, error)
|
||||||
GetByTypeQuestionID(ctx context.Context, typ UploadImageQuestionType, questionID uint) ([]*UploadImage, error)
|
GetByTypeQuestionID(ctx context.Context, typ UploadImageQuestionType, questionID uint) ([]*UploadImage, error)
|
||||||
// DeleteByQuestionID 删除某条提问时级联清理其全部上传图片记录与关联行。
|
// DeleteByQuestionID 先删除当前题目的关联,再物理删除不再被任何题目引用的主体。
|
||||||
//
|
// 实际删除的主体的对象 key 与元数据删除在同一事务中登记 R2 outbox。
|
||||||
// 用 Unscoped 真删:图片元数据(name/md5/key)随提问一并真删,与 questions.DeleteByID
|
|
||||||
// 对主表的真删口径一致。软删会留下带 deleted_at 的图片行,且关联行的 question_id
|
|
||||||
// 指向已不存在的提问。注意这里只清数据库记录,对象存储里的实际文件不在本函数删除
|
|
||||||
// (图片可能被多条记录复用,避免误删共享资源)。
|
|
||||||
DeleteByQuestionID(ctx context.Context, questionID uint) error
|
DeleteByQuestionID(ctx context.Context, questionID uint) error
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -87,9 +85,24 @@ func (db *uploadImages) Create(ctx context.Context, opts CreateUploadImageOption
|
|||||||
}
|
}
|
||||||
|
|
||||||
if opts.IsDeletingPrevious {
|
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")
|
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{
|
if err := tx.WithContext(ctx).Create(&UploadImageQuestion{
|
||||||
Type: opts.Type,
|
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)
|
return db.getBy(ctx, "type = ? AND question_id = ?", typ, questionID)
|
||||||
}
|
}
|
||||||
|
|
||||||
// DeleteByQuestionID 删除某条提问时级联清理其全部上传图片记录(软删除,与提问保持一致)。
|
// unreferencedUploadImagePredicate 是所有 UploadImage 删除入口共用的共享保护谓词。
|
||||||
// 注意:question_id 列在关联表 upload_image_questions 上,不在 upload_images 表上,
|
// 调用方必须先解除当前作用域的关联,再使用该谓词收集并删除真正的孤儿主体。
|
||||||
// 需先查关联行拿到 imageID 列表,再软删图片主体并删除关联行(同一事务内完成)。
|
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 {
|
func (db *uploadImages) DeleteByQuestionID(ctx context.Context, questionID uint) error {
|
||||||
return db.Transaction(func(tx *gorm.DB) error {
|
return db.Transaction(func(tx *gorm.DB) error {
|
||||||
var links []*UploadImageQuestion
|
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 {
|
imageIDs := lo.Map(links, func(item *UploadImageQuestion, _ int) uint {
|
||||||
return item.UploadImageID
|
return item.UploadImageID
|
||||||
})
|
})
|
||||||
if err := tx.WithContext(ctx).Where("id IN ?", imageIDs).Unscoped().Delete(&UploadImage{}).Error; err != nil {
|
// 关联表无 DeletedAt,Delete 本就是真删;先解除当前题目关联。
|
||||||
return errors.Wrap(err, "delete upload images by question id")
|
|
||||||
}
|
|
||||||
// 关联表 UploadImageQuestion 无 DeletedAt,Delete 本就是真删,无需 Unscoped。
|
|
||||||
if err := tx.WithContext(ctx).Where("question_id = ?", questionID).Delete(&UploadImageQuestion{}).Error; err != nil {
|
if err := tx.WithContext(ctx).Where("question_id = ?", questionID).Delete(&UploadImageQuestion{}).Error; err != nil {
|
||||||
return errors.Wrap(err, "delete image question links")
|
return errors.Wrap(err, "delete image question links")
|
||||||
}
|
}
|
||||||
return nil
|
return deleteUnreferencedUploadImages(ctx, tx, imageIDs)
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
+66
-23
@@ -386,10 +386,25 @@ func (db *users) Update(ctx context.Context, id uint, opts UpdateUserOptions) er
|
|||||||
"Language": opts.Language,
|
"Language": opts.Language,
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := db.WithContext(ctx).Model(&User{}).Where("id = ?", id).Updates(updates).Error; err != nil {
|
return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||||
return errors.Wrap(err, "update user")
|
var current User
|
||||||
}
|
if err := tx.WithContext(ctx).Select("id", "avatar", "background", "background_image").First(¤t, id).Error; err != nil {
|
||||||
return 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)。
|
// UpdateHarassmentType 仅更新骚扰防护类型(none / register_only)。
|
||||||
@@ -1021,6 +1036,19 @@ func (db *users) BanUserAndCleanup(ctx context.Context, userID uint) error {
|
|||||||
// 直接复用会把上一条语句的 Where 条件串到下一条上(曾把 messages 的条件串进 follow_ups,
|
// 直接复用会把上一条语句的 Where 条件串到下一条上(曾把 messages 的条件串进 follow_ups,
|
||||||
// 报「字段 conversation_id 不存在」)。
|
// 报「字段 conversation_id 不存在」)。
|
||||||
func purgeUserContent(ctx context.Context, tx *gorm.DB, userID uint) error {
|
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().
|
if err := tx.WithContext(ctx).Unscoped().
|
||||||
Where("question_id IN (SELECT id FROM questions WHERE user_id = ?)", userID).
|
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().
|
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 {
|
Delete(&MessageImage{}).Error; err != nil {
|
||||||
return errors.Wrap(err, "delete message images of user")
|
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 的会话
|
// 消息:属于该用户作为 owner 或 asker 的会话
|
||||||
if err := tx.WithContext(ctx).Unscoped().
|
if err := tx.WithContext(ctx).Unscoped().
|
||||||
Where("conversation_id IN (SELECT id FROM conversations WHERE owner_id = ? OR asker_id = ?)", userID, userID).
|
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")
|
return errors.Wrap(err, "delete conversations of user")
|
||||||
}
|
}
|
||||||
|
|
||||||
// 提问图片元数据:删除本用户上传的孤立图片,以及仅链接到本用户提问的图片。
|
// 先确定本用户上传或仅与当前作用域有关的候选主体,再解除本用户题目的关联。
|
||||||
// 仍被其他提问引用的图片必须保留;这里只清数据库元数据,不操作对象存储。
|
// 最终是否可删除由与 UploadImages.DeleteByQuestionID 完全相同的共享保护谓词决定。
|
||||||
if err := tx.WithContext(ctx).Unscoped().
|
var uploadImageIDs []uint
|
||||||
Where(`(
|
if err := tx.WithContext(ctx).Unscoped().Model(&UploadImage{}).
|
||||||
uploader_user_id = ?
|
Where(`uploader_user_id = ? OR upload_images.id IN (
|
||||||
OR upload_images.id IN (
|
SELECT upload_image_id FROM upload_image_questions
|
||||||
SELECT upload_image_id FROM upload_image_questions
|
WHERE question_id IN (SELECT id FROM questions WHERE user_id = ?)
|
||||||
WHERE question_id IN (SELECT id FROM questions WHERE user_id = ?)
|
)`, userID, userID).
|
||||||
)
|
Pluck("upload_images.id", &uploadImageIDs).Error; err != nil {
|
||||||
) AND NOT EXISTS (
|
return errors.Wrap(err, "get upload images of user")
|
||||||
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")
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// 提问-图片关联
|
|
||||||
if err := tx.WithContext(ctx).Unscoped().
|
if err := tx.WithContext(ctx).Unscoped().
|
||||||
Where("question_id IN (SELECT id FROM questions WHERE user_id = ?)", userID).
|
Where("question_id IN (SELECT id FROM questions WHERE user_id = ?)", userID).
|
||||||
Delete(&UploadImageQuestion{}).Error; err != nil {
|
Delete(&UploadImageQuestion{}).Error; err != nil {
|
||||||
return errors.Wrap(err, "delete upload image links of user")
|
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().
|
if err := tx.WithContext(ctx).Unscoped().
|
||||||
|
|||||||
@@ -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, "[email protected]", "outbox-owner")
|
||||||
|
other := createOutboxTestUser(t, ctx, users, "[email protected]", "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, "[email protected]", "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)
|
||||||
|
}
|
||||||
@@ -277,7 +277,8 @@ func SendMessage(ctx context.Context, cache cache.Cache) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// 上传图片(如有)
|
// 上传图片(如有)
|
||||||
var imageURL string
|
var imageURL, imageKey string
|
||||||
|
imageAttached := false
|
||||||
if imageFile != nil {
|
if imageFile != nil {
|
||||||
uploadedURL, err := storage.UploadPictureToS3(imageFile, imageHeader)
|
uploadedURL, err := storage.UploadPictureToS3(imageFile, imageHeader)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -291,6 +292,31 @@ func SendMessage(ctx context.Context, cache cache.Cache) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
imageURL = uploadedURL
|
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{
|
if _, err := db.MessageImages.Create(ctx.Request().Context(), db.CreateMessageImageOptions{
|
||||||
MessageID: msg.ID,
|
MessageID: msg.ID,
|
||||||
URL: imageURL,
|
URL: imageURL,
|
||||||
|
ObjectKey: imageKey,
|
||||||
FileName: imageHeader.Filename,
|
FileName: imageHeader.Filename,
|
||||||
FileSize: imageHeader.Size,
|
FileSize: imageHeader.Size,
|
||||||
}); err != nil {
|
}); err != nil {
|
||||||
logrus.WithError(err).Error("Failed to save message image")
|
logrus.WithError(err).Error("Failed to save message image")
|
||||||
|
ctx.SetInternalErrorFlash()
|
||||||
|
ctx.Redirect(fmt.Sprintf("/_/%s/chat/%s", domain, token))
|
||||||
|
return
|
||||||
}
|
}
|
||||||
|
imageAttached = true
|
||||||
}
|
}
|
||||||
|
|
||||||
// 达到上限时关闭对话(幂等,多实例同时执行也安全)
|
// 达到上限时关闭对话(幂等,多实例同时执行也安全)
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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://[email protected]/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",
|
||||||
|
"[email protected]",
|
||||||
|
} {
|
||||||
|
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)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -5,9 +5,6 @@ import (
|
|||||||
"sync"
|
"sync"
|
||||||
|
|
||||||
"github.com/aliyun/aliyun-oss-go-sdk/oss"
|
"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/aws/aws-sdk-go-v2/service/s3"
|
||||||
"github.com/pkg/errors"
|
"github.com/pkg/errors"
|
||||||
|
|
||||||
@@ -31,29 +28,13 @@ func GetS3Client() (*s3.Client, error) {
|
|||||||
s3ClientErr = errors.New("S3 endpoint not configured")
|
s3ClientErr = errors.New("S3 endpoint not configured")
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
s3Client, s3ClientErr = NewS3Client(
|
||||||
r2Resolver := aws.EndpointResolverWithOptionsFunc(func(service, region string, options ...interface{}) (aws.Endpoint, error) {
|
context.Background(),
|
||||||
return aws.Endpoint{
|
conf.Upload.ImageEndpoint,
|
||||||
URL: conf.Upload.ImageEndpoint,
|
conf.Upload.ImageBucket,
|
||||||
HostnameImmutable: true,
|
conf.Upload.ImageAccessID,
|
||||||
Source: aws.EndpointSourceCustom,
|
conf.Upload.ImageAccessSecret,
|
||||||
}, nil
|
|
||||||
})
|
|
||||||
|
|
||||||
cfg, err := config.LoadDefaultConfig(context.Background(),
|
|
||||||
config.WithEndpointResolverWithOptions(r2Resolver),
|
|
||||||
config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(
|
|
||||||
conf.Upload.ImageAccessID,
|
|
||||||
conf.Upload.ImageAccessSecret,
|
|
||||||
"",
|
|
||||||
)),
|
|
||||||
config.WithRegion("auto"),
|
|
||||||
)
|
)
|
||||||
if err != nil {
|
|
||||||
s3ClientErr = errors.Wrap(err, "load S3 config")
|
|
||||||
return
|
|
||||||
}
|
|
||||||
s3Client = s3.NewFromConfig(cfg)
|
|
||||||
})
|
})
|
||||||
return s3Client, s3ClientErr
|
return s3Client, s3ClientErr
|
||||||
}
|
}
|
||||||
|
|||||||
+11
-14
@@ -23,7 +23,6 @@ import (
|
|||||||
"github.com/disintegration/imaging"
|
"github.com/disintegration/imaging"
|
||||||
"github.com/pkg/errors"
|
"github.com/pkg/errors"
|
||||||
"github.com/sirupsen/logrus"
|
"github.com/sirupsen/logrus"
|
||||||
"github.com/wuhan005/gadget"
|
|
||||||
|
|
||||||
"github.com/tamakyi/TamaBox/internal/conf"
|
"github.com/tamakyi/TamaBox/internal/conf"
|
||||||
"github.com/tamakyi/TamaBox/internal/imagecheck"
|
"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")
|
return "", errors.Wrap(err, "get S3 client")
|
||||||
}
|
}
|
||||||
|
|
||||||
// 7. 上传到 R2
|
// 7. 上传到 R2。现有上传 API 不接收调用方 context,故显式标记为 TODO;
|
||||||
if err := gadget.Retry(5, func() error {
|
// PutObject 内部会为每次尝试设置独立超时,并对可取消 context 做退避。
|
||||||
_, err := client.PutObject(context.Background(), &s3.PutObjectInput{
|
if err := PutObject(context.TODO(), client, &s3.PutObjectInput{
|
||||||
Bucket: aws.String(conf.Upload.ImageBucket),
|
Bucket: aws.String(conf.Upload.ImageBucket),
|
||||||
Key: aws.String(key),
|
Key: aws.String(key),
|
||||||
Body: bytes.NewReader(uploadData),
|
Body: bytes.NewReader(uploadData),
|
||||||
ContentLength: aws.Int64(int64(len(uploadData))),
|
ContentLength: aws.Int64(int64(len(uploadData))),
|
||||||
ContentType: aws.String(contentType),
|
ContentType: aws.String(contentType),
|
||||||
// key 含纳秒时间戳且内容不可变 → 可安全启用 immutable 强缓存,
|
// key 含纳秒时间戳且内容不可变 → 可安全启用 immutable 强缓存,
|
||||||
// 避免每次页面加载都回源校验(曾导致头像/二维码 logo 每次重新下载)。
|
// 避免每次页面加载都回源校验(曾导致头像/二维码 logo 每次重新下载)。
|
||||||
CacheControl: aws.String("public, max-age=31536000, immutable"),
|
CacheControl: aws.String("public, max-age=31536000, immutable"),
|
||||||
})
|
|
||||||
return err
|
|
||||||
}); err != nil {
|
}); err != nil {
|
||||||
return "", errors.Wrap(err, "put object after retry")
|
return "", errors.Wrap(err, "put object after retry")
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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 可以不带路径,也可以带且仅带 /<bucket>;后者会被剥离,
|
||||||
|
// 再由 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 等编码形式绕过路径约束。
|
||||||
|
// 空路径保持原样;唯一允许的非空路径是 /<bucket>,包括 "/" 在内的
|
||||||
|
// 其他路径均拒绝,避免静默改变 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}
|
||||||
|
}
|
||||||
@@ -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://[email protected]", 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")
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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())
|
||||||
|
}
|
||||||
@@ -871,7 +871,7 @@ func uploadImage(ctx context.Context, opts uploadImageOptions) error {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return errors.Wrap(err, "get S3 client")
|
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),
|
Bucket: aws.String(conf.Upload.ImageBucket),
|
||||||
Key: aws.String(uploadKey),
|
Key: aws.String(uploadKey),
|
||||||
Body: bytes.NewReader(uploadData),
|
Body: bytes.NewReader(uploadData),
|
||||||
@@ -894,6 +894,9 @@ func uploadImage(ctx context.Context, opts uploadImageOptions) error {
|
|||||||
IsDeletingPrevious: opts.IsDeletingPrevious,
|
IsDeletingPrevious: opts.IsDeletingPrevious,
|
||||||
})
|
})
|
||||||
if err != nil {
|
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 errors.Wrap(err, "create upload image")
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
|
|||||||
@@ -5,6 +5,7 @@
|
|||||||
package user
|
package user
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
stdctx "context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"net/url"
|
"net/url"
|
||||||
@@ -121,6 +122,17 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) {
|
|||||||
return
|
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)
|
// 处理头像上传(强制 1:1,见 storage.UploadAvatarToS3)
|
||||||
avatarURL := currentUser.Avatar // 默认使用原值
|
avatarURL := currentUser.Avatar // 默认使用原值
|
||||||
avatarFile, avatarFileHeader, err := ctx.Request().FormFile("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")
|
logrus.WithContext(ctx.Request().Context()).WithError(err).Error("Failed to upload avatar")
|
||||||
} else {
|
} else {
|
||||||
avatarURL = uploadedURL
|
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")
|
logrus.WithContext(ctx.Request().Context()).WithError(err).Error("Failed to upload background")
|
||||||
} else {
|
} else {
|
||||||
backgroundURL = uploadedURL
|
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")
|
logrus.WithContext(ctx.Request().Context()).WithError(err).Error("Failed to upload QRbackgroundImage")
|
||||||
} else {
|
} else {
|
||||||
backgroundImageURL = uploadedURL
|
backgroundImageURL = uploadedURL
|
||||||
|
uploadedProfileImages = append(uploadedProfileImages, uploadedURL)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -268,6 +283,7 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) {
|
|||||||
}
|
}
|
||||||
ctx.SetInternalErrorFlash()
|
ctx.SetInternalErrorFlash()
|
||||||
} else if passwordChanged {
|
} else if passwordChanged {
|
||||||
|
profileSaved = true
|
||||||
// 密码修改成功且其余资料也保存成功:强制登出重登。
|
// 密码修改成功且其余资料也保存成功:强制登出重登。
|
||||||
// 先 Flush 清会话再写入 flash,保证提示能带到登录页显示。
|
// 先 Flush 清会话再写入 flash,保证提示能带到登录页显示。
|
||||||
ctx.Session.Flush()
|
ctx.Session.Flush()
|
||||||
@@ -275,6 +291,7 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) {
|
|||||||
ctx.Redirect("/login")
|
ctx.Redirect("/login")
|
||||||
return
|
return
|
||||||
} else {
|
} else {
|
||||||
|
profileSaved = true
|
||||||
ctx.SetSuccessFlash(i18n.T(ctx.Lang(), "admin.flash.profileUpdated"))
|
ctx.SetSuccessFlash(i18n.T(ctx.Lang(), "admin.flash.profileUpdated"))
|
||||||
}
|
}
|
||||||
ctx.Redirect("/user/profile")
|
ctx.Redirect("/user/profile")
|
||||||
|
|||||||
@@ -30,6 +30,8 @@ func main() {
|
|||||||
cmd.HidePublic,
|
cmd.HidePublic,
|
||||||
cmd.Honeypot,
|
cmd.Honeypot,
|
||||||
cmd.Backup,
|
cmd.Backup,
|
||||||
|
cmd.RestoreImages,
|
||||||
|
cmd.R2Outbox,
|
||||||
}
|
}
|
||||||
if err := app.Run(os.Args); err != nil {
|
if err := app.Run(os.Args); err != nil {
|
||||||
logrus.WithError(err).Fatal("Failed to start application")
|
logrus.WithError(err).Fatal("Failed to start application")
|
||||||
|
|||||||
Reference in New Issue
Block a user