Files
kefu/im/backend/cmd/import-generated-test-users/main.go
T
2026-09-03 08:38:17 +08:00

555 lines
22 KiB
Go

// Command import-generated-test-users validates a 100-person generated fixture
// set, publishes original/display/thumbnail objects to the configured Tencent
// COS bucket, and atomically creates explicitly labelled non-loginable users.
// It is operator-only and is never called during application startup.
package main
import (
"bytes"
"context"
"crypto/aes"
"crypto/cipher"
"crypto/sha256"
"database/sql"
"encoding/base64"
"encoding/hex"
"encoding/json"
"errors"
"flag"
"fmt"
"image"
"image/color"
"image/draw"
"image/jpeg"
"image/png"
"io"
"net/http"
"net/url"
"os"
"path"
"path/filepath"
"regexp"
"sort"
"strings"
"time"
"github.com/disintegration/imaging"
"github.com/go-sql-driver/mysql"
cos "github.com/tencentyun/cos-go-sdk-v5"
)
const (
batchName = "cn-realistic-18-30-20260902-v1"
recordCount = 100
)
type manifestRecord struct {
ID int `json:"id"`
Gender string `json:"gender"`
Age int `json:"age"`
}
type assetPlan struct {
ID int `json:"id"`
Gender string `json:"gender"`
Age int `json:"age"`
File string `json:"file"`
SHA256 string `json:"sha256"`
OriginalKey string `json:"originalKey"`
DisplayKey string `json:"displayKey"`
ThumbnailKey string `json:"thumbnailKey"`
OriginalURL string `json:"originalUrl"`
DisplayURL string `json:"displayUrl"`
ThumbnailURL string `json:"thumbnailUrl"`
}
type uploadResult struct {
OriginalUploaded bool `json:"originalUploaded"`
DisplayUploaded bool `json:"displayUploaded"`
ThumbnailUploaded bool `json:"thumbnailUploaded"`
}
type city struct {
Code string
Name string
Lat float64
Lng float64
}
var cities = []city{
{"310100", "上海", 31.2304, 121.4737},
{"440100", "广州", 23.1291, 113.2644},
{"440300", "深圳", 22.5431, 114.0579},
{"330100", "杭州", 30.2741, 120.1551},
{"510100", "成都", 30.5728, 104.0668},
{"420100", "武汉", 30.5928, 114.3055},
{"320100", "南京", 32.0603, 118.7969},
{"350200", "厦门", 24.4798, 118.0894},
{"610100", "西安", 34.3416, 108.9398},
{"370200", "青岛", 36.0671, 120.3826},
}
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")
}
for _, raw := range []string{configKey, jwtKey + ":integration-config"} {
if raw == "" {
continue
}
key := sha256.Sum256([]byte(raw))
block, _ := aes.NewCipher(key[:])
gcm, _ := cipher.NewGCM(block)
if len(payload) < gcm.NonceSize() {
continue
}
plain, openErr := gcm.Open(nil, payload[:gcm.NonceSize()], payload[gcm.NonceSize():], nil)
if openErr == 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 HTTPS without credentials, query, or fragment")
}
return u, nil
}
func loadManifest(file string) ([]manifestRecord, error) {
data, err := os.ReadFile(file)
if err != nil {
return nil, err
}
var records []manifestRecord
if err = json.Unmarshal(data, &records); err != nil {
return nil, errors.New("manifest is not valid JSON")
}
if len(records) != recordCount {
return nil, fmt.Errorf("manifest must contain exactly %d records", recordCount)
}
sort.Slice(records, func(i, j int) bool { return records[i].ID < records[j].ID })
female := 0
for index, item := range records {
if item.ID != index+1 || (item.Gender != "male" && item.Gender != "female") || item.Age < 18 || item.Age > 30 {
return nil, fmt.Errorf("invalid manifest record at position %d", index+1)
}
if item.Gender == "female" {
female++
}
}
if female <= recordCount/2 {
return nil, errors.New("manifest must contain more female than male accounts")
}
return records, nil
}
func avatarFilename(item manifestRecord) string {
return fmt.Sprintf("avatar-%03d-%s.png", item.ID, item.Gender)
}
func planAssets(records []manifestRecord, directory, prefix, publicBase string) ([]assetPlan, error) {
if !filepath.IsAbs(directory) {
return nil, errors.New("avatars 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(publicBase); err != nil {
return nil, err
}
base := strings.TrimRight(publicBase, "/")
plans := make([]assetPlan, 0, len(records))
for _, item := range records {
name := avatarFilename(item)
file := filepath.Join(directory, name)
info, err := os.Lstat(file)
if err != nil || !info.Mode().IsRegular() || info.Size() <= 0 || info.Size() > 12<<20 {
return nil, fmt.Errorf("missing or invalid avatar %s", name)
}
data, err := os.ReadFile(file)
if err != nil {
return nil, fmt.Errorf("cannot read avatar %s", name)
}
cfg, err := png.DecodeConfig(bytes.NewReader(data))
if err != nil || cfg.Width < 640 || cfg.Height < 640 || int64(cfg.Width)*int64(cfg.Height) > 24_000_000 {
return nil, fmt.Errorf("avatar must be a valid PNG between 640px and 24MP: %s", name)
}
hash := sha256.Sum256(data)
digest := hex.EncodeToString(hash[:])
stem := path.Join(prefix, "test-users", batchName, fmt.Sprintf("%03d-%s-%s", item.ID, item.Gender, digest[:20]))
plan := assetPlan{
ID: item.ID, Gender: item.Gender, Age: item.Age, File: name, SHA256: digest,
OriginalKey: stem + "-original.png", DisplayKey: stem + "-av1.jpg", ThumbnailKey: stem + "-av1-thumb.jpg",
}
plan.OriginalURL = base + "/" + plan.OriginalKey
plan.DisplayURL = base + "/" + plan.DisplayKey
plan.ThumbnailURL = base + "/" + plan.ThumbnailKey
plans = append(plans, plan)
}
return plans, nil
}
func jpegBytes(src image.Image, quality int) ([]byte, error) {
opaque := image.NewRGBA(src.Bounds())
draw.Draw(opaque, opaque.Bounds(), image.NewUniform(color.White), image.Point{}, draw.Src)
draw.Draw(opaque, opaque.Bounds(), src, src.Bounds().Min, draw.Over)
var out bytes.Buffer
err := jpeg.Encode(&out, opaque, &jpeg.Options{Quality: quality})
return out.Bytes(), err
}
func renderVariants(file string) (original, display, thumbnail []byte, err error) {
original, err = os.ReadFile(file)
if err != nil {
return nil, nil, nil, err
}
src, err := png.Decode(bytes.NewReader(original))
if err != nil {
return nil, nil, nil, err
}
displayImage := imaging.Fit(src, 640, 640, imaging.Lanczos)
thumbImage := imaging.Fit(displayImage, 256, 256, imaging.Lanczos)
display, err = jpegBytes(displayImage, 82)
if err == nil {
thumbnail, err = jpegBytes(thumbImage, 78)
}
return
}
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 were not logged", operation)
}
func verifyPublic(ctx context.Context, client *http.Client, objectURL, contentType string, expected []byte) (int, error) {
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, objectURL, nil)
res, err := client.Do(req)
if err != nil {
return 0, errors.New("public COS request failed")
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return res.StatusCode, fmt.Errorf("public COS URL returned HTTP %d", res.StatusCode)
}
actual, err := io.ReadAll(io.LimitReader(res.Body, int64(len(expected))+1))
if err != nil || !bytes.Equal(actual, expected) || !strings.HasPrefix(res.Header.Get("Content-Type"), contentType) {
return res.StatusCode, errors.New("public COS object content mismatch")
}
return res.StatusCode, nil
}
func uploadObject(ctx context.Context, client *cos.Client, public *http.Client, key, objectURL, contentType string, data []byte) (bool, error) {
uploaded := false
res, err := client.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 false, errors.New("existing COS object differs; no overwrite performed")
}
} else {
var apiErr *cos.ErrorResponse
if !errors.As(err, &apiErr) || apiErr.Code != "NoSuchKey" {
return false, cosFailure("lookup", err)
}
headers := http.Header{}
headers.Set("x-cos-forbid-overwrite", "true")
_, err = client.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 false, cosFailure("upload", err)
}
uploaded = true
}
status, err := verifyPublic(ctx, public, objectURL, contentType, data)
if err != nil && status == http.StatusForbidden && uploaded {
_, aclErr := client.Object.PutACL(ctx, key, &cos.ObjectPutACLOptions{Header: &cos.ACLHeaderOptions{XCosACL: "public-read"}})
if aclErr != nil {
return uploaded, cosFailure("object public-read ACL", aclErr)
}
_, err = verifyPublic(ctx, public, objectURL, contentType, data)
}
return uploaded, err
}
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 nickname(item manifestRecord) string {
surnames := []string{"陈", "林", "周", "许", "苏", "沈", "陆", "顾", "方", "季", "唐", "宋", "夏", "叶", "温", "乔", "江", "程", "简", "安"}
female := []string{"可欣", "清禾", "知夏", "若宁", "语桐", "星遥", "念安", "晚晴", "雨眠", "予柔"}
male := []string{"知远", "景行", "沐川", "星河", "予安", "亦辰", "云舟", "修远", "言澈", "嘉树"}
given := male
if item.Gender == "female" {
given = female
}
return "测试·" + surnames[(item.ID-1)%len(surnames)] + given[((item.ID-1)/len(surnames))%len(given)]
}
func createUsers(ctx context.Context, db *sql.DB, records []manifestRecord, plans []assetPlan, confirmDatabase string) (created, skipped int, ids []int64, err error) {
conn, err := db.Conn(ctx)
if err != nil {
return 0, 0, nil, err
}
defer conn.Close()
var database string
if err = conn.QueryRowContext(ctx, `SELECT DATABASE()`).Scan(&database); err != nil || confirmDatabase == "" || database != confirmDatabase {
return 0, 0, nil, errors.New("database confirmation does not match connected database")
}
var locked int
if err = conn.QueryRowContext(ctx, `SELECT GET_LOCK(?,10)`, "xingyu:testusers:"+batchName).Scan(&locked); err != nil || locked != 1 {
return 0, 0, nil, errors.New("could not acquire generated test-user batch lock")
}
defer func() {
_, _ = conn.ExecContext(context.Background(), `DO RELEASE_LOCK(?)`, "xingyu:testusers:"+batchName)
}()
tx, err := conn.BeginTx(ctx, nil)
if err != nil {
return 0, 0, nil, err
}
defer tx.Rollback()
var batchCount int
if err = tx.QueryRowContext(ctx, `SELECT COUNT(*) FROM users WHERE test_batch=?`, batchName).Scan(&batchCount); err != nil {
return 0, 0, nil, errors.New("test-user schema unavailable; migration 028 is required")
}
if batchCount != 0 && batchCount != recordCount {
return 0, 0, nil, fmt.Errorf("batch has %d users; refusing to modify partial data", batchCount)
}
jobs := []string{"设计师", "工程师", "教师", "摄影师", "产品经理", "运营", "编辑", "自由职业"}
hobbies := []string{"城市漫步", "阅读与咖啡", "跑步与音乐", "旅行与美食", "电影与绘画", "徒步与摄影"}
now := time.Now()
ids = make([]int64, 0, recordCount)
for index, item := range records {
publicID := fmt.Sprintf("TESTIM%06d", item.ID)
gender := 1
if item.Gender == "female" {
gender = 2
}
var id int64
var existingTest bool
var existingBatch string
var deleted sql.NullTime
queryErr := tx.QueryRowContext(ctx, `SELECT id,is_test,test_batch,deleted_at FROM users WHERE public_id=? FOR UPDATE`, publicID).Scan(&id, &existingTest, &existingBatch, &deleted)
if queryErr == nil {
if batchCount != recordCount || !existingTest || existingBatch != batchName || deleted.Valid {
return 0, 0, nil, fmt.Errorf("public ID %s conflicts with an existing user", publicID)
}
var actualGender int
var actualAvatar string
if err = tx.QueryRowContext(ctx, `SELECT gender,avatar_url FROM user_profiles WHERE user_id=?`, id).Scan(&actualGender, &actualAvatar); err != nil || actualGender != gender || actualAvatar != plans[index].DisplayURL {
return 0, 0, nil, fmt.Errorf("existing generated profile %s was modified or is incomplete", publicID)
}
skipped++
} else if queryErr == sql.ErrNoRows && batchCount == 0 {
result, insertErr := tx.ExecContext(ctx, `INSERT INTO users(public_id,password_hash,is_test,test_batch,status) VALUES(?,'!GENERATED_TEST_ACCOUNT_NO_LOGIN',1,?,1)`, publicID, batchName)
if insertErr != nil {
return 0, 0, nil, insertErr
}
id, err = result.LastInsertId()
if err != nil {
return 0, 0, nil, err
}
place := cities[(item.ID-1)%len(cities)]
birthday := now.AddDate(-item.Age, 0, -(item.ID%180 + 1)).Format("2006-01-02")
height := 168 + item.ID%12
if gender == 2 {
height = 158 + item.ID%11
}
bio := "【测试数据】虚构资料与AI生成头像,仅用于产品功能测试,非真实用户。兴趣示例:" + hobbies[(item.ID-1)%len(hobbies)] + "。"
_, err = tx.ExecContext(ctx, `INSERT INTO user_profiles(user_id,nickname,avatar_url,cover_url,gender,birthday,height_cm,city_code,city_name,occupation,bio,profile_score,last_active_at) VALUES(?,?,?,'',?,?,?,?,?,?,?,80,NOW(3))`, id, nickname(item), plans[index].DisplayURL, gender, birthday, height, place.Code, place.Name, jobs[(item.ID-1)%len(jobs)], bio)
if err != nil {
return 0, 0, nil, err
}
_, err = tx.ExecContext(ctx, `INSERT INTO user_privacy_settings(user_id,nearby_visible,distance_visible,online_visible,last_active_visible,allow_stranger_message,allow_profile_visit_record,allow_search) VALUES(?,1,1,1,0,1,0,1)`, id)
if err != nil {
return 0, 0, nil, err
}
lat := place.Lat + float64((item.ID%7)-3)*0.003
lng := place.Lng + float64((item.ID%9)-4)*0.003
_, err = tx.ExecContext(ctx, `INSERT INTO user_location_states(user_id,city_code,location_cell,latitude,longitude,source) VALUES(?,?,'fixture',?,?,'fixture')`, id, place.Code, lat, lng)
if err != nil {
return 0, 0, nil, err
}
created++
} else {
if queryErr == sql.ErrNoRows {
queryErr = errors.New("batch identity mismatch")
}
return 0, 0, nil, queryErr
}
ids = append(ids, id)
}
if err = tx.Commit(); err != nil {
return 0, 0, nil, errors.New("user transaction commit result uncertain; verify database before retry")
}
return created, skipped, ids, nil
}
func run() error {
apply := flag.Bool("apply", false, "upload objects and create the test-user batch")
confirmDB := flag.String("confirm-database", "", "exact production database name")
confirmBucket := flag.String("confirm-bucket", "", "exact configured COS bucket")
manifestFile := flag.String("manifest", "", "absolute path to manifest.json")
avatarsDir := flag.String("avatars-dir", "", "absolute directory containing generated PNG files")
backupDir := flag.String("backup-dir", "", "new scoped backup directory under /www/backup; required with --apply")
flag.Parse()
if os.Getenv("IM_ENV") != "production" {
return errors.New("this scoped importer requires IM_ENV=production")
}
if !filepath.IsAbs(*manifestFile) || !filepath.IsAbs(*avatarsDir) {
return errors.New("manifest and avatars directory must be absolute paths")
}
dsn := os.Getenv("IM_DB_DSN")
dsnConfig, err := mysql.ParseDSN(dsn)
if err != nil || dsn == "" || *confirmDB != "xim" || dsnConfig.DBName != *confirmDB || dsnConfig.Net != "tcp" || dsnConfig.Addr != "127.0.0.1:3306" {
return errors.New("database target confirmation failed")
}
dsnConfig.ParseTime = true
db, err := sql.Open("mysql", dsnConfig.FormatDSN())
if err != nil {
return errors.New("cannot open database")
}
defer db.Close()
ctx, cancel := context.WithTimeout(context.Background(), 45*time.Minute)
defer cancel()
if err = db.PingContext(ctx); err != nil {
return errors.New("database connection failed")
}
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 storage 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 storage provider/bucket 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 is not the configured Tencent COS bucket hostname")
}
if settings["storage.tencent_cos.secret_id"] == "" || settings["storage.tencent_cos.secret_key"] == "" {
return errors.New("COS credentials are not configured")
}
records, err := loadManifest(*manifestFile)
if err != nil {
return err
}
plans, err := planAssets(records, *avatarsDir, settings["storage.object_prefix"], settings["storage.tencent_cos.public_base_url"])
if err != nil {
return err
}
female := 0
for _, item := range records {
if item.Gender == "female" {
female++
}
}
if !*apply {
return json.NewEncoder(os.Stdout).Encode(map[string]any{"dryRun": true, "batch": batchName, "bucket": bucket, "users": len(records), "female": female, "male": len(records) - female, "objects": len(plans) * 3})
}
cleanBackup := filepath.Clean(*backupDir)
if filepath.Dir(cleanBackup) != "/www/backup" || !strings.HasPrefix(filepath.Base(cleanBackup), "xim-generated-users-") {
return errors.New("backup directory must be a new /www/backup/xim-generated-users-* directory")
}
if err = os.Mkdir(cleanBackup, 0700); err != nil {
return err
}
if err = writeJSON(filepath.Join(cleanBackup, "plan.json"), map[string]any{"batch": batchName, "bucket": bucket, "records": records, "objects": plans}); err != nil {
return err
}
fmt.Println("BACKUP=" + cleanBackup)
noRedirect := func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }
public := &http.Client{Timeout: 60 * time.Second, CheckRedirect: noRedirect}
endpoint.Path = ""
cosClient := cos.NewClient(&cos.BaseURL{BucketURL: endpoint}, &http.Client{Timeout: 90 * time.Second, CheckRedirect: noRedirect, Transport: &cos.AuthorizationTransport{SecretID: settings["storage.tencent_cos.secret_id"], SecretKey: settings["storage.tencent_cos.secret_key"]}})
uploads := make([]uploadResult, 0, len(plans))
for _, plan := range plans {
original, display, thumbnail, renderErr := renderVariants(filepath.Join(*avatarsDir, plan.File))
if renderErr != nil {
return fmt.Errorf("%s: cannot render avatar variants; database unchanged", plan.File)
}
originalUploaded, uploadErr := uploadObject(ctx, cosClient, public, plan.OriginalKey, plan.OriginalURL, "image/png", original)
if uploadErr != nil {
return fmt.Errorf("%s original: %w; database unchanged", plan.File, uploadErr)
}
displayUploaded, uploadErr := uploadObject(ctx, cosClient, public, plan.DisplayKey, plan.DisplayURL, "image/jpeg", display)
if uploadErr != nil {
return fmt.Errorf("%s display: %w; database unchanged", plan.File, uploadErr)
}
thumbUploaded, uploadErr := uploadObject(ctx, cosClient, public, plan.ThumbnailKey, plan.ThumbnailURL, "image/jpeg", thumbnail)
if uploadErr != nil {
return fmt.Errorf("%s thumbnail: %w; database unchanged", plan.File, uploadErr)
}
uploads = append(uploads, uploadResult{originalUploaded, displayUploaded, thumbUploaded})
fmt.Printf("VERIFIED %03d/100 original=%t display=%t thumbnail=%t\n", plan.ID, originalUploaded, displayUploaded, thumbUploaded)
}
if err = writeJSON(filepath.Join(cleanBackup, "uploaded.json"), map[string]any{"plans": plans, "results": uploads}); err != nil {
return err
}
created, skipped, ids, err := createUsers(ctx, db, records, plans, *confirmDB)
if err != nil {
return err
}
result := map[string]any{"batch": batchName, "bucket": bucket, "created": created, "skipped": skipped, "female": female, "male": len(records) - female, "objectsVerified": len(plans) * 3, "ids": ids, "backup": cleanBackup, "loginDisabled": true, "globalConfigUnchanged": true}
if err = writeJSON(filepath.Join(cleanBackup, "result.json"), result); err != nil {
return errors.New("users 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)
}
}