tamakyi/tamabox
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

可靠性: 完善 R2 备份恢复与对象生命周期

Browse Source
This commit is contained in:
tamakyi
2026-11-05 08:41:55 +08:00
parent 3e7719f200
commit 937e688cfa
32 changed files with 4151 additions and 231 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/tamabox.go
+2
View File
@@ -30,6 +30,8 @@ func main() {
cmd.HidePublic,
cmd.Honeypot,
cmd.Backup,
cmd.RestoreImages,
cmd.R2Outbox,
}
if err := app.Run(os.Args); err != nil {
logrus.WithError(err).Fatal("Failed to start application")
docs/cli.md
+56 -5
View File
@@ -32,7 +32,7 @@ ps -eo pid,lstart,cmd | grep '[t]amabox' # 看进程启动时间，与二进
```bash
./tamabox --help | sed -n '/COMMANDS/,/GLOBAL/p'
# web / censor / uid / ban / restore / reports / restrict / unrestrict / hide-public / honeypot
# web / censor / uid / ban / restore / reports / restrict / unrestrict / hide-public / honeypot / backup / restore-images
```
> 想确认某个字符串有没有编进二进制，**不要用 `strings`** —— 它默认只输出可打印 ASCII，
@@ -77,10 +77,13 @@ pg_dump -Fc <dbname> > before-$(date +%F-%H%M).dump
| `censor` | 重新审核历史内容 | 是 | 审核规则变更后复检 |
| `uid` | 为用户生成/补齐 uid | 是 | 数据修复 |
| `honeypot` | 扫描器诱捕取证（假登录页） | 部分是（`prune` / `purge-creds`） | 看谁在扫、拿到了什么字典 |
| `backup` | 备份 PostgreSQL 与配置 bucket 的全部 R2 对象 | 写本地备份目录 | 灾难恢复前留档 |
| `restore-images` | 校验并恢复 R2 图片及元数据 | 是（除 `--dry-run`） | 图片灾难恢复、隔离恢复演练 |
| `web` | 启动 Web 服务 | — | 生产入口 |
除 `reports`、`honeypot hits/stats/paths/size` 外，所有写操作都要求交互式输入大写 `YES` 确认；
除 `reports`、`honeypot hits/stats/paths/size` 外，既有运营写操作要求交互式输入大写 `YES` 确认；
脚本化时加 `--yes`。注意 `honeypot prune` / `purge-creds` 的确认词是 **`yes`**（小写）。
`restore-images` 不做交互确认：实际写入必须显式提供 `--yes`，`--dry-run` 除外。
---
@@ -329,7 +332,55 @@ pg_dump -Fc <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 次」。
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 写入。
internal/cmd/backup.go
+391 -122
View File
@@ -6,40 +6,46 @@ package cmd
import (
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"time"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/config"
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/pkg/errors"
"github.com/sirupsen/logrus"
"github.com/urfave/cli/v2"
"github.com/tamakyi/TamaBox/internal/conf"
"github.com/tamakyi/TamaBox/internal/storage"
)
const (
imageObjectAttemptTimeout = 60 * time.Second
imageObjectMaxAttempts = 3
imageObjectRetryBackoff = 100 * time.Millisecond
)
// Backup 全量备份数据库与对象存储（Cloudflare R2）图片到本地目录。
//
// 为什么必须同时备份图片：站点的所有图片（头像 / 背景 / 二维码 logo / 提问图 / 回答图）
// 都以 picture/YYYY/MM/DD/<hex>（新）或 tamabox/YYYY/MM/<hex>.<ext>（历史）的形式存在
// R2 上，数据库里只存 key（upload_images.key）或 URL（users.avatar/background 里的完整
// CDN 地址）。只备份数据库的话，恢复后这些引用指向的图片若已不在 R2（bucket 误删 /
// 账号丢失 / 图片被清理），页面就全是裂图。因此备份必须把 R2 上的图片一并落盘。
// 都以根日期 key、picture/ 或 tamabox/ 等形式存在 R2 上，数据库里只存 key 或 CDN URL。
// 只备份数据库的话，恢复后这些引用指向的图片若已不在 R2，页面就全是裂图。
//
// 输出目录结构：
//
// <output>/
// db-<20060102-150405>.sql # pg_dump 全量导出（schema + data，可用 psql 恢复）
// images/<key> # R2 对象，保持原 key 路径
// images-manifest.json # key → md5 / size 清单，用于恢复后完整性校验
// images-manifest.json # key → 实际大小 / SHA-256 / GET 元数据
//
// 用法：TamaBox backup --output /var/backups/tamabox
//
@@ -47,7 +53,7 @@ import (
// TamaBox backup --output . --images-only # 只备份图片
var Backup = &cli.Command{
Name: "backup",
Usage: "Back up the database and R2 images to a local directory",
Usage: "Back up the database and all R2 objects to a local directory",
Action: runBackup,
Flags: []cli.Flag{
&cli.StringFlag{
@@ -67,6 +73,9 @@ var Backup = &cli.Command{
}
func runBackup(ctx *cli.Context) error {
if ctx.Bool("db-only") && ctx.Bool("images-only") {
return errors.New("--db-only and --images-only cannot be used together")
}
if err := conf.Init(); err != nil {
return errors.Wrap(err, "load configuration")
}
@@ -127,190 +136,450 @@ func backupDatabase(ctx context.Context, outDir string) error {
)
cmd.Env = append(os.Environ(), "PGPASSWORD="+conf.Database.Password)
f, err := os.OpenFile(outFile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
if err != nil {
return errors.Wrapf(err, "create backup file %q", outFile)
}
defer f.Close()
var dumpErr strings.Builder
cmd.Stdout = f
if err := writeFileAtomically(outFile, func(w io.Writer) error {
cmd.Stdout = w
cmd.Stderr = &dumpErr
if err := cmd.Run(); err != nil {
return errors.Wrapf(err, "pg_dump failed: %s", strings.TrimSpace(dumpErr.String()))
}
return nil
}); err != nil {
return err
}
logrus.WithField("file", outFile).Info("数据库已备份")
return nil
}
// imageManifestEntry 图片备份清单里的一条：key → 大小 → md5（用于恢复后校验完整性）。
// writeFileAtomically 保留数据库、manifest 与对象下载共用的原子写修复：
// 只有完整写入、Sync 并关闭成功后才替换目标，失败时清理临时 .part 文件。
func writeFileAtomically(outFile string, write func(io.Writer) error) error {
f, err := os.CreateTemp(filepath.Dir(outFile), "."+filepath.Base(outFile)+"-*.part")
if err != nil {
return errors.Wrapf(err, "create temporary file for %q", outFile)
}
tmpFile := f.Name()
defer os.Remove(tmpFile)
defer f.Close()
if err := write(f); err != nil {
return err
}
if err := f.Sync(); err != nil {
return errors.Wrapf(err, "sync temporary file for %q", outFile)
}
if err := f.Close(); err != nil {
return errors.Wrapf(err, "close temporary file for %q", outFile)
}
if err := os.Rename(tmpFile, outFile); err != nil {
return errors.Wrapf(err, "finalize %q", outFile)
}
return nil
}
// imageManifestEntry 记录实际下载内容的大小和 SHA-256，以及 GET 返回的原始元数据。
// ETag 只是服务商元数据，不能当作内容 MD5 或完整性校验的替代品。
type imageManifestEntry struct {
Key string `json:"key"`
Size int64 `json:"size"`
MD5 string `json:"md5,omitempty"`
SHA256 string `json:"sha256"`
ETag string `json:"etag,omitempty"`
ContentType string `json:"content_type,omitempty"`
CacheControl string `json:"cache_control,omitempty"`
}
// backupImages 从 R2 下载 tamabox 站的全部图片对象到本地，并生成 manifest。
//
// R2 里 tamabox 站图片的 key 都带 `tamabox/` 前缀（如 tamabox/2025/12/<ts>.jpg），
// 这是「bucket 用 path-style 访问时 List 返回的 key 带 bucket 名作前缀」这一 R2 行为。
// 因此只需列出 `tamabox/` 前缀即可覆盖全部图片。blog/ 是同一 R2 account 下另一个站点
// 的内容，不在 tamabox 的备份范围。
type backupS3API interface {
ListObjectsV2(context.Context, *s3.ListObjectsV2Input, ...func(*s3.Options)) (*s3.ListObjectsV2Output, error)
GetObject(context.Context, *s3.GetObjectInput, ...func(*s3.Options)) (*s3.GetObjectOutput, error)
}
type backupDownloadPlan struct {
key string
destination string
}
type downloadedImageObject struct {
size int64
sha256 string
etag string
contentType string
cacheControl string
}
// backupImages 默认使用空前缀列举配置 bucket 的全部对象，而不是假定图片都在 tamabox/ 下。
func backupImages(ctx context.Context, outDir string) error {
client, err := backupS3Client()
client, err := backupS3Client(ctx)
if err != nil {
return errors.Wrap(err, "get R2 client for backup")
}
imageDir := filepath.Join(outDir, "images")
if err := os.MkdirAll(imageDir, 0o700); err != nil {
return errors.Wrap(err, "create images directory")
return backupImagesWithClient(ctx, client, conf.Upload.ImageBucket, outDir, "")
}
var manifest []imageManifestEntry
var count int64
var totalBytes int64
// 见函数注释：tamabox/ 前缀即覆盖全部图片（含新旧 key 格式）。
if err := backupPrefix(ctx, client, imageDir, "tamabox/", &manifest, &count, &totalBytes); err != nil {
// backupImagesWithClient 允许测试和隔离演练显式指定 prefix。
// 所有分页 List、key 校验和本地路径预检都在第一次 GetObject 之前完成。
func backupImagesWithClient(ctx context.Context, client backupS3API, bucket, outDir, prefix string) error {
keys, err := listAllBackupKeys(ctx, client, bucket, prefix)
if err != nil {
return err
}
// 写 manifest：记录 key → 大小 → md5，恢复后可用于校验完整性。
mf, err := os.Create(filepath.Join(outDir, "images-manifest.json"))
imageDir := filepath.Join(outDir, "images")
plans, err := planBackupDownloads(imageDir, keys)
if err != nil {
return errors.Wrap(err, "create manifest file")
return err
}
defer mf.Close()
enc := json.NewEncoder(mf)
manifest := make([]imageManifestEntry, 0, len(plans))
var totalBytes int64
for _, plan := range plans {
if err := ensureBackupDestinationDirectory(imageDir, plan.destination); err != nil {
return err
}
downloaded, err := downloadObject(ctx, client, bucket, plan.key, plan.destination)
if err != nil {
return err
}
manifest = append(manifest, imageManifestEntry{
Key: plan.key,
Size: downloaded.size,
SHA256: downloaded.sha256,
ETag: downloaded.etag,
ContentType: downloaded.contentType,
CacheControl: downloaded.cacheControl,
})
totalBytes += downloaded.size
}
manifestFile := filepath.Join(outDir, "images-manifest.json")
if err := writeFileAtomically(manifestFile, func(w io.Writer) error {
enc := json.NewEncoder(w)
enc.SetIndent("", " ")
if err := enc.Encode(manifest); err != nil {
return errors.Wrap(err, "write manifest")
}
return nil
}); err != nil {
return err
}
logrus.WithFields(logrus.Fields{
"objects": count,
"objects": len(manifest),
"bytes": totalBytes,
"dir": imageDir,
"prefix": prefix,
}).Info("图片已备份")
return nil
}
// backupPrefix 分页列出并下载指定前缀下的全部对象。
func backupPrefix(ctx context.Context, client *s3.Client, imageDir, prefix string, manifest *[]imageManifestEntry, count, totalBytes *int64) error {
paginator := s3.NewListObjectsV2Paginator(client, &s3.ListObjectsV2Input{
Bucket: aws.String(conf.Upload.ImageBucket),
func listAllBackupKeys(ctx context.Context, client backupS3API, bucket, prefix string) ([]string, error) {
keys := make([]string, 0)
var continuationToken *string
seenTokens := make(map[string]struct{})
for {
page, err := client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
Bucket: aws.String(bucket),
Prefix: aws.String(prefix),
ContinuationToken: continuationToken,
})
for paginator.HasMorePages() {
page, err := paginator.NextPage(ctx)
if err != nil {
return errors.Wrapf(err, "list R2 objects prefix %q", prefix)
return nil, errors.Wrapf(err, "list R2 objects prefix %q", prefix)
}
for _, obj := range page.Contents {
if obj.Key == nil {
continue
if page == nil {
return nil, errors.New("list R2 objects returned a nil page")
}
key := *obj.Key
// 落盘路径 = imageDir + "/" + key，保持对象 key 结构（含 tamabox/ 前缀），
// 恢复时可直接用同一 key 回传 R2。
dst := filepath.Join(imageDir, filepath.FromSlash(key))
if err := os.MkdirAll(filepath.Dir(dst), 0o700); err != nil {
return errors.Wrapf(err, "create dir for %q", key)
}
if err := downloadObject(ctx, client, conf.Upload.ImageBucket, key, dst); err != nil {
return err
for _, object := range page.Contents {
// nil key 与空 key 一样是无效输入；统一转成空串，由完整 List 后的预检拒绝。
keys = append(keys, aws.ToString(object.Key))
}
size := int64(0)
if obj.Size != nil {
size = *obj.Size
if !aws.ToBool(page.IsTruncated) {
return keys, nil
}
next := aws.ToString(page.NextContinuationToken)
if next == "" {
return nil, errors.New("truncated R2 list response has no continuation token")
}
if _, exists := seenTokens[next]; exists {
return nil, errors.Errorf("R2 list continuation token %q repeated", next)
}
seenTokens[next] = struct{}{}
continuationToken = aws.String(next)
}
md5 := ""
if obj.ETag != nil {
md5 = strings.Trim(*obj.ETag, `"`)
}
*manifest = append(*manifest, imageManifestEntry{Key: key, Size: size, MD5: md5})
*count++
*totalBytes += size
func planBackupDownloads(imageDir string, keys []string) ([]backupDownloadPlan, error) {
plans := make([]backupDownloadPlan, 0, len(keys))
files := make(map[string]string, len(keys))
directories := make(map[string]string)
root, err := filepath.Abs(imageDir)
if err != nil {
return nil, errors.Wrap(err, "resolve images directory")
}
directories[canonicalLocalPath(root)] = ""
for _, key := range keys {
destination, err := localImagePath(root, key)
if err != nil {
return nil, err
}
canonicalDestination := canonicalLocalPath(destination)
if previous, exists := files[canonicalDestination]; exists {
return nil, errors.Errorf("objects %q and %q map to the same local file", previous, key)
}
if previous, exists := directories[canonicalDestination]; exists {
return nil, errors.Errorf("file/directory collision between objects %q and %q", previous, key)
}
for parent := filepath.Dir(destination); pathWithinDirectory(root, parent); parent = filepath.Dir(parent) {
canonicalParent := canonicalLocalPath(parent)
if previous, exists := files[canonicalParent]; exists {
return nil, errors.Errorf("file/directory collision between objects %q and %q", previous, key)
}
if _, exists := directories[canonicalParent]; !exists {
directories[canonicalParent] = key
}
if parent == root {
break
}
}
files[canonicalDestination] = key
plans = append(plans, backupDownloadPlan{key: key, destination: destination})
}
if err := prepareBackupImageRoot(root); err != nil {
return nil, err
}
for _, plan := range plans {
if err := preflightBackupLocalPath(root, plan.destination); err != nil {
return nil, err
}
}
return plans, nil
}
func validateObjectKey(key string) error {
if key == "" {
return errors.New("R2 object key must not be empty")
}
if strings.ContainsRune(key, 0) {
return errors.Errorf("R2 object key %q contains NUL", key)
}
if strings.Contains(key, `\`) {
return errors.Errorf("R2 object key %q contains a backslash", key)
}
if strings.HasSuffix(key, "/") {
return errors.Errorf("R2 object key %q is a directory key", key)
}
if strings.HasPrefix(key, "/") || filepath.IsAbs(filepath.FromSlash(key)) || filepath.VolumeName(filepath.FromSlash(key)) != "" ||
(len(key) >= 2 && ((key[0] >= 'A' && key[0] <= 'Z') || (key[0] >= 'a' && key[0] <= 'z')) && key[1] == ':') {
return errors.Errorf("R2 object key %q is an absolute or volume-qualified path", key)
}
for _, segment := range strings.Split(key, "/") {
if segment == "" || segment == "." || segment == ".." {
return errors.Errorf("R2 object key %q contains an empty or dot path segment", key)
}
}
return nil
}
// backupS3Client 构造用于备份的 R2 客户端。
//
// 不能直接复用 storage.GetS3Client 做 List：那个客户端把 bucket 路径拼进了 endpoint
// （image_endpoint 形如 https://<account>.r2.cloudflarestorage.com/<bucket>），
// 它对 GetObject/PutObject 恰好可用（R2 容错），但 ListObjectsV2 会把 URL 拼成
// /<bucket>/<bucket>?list-type=2，R2 找不到这个 key 前缀而返回 404 NoSuchKey。
//
// 这里剥离 endpoint 的 bucket 后缀，得到纯 account endpoint，并启用 path-style：
// List/Get 都用 Bucket 参数 + 完整 key（List 返回的 key 自带 bucket 名作前缀，
// 如 tamabox/2025/12/<ts>.jpg），这是 R2 上唯一稳定可用的姿势。
func backupS3Client() (*s3.Client, error) {
endpoint := strings.TrimRight(conf.Upload.ImageEndpoint, "/")
// 剥离最后一个路径段（bucket 名）。endpoint 形如
// https://<account>.r2.cloudflarestorage.com/<bucket>，去掉 /<bucket>。
if i := strings.LastIndex(endpoint, "/"); i > len("https://") {
endpoint = endpoint[:i]
func localImagePath(imageDir, key string) (string, error) {
if err := validateObjectKey(key); err != nil {
return "", err
}
root, err := filepath.Abs(imageDir)
if err != nil {
return "", errors.Wrap(err, "resolve images directory")
}
destination := filepath.Join(root, filepath.FromSlash(key))
if !pathWithinDirectory(root, destination) || destination == root {
return "", errors.Errorf("R2 object key %q escapes the images directory", key)
}
return destination, nil
}
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(
func pathWithinDirectory(root, candidate string) bool {
rel, err := filepath.Rel(root, candidate)
if err != nil || filepath.IsAbs(rel) {
return false
}
return rel == "." || rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))
}
func canonicalLocalPath(path string) string {
cleaned := filepath.Clean(path)
if runtime.GOOS == "windows" {
return strings.ToLower(cleaned)
}
return cleaned
}
func prepareBackupImageRoot(imageDir string) error {
info, err := os.Lstat(imageDir)
if os.IsNotExist(err) {
if err := os.MkdirAll(imageDir, 0o700); err != nil {
return errors.Wrap(err, "create images directory")
}
info, err = os.Lstat(imageDir)
}
if err != nil {
return errors.Wrap(err, "inspect images directory")
}
if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
return errors.Errorf("images path %q must be a real directory, not a symlink or file", imageDir)
}
return nil
}
func preflightBackupLocalPath(imageDir, destination string) error {
info, err := os.Lstat(destination)
if err == nil && (info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular()) {
return errors.Errorf("backup destination %q is not a regular file", destination)
}
if err != nil && !os.IsNotExist(err) {
return errors.Wrapf(err, "inspect backup destination %q", destination)
}
for parent := filepath.Dir(destination); parent != imageDir; parent = filepath.Dir(parent) {
if !pathWithinDirectory(imageDir, parent) {
return errors.Errorf("backup destination %q escapes images directory", destination)
}
info, err := os.Lstat(parent)
if os.IsNotExist(err) {
continue
}
if err != nil {
return errors.Wrapf(err, "inspect backup parent %q", parent)
}
if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
return errors.Errorf("backup parent %q is not a real directory", parent)
}
}
return nil
}
func ensureBackupDestinationDirectory(imageDir, destination string) error {
root, err := filepath.Abs(imageDir)
if err != nil {
return err
}
parent := filepath.Dir(destination)
rel, err := filepath.Rel(root, parent)
if err != nil || !pathWithinDirectory(root, parent) {
return errors.Errorf("invalid backup destination %q", destination)
}
if rel == "." {
return nil
}
current := root
for _, segment := range strings.Split(rel, string(filepath.Separator)) {
current = filepath.Join(current, segment)
info, err := os.Lstat(current)
if os.IsNotExist(err) {
if err := os.Mkdir(current, 0o700); err != nil && !os.IsExist(err) {
return errors.Wrapf(err, "create backup directory %q", current)
}
info, err = os.Lstat(current)
}
if err != nil {
return errors.Wrapf(err, "inspect backup directory %q", current)
}
if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
return errors.Errorf("backup parent %q is not a real directory", current)
}
}
return nil
}
// backupS3Client 与上传、恢复共用 storage.NewS3Client 的 endpoint 语义：
// endpoint 可不带路径或只带 /<bucket>，请求始终以 Bucket 参数和原始完整 key 寻址。
// List 返回什么 key 就下载什么 key，不额外添加或剥离 bucket 前缀。
func backupS3Client(ctx context.Context) (*s3.Client, error) {
return storage.NewS3Client(
ctx,
conf.Upload.ImageEndpoint,
conf.Upload.ImageBucket,
conf.Upload.ImageAccessID,
conf.Upload.ImageAccessSecret,
"",
)),
config.WithRegion("auto"),
)
if err != nil {
return nil, errors.Wrap(err, "load R2 config")
}
return s3.NewFromConfig(cfg, func(o *s3.Options) {
o.UsePathStyle = true
}), nil
}
// downloadObject 下载单个对象到本地文件，失败时删除半成品避免留下损坏文件。
func downloadObject(ctx context.Context, client *s3.Client, bucket, key, dst string) error {
// downloadObject 原子下载单个对象，并返回实际落盘内容的 SHA-256 和 GET 元数据。
// 每次尝试都有独立 60 秒 context；流中断会从头重试整个对象，最多 3 次。
func downloadObject(ctx context.Context, client backupS3API, bucket, key, destination string) (downloadedImageObject, error) {
var lastErr error
for attempt := 0; attempt < imageObjectMaxAttempts; attempt++ {
if err := ctx.Err(); err != nil {
return downloadedImageObject{}, err
}
attemptCtx, cancel := context.WithTimeout(ctx, imageObjectAttemptTimeout)
downloaded, err := downloadObjectOnce(attemptCtx, client, bucket, key, destination)
cancel()
if err == nil {
return downloaded, nil
}
lastErr = err
if ctx.Err() != nil {
return downloadedImageObject{}, ctx.Err()
}
if attempt == imageObjectMaxAttempts-1 {
break
}
if err := waitForImageObjectRetry(ctx, imageObjectRetryBackoff<<attempt); err != nil {
return downloadedImageObject{}, err
}
}
return downloadedImageObject{}, errors.Wrapf(lastErr, "download %q after %d attempts", key, imageObjectMaxAttempts)
}
func downloadObjectOnce(ctx context.Context, client backupS3API, bucket, key, destination string) (downloadedImageObject, error) {
out, err := client.GetObject(ctx, &s3.GetObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(key),
})
if err != nil {
return errors.Wrapf(err, "get R2 object %q", key)
return downloadedImageObject{}, errors.Wrapf(err, "get R2 object %q", key)
}
if out == nil || out.Body == nil {
return downloadedImageObject{}, errors.Errorf("get R2 object %q returned an empty response body", key)
}
defer out.Body.Close()
tmp := dst + ".part"
f, err := os.OpenFile(tmp, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
hash := sha256.New()
var size int64
writeErr := writeFileAtomically(destination, func(w io.Writer) error {
var err error
size, err = io.Copy(io.MultiWriter(w, hash), out.Body)
if err != nil {
return errors.Wrapf(err, "create %q", tmp)
}
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)
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
}
}
internal/cmd/backup_restore.go
+437
View File
@@ -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)
}
internal/cmd/backup_restore_integration_test.go
+184
View File
@@ -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 下对象数不应触发分页")
}
internal/cmd/backup_restore_test.go
+603
View File
@@ -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)
}
internal/cmd/backup_test.go
+96
View File
@@ -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")
}
internal/cmd/ban.go
+3 -1
View File
@@ -95,7 +95,9 @@ func runBan(ctx *cli.Context) error {
return errors.Wrap(err, "ban user and cleanup")
}
logrus.WithField("user_id", user.ID).Info("用户已封禁，关联数据已清理")
// 数据库事务已经提交，先明确输出这一事实，再进行非事务性的 R2 清理。
logrus.WithField("user_id", user.ID).Info("数据库封禁与关联数据清理已提交")
drainR2OutboxAfterBan(ctx.Context)
return nil
}
internal/cmd/r2_outbox.go
+365
View File
@@ -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)
}
internal/cmd/r2_outbox_test.go
+270
View File
@@ -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))
}
}
internal/cmd/web.go
+10 -1
View File
@@ -10,6 +10,7 @@ import (
"net/http"
"os"
"os/signal"
"sync"
"syscall"
"time"
@@ -209,7 +210,15 @@ func runWeb(ctx *cli.Context) error {
// 「提问已 INSERT、图片还没传完」「回答已写入、配额还没扣」这类半成品。
// systemd Restart=always 与 docker restart 都救不了这个。
runCtx, stop := signal.NotifyContext(ctx.Context, syscall.SIGTERM, os.Interrupt)
defer stop()
var r2WorkerWG sync.WaitGroup
defer func() {
// 先取消 worker 的 context，再等待当前一轮退出，避免无人等待的后台删除。
stop()
r2WorkerWG.Wait()
}()
// db.Init 已成功；R2 未完整配置时此函数直接返回，不影响 Web 启动。
startR2OutboxWorker(runCtx, &r2WorkerWG)
serveErr := make(chan error, 1)
go func() {
internal/db/db.go
+2 -1
View File
@@ -36,7 +36,7 @@ func Ping(ctx context.Context) error {
var AllTables = []interface{}{
&User{}, &Question{}, &FollowUp{}, &CensorLog{}, &UploadImage{}, &UploadImageQuestion{},
&Conversation{}, &Message{}, &MessageImage{}, &UserBlock{},
&Conversation{}, &Message{}, &MessageImage{}, &UserBlock{}, &R2DeleteTask{},
// 匿名投票
&Vote{}, &VoteOption{}, &VoteFormField{}, &VoteSubmission{},
// 功能开关（提问 / 私聊 / 投票 的开关、定时、定量收集）
@@ -132,6 +132,7 @@ func Init(typ, dsn string) (*gorm.DB, error) {
Conversations = NewConversationStore(db)
Messages = NewMessageStore(db)
MessageImages = NewMessageImageStore(db)
R2DeleteTasks = NewR2DeleteTaskStore(db)
Votes = NewVotesStore(db)
VoteOptions = NewVoteOptionsStore(db)
VoteFormFields = NewVoteFormFieldsStore(db)
internal/db/message_image.go
+3
View File
@@ -26,6 +26,7 @@ type MessageImage struct {
gorm.Model
MessageID uint `gorm:"index" json:"message_id"`
URL string `gorm:"size:500" json:"url"`
ObjectKey string `gorm:"type:text" json:"-"`
FileName string `gorm:"size:200" json:"file_name"`
FileSize int64 `json:"file_size"`
}
@@ -33,6 +34,7 @@ type MessageImage struct {
type CreateMessageImageOptions struct {
MessageID uint
URL string
ObjectKey string
FileName string
FileSize int64
}
@@ -41,6 +43,7 @@ func (db *messageImages) Create(ctx context.Context, opts CreateMessageImageOpti
img := &MessageImage{
MessageID: opts.MessageID,
URL: opts.URL,
ObjectKey: opts.ObjectKey,
FileName: opts.FileName,
FileSize: opts.FileSize,
}
internal/db/message_image_test.go
+39
View File
@@ -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")
}
internal/db/r2_delete_task.go
+277
View File
@@ -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)
}
internal/db/r2_delete_task_test.go
+141
View File
@@ -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)
}
internal/db/upload_images.go
+60 -16
View File
@@ -10,6 +10,8 @@ import (
"github.com/pkg/errors"
"github.com/samber/lo"
"gorm.io/gorm"
"github.com/tamakyi/TamaBox/internal/conf"
)
var UploadImages UploadImagesStore
@@ -23,12 +25,8 @@ type UploadImagesStore interface {
// 用于账号导出等场景，将 N 次单题查询合并为少量 IN 查询，避免 N+1。
GetByQuestionIDs(ctx context.Context, questionIDs []uint) (map[uint][]*UploadImage, error)
GetByTypeQuestionID(ctx context.Context, typ UploadImageQuestionType, questionID uint) ([]*UploadImage, error)
// DeleteByQuestionID 删除某条提问时级联清理其全部上传图片记录与关联行。
//
// 用 Unscoped 真删：图片元数据（name/md5/key）随提问一并真删，与 questions.DeleteByID
// 对主表的真删口径一致。软删会留下带 deleted_at 的图片行，且关联行的 question_id
// 指向已不存在的提问。注意这里只清数据库记录，对象存储里的实际文件不在本函数删除
// （图片可能被多条记录复用，避免误删共享资源）。
// DeleteByQuestionID 先删除当前题目的关联，再物理删除不再被任何题目引用的主体。
// 实际删除的主体的对象 key 与元数据删除在同一事务中登记 R2 outbox。
DeleteByQuestionID(ctx context.Context, questionID uint) error
}
@@ -87,9 +85,24 @@ func (db *uploadImages) Create(ctx context.Context, opts CreateUploadImageOption
}
if opts.IsDeletingPrevious {
if err := tx.WithContext(ctx).Where("type = ? AND question_id = ?", opts.Type, opts.QuestionID).Delete(&UploadImageQuestion{}).Error; err != nil {
var previousLinks []*UploadImageQuestion
if err := tx.WithContext(ctx).
Where("type = ? AND question_id = ?", opts.Type, opts.QuestionID).
Find(&previousLinks).Error; err != nil {
return errors.Wrap(err, "get previous image question links")
}
previousImageIDs := lo.Map(previousLinks, func(item *UploadImageQuestion, _ int) uint {
return item.UploadImageID
})
if err := tx.WithContext(ctx).
Where("type = ? AND question_id = ?", opts.Type, opts.QuestionID).
Delete(&UploadImageQuestion{}).Error; err != nil {
return errors.Wrap(err, "delete previous image question link")
}
// 新图已经创建但尚未关联；即使它与旧图使用相同 key，也绝不加入删除队列。
if err := deleteUnreferencedUploadImages(ctx, tx, previousImageIDs, opts.Key); err != nil {
return err
}
}
if err := tx.WithContext(ctx).Create(&UploadImageQuestion{
Type: opts.Type,
@@ -167,10 +180,44 @@ func (db *uploadImages) GetByTypeQuestionID(ctx context.Context, typ UploadImage
return db.getBy(ctx, "type = ? AND question_id = ?", typ, questionID)
}
// DeleteByQuestionID 删除某条提问时级联清理其全部上传图片记录（软删除，与提问保持一致）。
// 注意：question_id 列在关联表 upload_image_questions 上，不在 upload_images 表上，
// 需先查关联行拿到 imageID 列表，再软删图片主体并删除关联行（同一事务内完成）。
// 注：仅清理数据库记录，对象存储中的实际文件不在此处删除（避免误删复用资源）。
// unreferencedUploadImagePredicate 是所有 UploadImage 删除入口共用的共享保护谓词。
// 调用方必须先解除当前作用域的关联，再使用该谓词收集并删除真正的孤儿主体。
const unreferencedUploadImagePredicate = `NOT EXISTS (
SELECT 1 FROM upload_image_questions AS shared_links
WHERE shared_links.upload_image_id = upload_images.id
)`
func deleteUnreferencedUploadImages(ctx context.Context, tx *gorm.DB, imageIDs []uint, excludedKeys ...string) error {
if len(imageIDs) == 0 {
return nil
}
var images []*UploadImage
if err := tx.WithContext(ctx).Unscoped().
Where("id IN ?", imageIDs).
Where(unreferencedUploadImagePredicate).
Find(&images).Error; err != nil {
return errors.Wrap(err, "get unreferenced upload images")
}
if len(images) == 0 {
return nil
}
ids := lo.Map(images, func(image *UploadImage, _ int) uint { return image.ID })
if err := tx.WithContext(ctx).Unscoped().
Where("id IN ?", ids).
Where(unreferencedUploadImagePredicate).
Delete(&UploadImage{}).Error; err != nil {
return errors.Wrap(err, "delete unreferenced upload images")
}
values := lo.Map(images, func(image *UploadImage, _ int) string { return image.Key })
if err := enqueueR2DeleteTasksForValues(tx.WithContext(ctx), conf.Upload.ImageBucket, values, excludedKeys...); err != nil {
return errors.Wrap(err, "enqueue deleted upload images")
}
return nil
}
func (db *uploadImages) DeleteByQuestionID(ctx context.Context, questionID uint) error {
return db.Transaction(func(tx *gorm.DB) error {
var links []*UploadImageQuestion
@@ -184,13 +231,10 @@ func (db *uploadImages) DeleteByQuestionID(ctx context.Context, questionID uint)
imageIDs := lo.Map(links, func(item *UploadImageQuestion, _ int) uint {
return item.UploadImageID
})
if err := tx.WithContext(ctx).Where("id IN ?", imageIDs).Unscoped().Delete(&UploadImage{}).Error; err != nil {
return errors.Wrap(err, "delete upload images by question id")
}
// 关联表 UploadImageQuestion 无 DeletedAt，Delete 本就是真删，无需 Unscoped。
// 关联表无 DeletedAt，Delete 本就是真删；先解除当前题目关联。
if err := tx.WithContext(ctx).Where("question_id = ?", questionID).Delete(&UploadImageQuestion{}).Error; err != nil {
return errors.Wrap(err, "delete image question links")
}
return nil
return deleteUnreferencedUploadImages(ctx, tx, imageIDs)
})
}
internal/db/upload_images_outbox_test.go
+95
View File
@@ -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)
}
internal/db/users.go
+61 -18
View File
@@ -386,10 +386,25 @@ func (db *users) Update(ctx context.Context, id uint, opts UpdateUserOptions) er
"Language": opts.Language,
}
if err := db.WithContext(ctx).Model(&User{}).Where("id = ?", id).Updates(updates).Error; err != nil {
return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
var current User
if err := tx.WithContext(ctx).Select("id", "avatar", "background", "background_image").First(&current, id).Error; err != nil {
return errors.Wrap(err, "get current profile images")
}
if err := tx.WithContext(ctx).Model(&User{}).Where("id = ?", id).Updates(updates).Error; err != nil {
return errors.Wrap(err, "update user")
}
oldValues := make([]string, 0, 3)
for _, pair := range [][2]string{{current.Avatar, opts.Avatar}, {current.Background, opts.Background}, {current.BackgroundImage, opts.BackgroundImage}} {
if pair[0] != pair[1] {
oldValues = append(oldValues, pair[0])
}
}
if err := enqueueR2DeleteTasksForValues(tx.WithContext(ctx), conf.Upload.ImageBucket, oldValues, opts.Avatar, opts.Background, opts.BackgroundImage); err != nil {
return errors.Wrap(err, "enqueue replaced profile images")
}
return nil
})
}
// UpdateHarassmentType 仅更新骚扰防护类型（none / register_only）。
@@ -1021,6 +1036,19 @@ func (db *users) BanUserAndCleanup(ctx context.Context, userID uint) error {
// 直接复用会把上一条语句的 Where 条件串到下一条上（曾把 messages 的条件串进 follow_ups，
// 报「字段 conversation_id 不存在」）。
func purgeUserContent(ctx context.Context, tx *gorm.DB, userID uint) error {
var profile User
if err := tx.WithContext(ctx).Select("id", "avatar", "background", "background_image").First(&profile, userID).Error; err != nil {
return errors.Wrap(err, "get profile images of user")
}
if err := enqueueR2DeleteTasksForValues(tx.WithContext(ctx), conf.Upload.ImageBucket, []string{profile.Avatar, profile.Background, profile.BackgroundImage}); err != nil {
return errors.Wrap(err, "enqueue profile images of user")
}
if err := tx.WithContext(ctx).Model(&User{}).Where("id = ?", userID).Updates(map[string]interface{}{
"avatar": "", "background": "", "background_image": "",
}).Error; err != nil {
return errors.Wrap(err, "clear profile images of user")
}
// 追问：属于该用户提问下的追问
if err := tx.WithContext(ctx).Unscoped().
Where("question_id IN (SELECT id FROM questions WHERE user_id = ?)", userID).
@@ -1029,11 +1057,31 @@ func purgeUserContent(ctx context.Context, tx *gorm.DB, userID uint) error {
}
// 消息图片必须先于消息删除，否则消息子查询将无法再匹配待删图片。
// 新数据优先使用物化 ObjectKey，旧数据回退到 URL；外部 URL 不会进入 R2 outbox。
messageImageScope := "message_id IN (SELECT id FROM messages WHERE conversation_id IN (SELECT id FROM conversations WHERE owner_id = ? OR asker_id = ?))"
var messageImages []*MessageImage
if err := tx.WithContext(ctx).Unscoped().
Where("message_id IN (SELECT id FROM messages WHERE conversation_id IN (SELECT id FROM conversations WHERE owner_id = ? OR asker_id = ?))", userID, userID).
Where(messageImageScope, userID, userID).
Find(&messageImages).Error; err != nil {
return errors.Wrap(err, "get message images of user")
}
messageImageValues := make([]string, 0, len(messageImages))
for _, image := range messageImages {
value := image.ObjectKey
if value == "" {
value = image.URL
}
messageImageValues = append(messageImageValues, value)
}
if err := tx.WithContext(ctx).Unscoped().
Where(messageImageScope, userID, userID).
Delete(&MessageImage{}).Error; err != nil {
return errors.Wrap(err, "delete message images of user")
}
if err := enqueueR2DeleteTasksForValues(tx.WithContext(ctx), conf.Upload.ImageBucket, messageImageValues); err != nil {
return errors.Wrap(err, "enqueue deleted message images")
}
// 消息：属于该用户作为 owner 或 asker 的会话
if err := tx.WithContext(ctx).Unscoped().
Where("conversation_id IN (SELECT id FROM conversations WHERE owner_id = ? OR asker_id = ?)", userID, userID).
@@ -1048,30 +1096,25 @@ func purgeUserContent(ctx context.Context, tx *gorm.DB, userID uint) error {
return errors.Wrap(err, "delete conversations of user")
}
// 提问图片元数据：删除本用户上传的孤立图片，以及仅链接到本用户提问的图片。
// 仍被其他提问引用的图片必须保留；这里只清数据库元数据，不操作对象存储。
if err := tx.WithContext(ctx).Unscoped().
Where(`(
uploader_user_id = ?
OR upload_images.id IN (
// 先确定本用户上传或仅与当前作用域有关的候选主体，再解除本用户题目的关联。
// 最终是否可删除由与 UploadImages.DeleteByQuestionID 完全相同的共享保护谓词决定。
var uploadImageIDs []uint
if err := tx.WithContext(ctx).Unscoped().Model(&UploadImage{}).
Where(`uploader_user_id = ? OR upload_images.id IN (
SELECT upload_image_id FROM upload_image_questions
WHERE question_id IN (SELECT id FROM questions WHERE user_id = ?)
)
) AND NOT EXISTS (
SELECT 1 FROM upload_image_questions AS shared_links
WHERE shared_links.upload_image_id = upload_images.id
AND shared_links.question_id NOT IN (SELECT id FROM questions WHERE user_id = ?)
)`, userID, userID, userID).
Delete(&UploadImage{}).Error; err != nil {
return errors.Wrap(err, "delete upload images of user")
)`, userID, userID).
Pluck("upload_images.id", &uploadImageIDs).Error; err != nil {
return errors.Wrap(err, "get upload images of user")
}
// 提问-图片关联
if err := tx.WithContext(ctx).Unscoped().
Where("question_id IN (SELECT id FROM questions WHERE user_id = ?)", userID).
Delete(&UploadImageQuestion{}).Error; err != nil {
return errors.Wrap(err, "delete upload image links of user")
}
if err := deleteUnreferencedUploadImages(ctx, tx, uploadImageIDs); err != nil {
return errors.Wrap(err, "delete upload images of user")
}
// 提问
if err := tx.WithContext(ctx).Unscoped().
internal/db/users_outbox_test.go
+142
View File
@@ -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)
}
internal/route/chat/chat_handler.go
+32 -1
View File
@@ -277,7 +277,8 @@ func SendMessage(ctx context.Context, cache cache.Cache) {
}
// 上传图片（如有）
var imageURL string
var imageURL, imageKey string
imageAttached := false
if imageFile != nil {
uploadedURL, err := storage.UploadPictureToS3(imageFile, imageHeader)
if err != nil {
@@ -291,6 +292,31 @@ func SendMessage(ctx context.Context, cache cache.Cache) {
return
}
imageURL = uploadedURL
defer func() {
if imageAttached || imageKey != "" {
return
}
if err := db.R2DeleteTasks.EnqueueValues(stdctx.Background(), conf.Upload.ImageBucket, []string{imageURL}); err != nil {
logrus.WithError(err).Error("Failed to enqueue unresolved uploaded message image compensation")
}
}()
parsedKey, owned, err := storage.ObjectKeyFromCDNURL(imageURL, conf.Upload.ImageBucketCDNHost)
if err != nil || !owned || parsedKey == "" {
logrus.WithError(err).WithField("owned", owned).Error("Failed to resolve uploaded message image key")
ctx.SetInternalErrorFlash()
ctx.Redirect(fmt.Sprintf("/_/%s/chat/%s", domain, token))
return
}
imageKey = parsedKey
// 只有消息图片元数据成功落库才标记 attached；此前任意失败出口都补偿新对象。
defer func() {
if imageAttached {
return
}
if err := db.R2DeleteTasks.Enqueue(stdctx.Background(), conf.Upload.ImageBucket, []string{imageKey}); err != nil {
logrus.WithError(err).WithField("key", imageKey).Error("Failed to enqueue uploaded message image compensation")
}
}()
}
// 内容安全检测（仅文本）
@@ -341,11 +367,16 @@ func SendMessage(ctx context.Context, cache cache.Cache) {
if _, err := db.MessageImages.Create(ctx.Request().Context(), db.CreateMessageImageOptions{
MessageID: msg.ID,
URL: imageURL,
ObjectKey: imageKey,
FileName: imageHeader.Filename,
FileSize: imageHeader.Size,
}); err != nil {
logrus.WithError(err).Error("Failed to save message image")
ctx.SetInternalErrorFlash()
ctx.Redirect(fmt.Sprintf("/_/%s/chat/%s", domain, token))
return
}
imageAttached = true
}
// 达到上限时关闭对话（幂等，多实例同时执行也安全）
internal/storage/cdn_url.go
+128
View File
@@ -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
}
internal/storage/cdn_url_test.go
+112
View File
@@ -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)
})
}
}
internal/storage/client.go
+4 -23
View File
@@ -5,9 +5,6 @@ import (
"sync"
"github.com/aliyun/aliyun-oss-go-sdk/oss"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/config"
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/pkg/errors"
@@ -31,29 +28,13 @@ func GetS3Client() (*s3.Client, error) {
s3ClientErr = errors.New("S3 endpoint not configured")
return
}
r2Resolver := aws.EndpointResolverWithOptionsFunc(func(service, region string, options ...interface{}) (aws.Endpoint, error) {
return aws.Endpoint{
URL: conf.Upload.ImageEndpoint,
HostnameImmutable: true,
Source: aws.EndpointSourceCustom,
}, nil
})
cfg, err := config.LoadDefaultConfig(context.Background(),
config.WithEndpointResolverWithOptions(r2Resolver),
config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(
s3Client, s3ClientErr = NewS3Client(
context.Background(),
conf.Upload.ImageEndpoint,
conf.Upload.ImageBucket,
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
}
internal/storage/s3.go
+3 -6
View File
@@ -23,7 +23,6 @@ import (
"github.com/disintegration/imaging"
"github.com/pkg/errors"
"github.com/sirupsen/logrus"
"github.com/wuhan005/gadget"
"github.com/tamakyi/TamaBox/internal/conf"
"github.com/tamakyi/TamaBox/internal/imagecheck"
@@ -189,9 +188,9 @@ func uploadPictureToS3(file multipart.File, fileHeader *multipart.FileHeader, as
return "", errors.Wrap(err, "get S3 client")
}
// 7. 上传到 R2
if err := gadget.Retry(5, func() error {
_, err := client.PutObject(context.Background(), &s3.PutObjectInput{
// 7. 上传到 R2。现有上传 API 不接收调用方 context，故显式标记为 TODO；
// PutObject 内部会为每次尝试设置独立超时，并对可取消 context 做退避。
if err := PutObject(context.TODO(), client, &s3.PutObjectInput{
Bucket: aws.String(conf.Upload.ImageBucket),
Key: aws.String(key),
Body: bytes.NewReader(uploadData),
@@ -200,8 +199,6 @@ func uploadPictureToS3(file multipart.File, fileHeader *multipart.FileHeader, as
// key 含纳秒时间戳且内容不可变 → 可安全启用 immutable 强缓存，
// 避免每次页面加载都回源校验（曾导致头像/二维码 logo 每次重新下载）。
CacheControl: aws.String("public, max-age=31536000, immutable"),
})
return err
}); err != nil {
return "", errors.Wrap(err, "put object after retry")
}
internal/storage/s3_client.go
+112
View File
@@ -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}
}
internal/storage/s3_client_test.go
+126
View File
@@ -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")
}
internal/storage/s3_retry.go
+167
View File
@@ -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
}
internal/storage/s3_retry_test.go
+171
View File
@@ -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())
}
route/question/page.go
+4 -1
View File
@@ -871,7 +871,7 @@ func uploadImage(ctx context.Context, opts uploadImageOptions) error {
if err != nil {
return errors.Wrap(err, "get S3 client")
}
if _, err := client.PutObject(ctx.Request().Context(), &s3.PutObjectInput{
if err := storage.PutObject(ctx.Request().Context(), client, &s3.PutObjectInput{
Bucket: aws.String(conf.Upload.ImageBucket),
Key: aws.String(uploadKey),
Body: bytes.NewReader(uploadData),
@@ -894,6 +894,9 @@ func uploadImage(ctx context.Context, opts uploadImageOptions) error {
IsDeletingPrevious: opts.IsDeletingPrevious,
})
if err != nil {
if enqueueErr := db.R2DeleteTasks.Enqueue(ctx.Request().Context(), conf.Upload.ImageBucket, []string{uploadKey}); enqueueErr != nil {
return errors.Wrapf(enqueueErr, "compensate upload image after metadata failure: %v", err)
}
return errors.Wrap(err, "create upload image")
}
return nil
route/user/profile.go
+17
View File
@@ -5,6 +5,7 @@
package user
import (
stdctx "context"
"encoding/json"
"fmt"
"net/url"
@@ -121,6 +122,17 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) {
return
}
var uploadedProfileImages []string
profileSaved := false
defer func() {
if profileSaved || len(uploadedProfileImages) == 0 {
return
}
if err := db.R2DeleteTasks.EnqueueValues(stdctx.Background(), conf.Upload.ImageBucket, uploadedProfileImages); err != nil {
logrus.WithError(err).Error("Failed to enqueue uploaded profile image compensation")
}
}()
// 处理头像上传（强制 1:1，见 storage.UploadAvatarToS3）
avatarURL := currentUser.Avatar // 默认使用原值
avatarFile, avatarFileHeader, err := ctx.Request().FormFile("avatar")
@@ -146,6 +158,7 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) {
logrus.WithContext(ctx.Request().Context()).WithError(err).Error("Failed to upload avatar")
} else {
avatarURL = uploadedURL
uploadedProfileImages = append(uploadedProfileImages, uploadedURL)
}
}
@@ -177,6 +190,7 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) {
logrus.WithContext(ctx.Request().Context()).WithError(err).Error("Failed to upload background")
} else {
backgroundURL = uploadedURL
uploadedProfileImages = append(uploadedProfileImages, uploadedURL)
}
}
@@ -200,6 +214,7 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) {
logrus.WithContext(ctx.Request().Context()).WithError(err).Error("Failed to upload QRbackgroundImage")
} else {
backgroundImageURL = uploadedURL
uploadedProfileImages = append(uploadedProfileImages, uploadedURL)
}
}
@@ -268,6 +283,7 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) {
}
ctx.SetInternalErrorFlash()
} else if passwordChanged {
profileSaved = true
// 密码修改成功且其余资料也保存成功：强制登出重登。
// 先 Flush 清会话再写入 flash，保证提示能带到登录页显示。
ctx.Session.Flush()
@@ -275,6 +291,7 @@ func UpdateProfile(ctx context.Context, f form.UpdateProfile) {
ctx.Redirect("/login")
return
} else {
profileSaved = true
ctx.SetSuccessFlash(i18n.T(ctx.Lang(), "admin.flash.profileUpdated"))
}
ctx.Redirect("/user/profile")
tamabox.go
+2
View File
@@ -30,6 +30,8 @@ func main() {
cmd.HidePublic,
cmd.Honeypot,
cmd.Backup,
cmd.RestoreImages,
cmd.R2Outbox,
}
if err := app.Run(os.Args); err != nil {
logrus.WithError(err).Fatal("Failed to start application")