kefu/im/backend/cmd/migrate-test-avatars-cos/main.go

// Migrate only the explicitly labelled fixture batch to the configured COS bucket.
// This command never deletes files, changes bucket ACLs, or updates global config.
package main

import (
	"bytes"
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
	"database/sql"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"image/jpeg"
	"image/png"
	"io"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/disintegration/imaging"
	"github.com/example/xingyu/internal/testusers"
	"github.com/go-sql-driver/mysql"
	cos "github.com/tencentyun/cos-go-sdk-v5"
)

const oldBase = "https://im.bchongw.com/uploads"

type asset struct {
	File                   string `json:"file"`
	Key                    string `json:"key"`
	URL                    string `json:"url"`
	SHA256                 string `json:"sha256"`
	Data                   []byte `json:"-"`
	ThumbnailKey           string `json:"thumbnailKey"`
	ThumbnailURL           string `json:"thumbnailUrl"`
	ThumbnailData          []byte `json:"-"`
	Uploaded               bool   `json:"uploaded"`
	ThumbnailUploaded      bool   `json:"thumbnailUploaded"`
	PublicReadSet          bool   `json:"publicReadSet"`
	ThumbnailPublicReadSet bool   `json:"thumbnailPublicReadSet"`
}

type replacement struct {
	ID       int64  `json:"id"`
	PublicID string `json:"publicId"`
	OldURL   string `json:"oldUrl"`
	NewURL   string `json:"newUrl"`
}

func decrypt(value, configKey, jwtKey string) (string, error) {
	if !strings.HasPrefix(value, "enc:v1:") {
		return value, nil
	}
	payload, err := base64.RawStdEncoding.DecodeString(strings.TrimPrefix(value, "enc:v1:"))
	if err != nil {
		return "", errors.New("invalid encrypted COS credential")
	}
	keys := []string{configKey, jwtKey + ":integration-config"}
	for _, raw := range keys {
		if raw == "" {
			continue
		}
		key := sha256.Sum256([]byte(raw))
		block, _ := aes.NewCipher(key[:])
		gcm, _ := cipher.NewGCM(block)
		if len(payload) < gcm.NonceSize() {
			continue
		}
		plain, err := gcm.Open(nil, payload[:gcm.NonceSize()], payload[gcm.NonceSize():], nil)
		if err == nil {
			return string(plain), nil
		}
	}
	return "", errors.New("cannot decrypt COS credential with server encryption key")
}

func validHTTPS(raw string) (*url.URL, error) {
	u, err := url.Parse(strings.TrimSpace(raw))
	if err != nil || u.Scheme != "https" || u.Host == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" {
		return nil, errors.New("COS endpoint/public base must be an HTTPS URL without credentials or query parameters")
	}
	return u, nil
}

func loadAssets(directory, prefix, base string) ([]*asset, error) {
	if !filepath.IsAbs(directory) {
		return nil, errors.New("media directory must be absolute")
	}
	if !regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9/_-]{0,119}$`).MatchString(prefix) || strings.Contains(prefix, "//") || path.Clean(prefix) != prefix {
		return nil, errors.New("invalid configured object prefix")
	}
	if _, err := validHTTPS(base); err != nil {
		return nil, err
	}
	assets := []*asset{}
	for _, name := range testusers.Avatars() {
		file := filepath.Join(directory, name)
		info, err := os.Lstat(file)
		if err != nil || !info.Mode().IsRegular() || info.Size() > 10<<20 {
			return nil, fmt.Errorf("missing/invalid local avatar: %s", name)
		}
		data, err := os.ReadFile(file)
		if err != nil {
			return nil, fmt.Errorf("cannot read local avatar: %s", name)
		}
		cfg, err := png.DecodeConfig(bytes.NewReader(data))
		if err != nil || cfg.Width < 256 || cfg.Height < 256 {
			return nil, fmt.Errorf("invalid PNG: %s", name)
		}
		hash := sha256.Sum256(data)
		digest := hex.EncodeToString(hash[:])
		key := path.Join(prefix, "test-users", testusers.Batch, digest[:20]+"-"+name)
		decoded, err := png.Decode(bytes.NewReader(data))
		if err != nil {
			return nil, fmt.Errorf("cannot decode PNG: %s", name)
		}
		thumbnail := imaging.Fit(decoded, 256, 256, imaging.Lanczos)
		var thumbnailData bytes.Buffer
		if err = jpeg.Encode(&thumbnailData, thumbnail, &jpeg.Options{Quality: 78}); err != nil {
			return nil, fmt.Errorf("cannot encode thumbnail: %s", name)
		}
		thumbnailKey := strings.TrimSuffix(key, path.Ext(key)) + "-thumb.jpg"
		base = strings.TrimRight(base, "/")
		assets = append(assets, &asset{File: name, Key: key, URL: base + "/" + key, SHA256: digest, Data: data, ThumbnailKey: thumbnailKey, ThumbnailURL: base + "/" + thumbnailKey, ThumbnailData: thumbnailData.Bytes()})
	}
	return assets, nil
}

func planRows(ctx context.Context, db *sql.DB, assets []*asset) ([]replacement, error) {
	fixtures, _ := testusers.Generate(oldBase)
	byPublicID := map[string]string{}
	byFile := map[string]*asset{}
	for _, a := range assets {
		byFile[a.File] = a
	}
	for _, p := range fixtures {
		byPublicID[p.PublicID] = p.AvatarFile
	}
	rows, err := db.QueryContext(ctx, `SELECT u.id,u.public_id,p.avatar_url FROM users u JOIN user_profiles p ON p.user_id=u.id WHERE u.is_test=1 AND u.test_batch=? AND u.deleted_at IS NULL ORDER BY u.id`, testusers.Batch)
	if err != nil {
		return nil, errors.New("cannot load fixture profiles")
	}
	defer rows.Close()
	plan := []replacement{}
	for rows.Next() {
		var r replacement
		if err = rows.Scan(&r.ID, &r.PublicID, &r.OldURL); err != nil {
			return nil, err
		}
		a := byFile[byPublicID[r.PublicID]]
		if a == nil {
			return nil, errors.New("fixture identity differs from expected batch")
		}
		r.NewURL = a.URL
		if r.OldURL != oldBase+"/"+a.File && r.OldURL != r.NewURL {
			return nil, fmt.Errorf("profile %d has an edited avatar; refusing to overwrite it", r.ID)
		}
		plan = append(plan, r)
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}
	if len(plan) != testusers.Count {
		return nil, fmt.Errorf("expected 100 fixture profiles; found %d", len(plan))
	}
	return plan, nil
}

// Verify bytes anonymously, without signed query parameters or Authorization.
func publicObjectCheck(ctx context.Context, client *http.Client, objectURL string, data []byte, contentType string) (int, error) {
	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, objectURL, nil)
	res, err := client.Do(req)
	if err != nil {
		return 0, errors.New("COS public URL request failed")
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		return res.StatusCode, fmt.Errorf("COS public URL returned HTTP %d", res.StatusCode)
	}
	actual, err := io.ReadAll(io.LimitReader(res.Body, int64(len(data))+1))
	if err != nil || !bytes.Equal(actual, data) || !strings.HasPrefix(res.Header.Get("Content-Type"), contentType) {
		return res.StatusCode, errors.New("COS public image content differs from generated asset")
	}
	return res.StatusCode, nil
}

func publicCheck(ctx context.Context, client *http.Client, a *asset) (int, error) {
	return publicObjectCheck(ctx, client, a.URL, a.Data, "image/png")
}

func cosFailure(operation string, err error) error {
	var apiErr *cos.ErrorResponse
	if errors.As(err, &apiErr) {
		return fmt.Errorf("COS %s failed: %s (request %s)", operation, apiErr.Code, apiErr.RequestID)
	}
	return fmt.Errorf("COS %s failed (credentials and request headers not logged)", operation)
}

func uploadObject(ctx context.Context, c *cos.Client, public *http.Client, key, objectURL, contentType string, data []byte, uploaded, publicReadSet *bool) error {
	// Never replace a different existing object, even under this batch prefix.
	res, err := c.Object.Get(ctx, key, nil)
	if err == nil {
		actual, readErr := io.ReadAll(io.LimitReader(res.Body, int64(len(data))+1))
		res.Body.Close()
		if readErr != nil || !bytes.Equal(actual, data) {
			return errors.New("existing COS object differs; no overwrite performed")
		}
	} else {
		var apiErr *cos.ErrorResponse
		if !errors.As(err, &apiErr) || apiErr.Code != "NoSuchKey" {
			return cosFailure("lookup", err)
		}
		headers := http.Header{}
		headers.Set("x-cos-forbid-overwrite", "true")
		_, err = c.Object.Put(ctx, key, bytes.NewReader(data), &cos.ObjectPutOptions{ObjectPutHeaderOptions: &cos.ObjectPutHeaderOptions{ContentType: contentType, ContentLength: int64(len(data)), CacheControl: "public, max-age=31536000, immutable", XOptionHeader: &headers}})
		if err != nil {
			return cosFailure("upload", err)
		}
		*uploaded = true
	}
	status, err := publicObjectCheck(ctx, public, objectURL, data, contentType)
	if err != nil && status == http.StatusForbidden && *uploaded {
		// Only these newly uploaded, explicitly public test avatars may receive
		// object-level public read. Never change the bucket or existing object ACLs.
		_, aclErr := c.Object.PutACL(ctx, key, &cos.ObjectPutACLOptions{Header: &cos.ACLHeaderOptions{XCosACL: "public-read"}})
		if aclErr != nil {
			return cosFailure("test-avatar public-read ACL", aclErr)
		}
		*publicReadSet = true
		_, err = publicObjectCheck(ctx, public, objectURL, data, contentType)
	}
	return err
}

func upload(ctx context.Context, c *cos.Client, public *http.Client, a *asset) error {
	if err := uploadObject(ctx, c, public, a.Key, a.URL, "image/png", a.Data, &a.Uploaded, &a.PublicReadSet); err != nil {
		return err
	}
	if len(a.ThumbnailData) == 0 {
		return nil
	}
	return uploadObject(ctx, c, public, a.ThumbnailKey, a.ThumbnailURL, "image/jpeg", a.ThumbnailData, &a.ThumbnailUploaded, &a.ThumbnailPublicReadSet)
}

func writeJSON(file string, value any) error {
	data, err := json.MarshalIndent(value, "", "  ")
	if err != nil {
		return err
	}
	f, err := os.OpenFile(file, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	defer f.Close()
	if _, err = f.Write(append(data, '\n')); err != nil {
		return err
	}
	return f.Sync()
}

func run() error {
	apply := flag.Bool("apply", false, "upload objects and update fixture links; default is read-only")
	confirmDB := flag.String("confirm-database", "", "exact production database name")
	confirmBucket := flag.String("confirm-bucket", "", "exact configured COS bucket")
	media := flag.String("media-dir", "", "absolute directory containing the ten original fixture PNGs")
	backup := flag.String("backup-dir", "", "new private backup directory under /www/backup, required with --apply")
	flag.Parse()
	if os.Getenv("IM_ENV") != "production" {
		return errors.New("this scoped tool requires IM_ENV=production")
	}
	dsn := os.Getenv("IM_DB_DSN")
	cfg, err := mysql.ParseDSN(dsn)
	if err != nil || dsn == "" || *confirmDB != "im" || cfg.DBName != *confirmDB || cfg.Net != "tcp" || cfg.Addr != "127.0.0.1:3307" {
		return errors.New("database target confirmation failed")
	}
	cfg.ParseTime = true
	db, err := sql.Open("mysql", cfg.FormatDSN())
	if err != nil {
		return errors.New("cannot open database")
	}
	defer db.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()
	settings := map[string]string{}
	rows, err := db.QueryContext(ctx, `SELECT config_key,config_value,value_type FROM system_configs WHERE config_key LIKE 'storage.tencent_cos.%' OR config_key IN ('storage.provider','storage.object_prefix')`)
	if err != nil {
		return errors.New("cannot read COS configuration")
	}
	for rows.Next() {
		var key, value, kind string
		if err = rows.Scan(&key, &value, &kind); err != nil {
			rows.Close()
			return err
		}
		if kind == "secret" {
			value, err = decrypt(value, os.Getenv("IM_CONFIG_ENCRYPTION_KEY"), os.Getenv("IM_JWT_SECRET"))
			if err != nil {
				rows.Close()
				return err
			}
		}
		settings[key] = strings.TrimSpace(value)
	}
	err = rows.Err()
	rows.Close()
	if err != nil {
		return err
	}
	bucket := settings["storage.tencent_cos.bucket"]
	if settings["storage.provider"] != "tencent_cos" || bucket == "" || bucket != *confirmBucket {
		return errors.New("configured COS bucket/provider does not match confirmation")
	}
	endpoint, err := validHTTPS(settings["storage.tencent_cos.endpoint"])
	if err != nil {
		return err
	}
	if endpoint.Port() != "" || !regexp.MustCompile(`^`+regexp.QuoteMeta(bucket)+`\.cos\.[a-z0-9-]+\.myqcloud\.com$`).MatchString(endpoint.Hostname()) || strings.Trim(endpoint.Path, "/") != "" {
		return errors.New("authenticated endpoint must be the configured Tencent COS bucket hostname")
	}
	endpoint.Path = ""
	for _, key := range []string{"storage.tencent_cos.secret_id", "storage.tencent_cos.secret_key"} {
		if settings[key] == "" {
			return errors.New("COS credentials are not configured")
		}
	}
	assets, err := loadAssets(*media, settings["storage.object_prefix"], settings["storage.tencent_cos.public_base_url"])
	if err != nil {
		return err
	}
	plan, err := planRows(ctx, db, assets)
	if err != nil {
		return err
	}
	if !*apply {
		return json.NewEncoder(os.Stdout).Encode(map[string]any{"dryRun": true, "batch": testusers.Batch, "bucket": bucket, "users": len(plan), "objects": assets})
	}
	if filepath.Dir(filepath.Clean(*backup)) != "/www/backup" || !strings.HasPrefix(filepath.Base(*backup), "xingyu-cos-avatars-") {
		return errors.New("backup directory must be a new scoped /www/backup/xingyu-cos-avatars-* directory")
	}
	if err = os.Mkdir(*backup, 0700); err != nil {
		return err
	}
	if err = writeJSON(filepath.Join(*backup, "before.json"), map[string]any{"batch": testusers.Batch, "bucket": bucket, "replacements": plan, "objects": assets}); err != nil {
		return err
	}
	fmt.Println("BACKUP=" + *backup)
	noRedirect := func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }
	public := &http.Client{Timeout: 45 * time.Second, CheckRedirect: noRedirect}
	c := cos.NewClient(&cos.BaseURL{BucketURL: endpoint}, &http.Client{Timeout: 60 * time.Second, CheckRedirect: noRedirect, Transport: &cos.AuthorizationTransport{SecretID: settings["storage.tencent_cos.secret_id"], SecretKey: settings["storage.tencent_cos.secret_key"]}})
	for _, a := range assets {
		if err = upload(ctx, c, public, a); err != nil {
			return fmt.Errorf("%s: %w; database links unchanged", a.File, err)
		}
		fmt.Printf("VERIFIED %s original_uploaded=%t thumbnail_uploaded=%t\n", a.File, a.Uploaded, a.ThumbnailUploaded)
	}
	if err = writeJSON(filepath.Join(*backup, "uploaded.json"), assets); err != nil {
		return err
	}
	// Lock and validate all scoped rows before changing any link. Any concurrent
	// avatar edit causes rollback rather than overwriting a user's change.
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()
	for _, r := range plan {
		var current string
		if err = tx.QueryRowContext(ctx, `SELECT p.avatar_url FROM user_profiles p JOIN users u ON u.id=p.user_id WHERE u.id=? AND u.public_id=? AND u.is_test=1 AND u.test_batch=? AND u.deleted_at IS NULL FOR UPDATE`, r.ID, r.PublicID, testusers.Batch).Scan(&current); err != nil || current != r.OldURL {
			return errors.New("fixture avatar changed during upload; no database links updated")
		}
	}
	changed := int64(0)
	for _, r := range plan {
		if r.OldURL == r.NewURL {
			continue
		}
		res, updateErr := tx.ExecContext(ctx, `UPDATE user_profiles SET avatar_url=? WHERE user_id=? AND avatar_url=?`, r.NewURL, r.ID, r.OldURL)
		if updateErr != nil {
			return errors.New("avatar link update failed; transaction rolled back")
		}
		n, _ := res.RowsAffected()
		if n != 1 {
			return errors.New("avatar update count mismatch; transaction rolled back")
		}
		changed += n
	}
	if err = tx.Commit(); err != nil {
		return errors.New("avatar transaction commit result uncertain; inspect backup and database before retrying")
	}
	verified, err := planRows(ctx, db, assets)
	if err != nil {
		return err
	}
	for _, r := range verified {
		if r.OldURL != r.NewURL {
			return errors.New("post-commit verification failed")
		}
	}
	result := map[string]any{"batch": testusers.Batch, "bucket": bucket, "users": len(verified), "updated": changed, "objects": len(assets) * 2, "publicURLsVerified": len(assets) * 2, "backup": *backup, "localOriginalsKept": true, "globalConfigUnchanged": true}
	if err = writeJSON(filepath.Join(*backup, "result.json"), result); err != nil {
		return errors.New("links committed but result record could not be written; inspect database")
	}
	return json.NewEncoder(os.Stdout).Encode(result)
}

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}