438 lines
14 KiB
Go
438 lines
14 KiB
Go
// 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)
|
|
}
|