hyapp-server/tools/im-c2c-broadcast/main.go

package main

import (
	"bytes"
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/binary"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"hyapp/pkg/configx"
	"hyapp/pkg/tencentim"
)

const (
	sendC2CMsgCommand = "v4/openim/sendmsg"
	confirmSendValue  = "SEND_REAL_TENCENT_IM_C2C"
)

type serviceConfig struct {
	MySQLDSN  string          `yaml:"mysql_dsn"`
	TencentIM tencentIMConfig `yaml:"tencent_im"`
}

type tencentIMConfig struct {
	Enabled         bool   `yaml:"enabled"`
	SDKAppID        int64  `yaml:"sdk_app_id"`
	SecretKey       string `yaml:"secret_key"`
	AdminIdentifier string `yaml:"admin_identifier"`
	Endpoint        string `yaml:"endpoint"`
}

type sendTextC2CRequest struct {
	SyncOtherMachine int              `json:"SyncOtherMachine"`
	FromAccount      string           `json:"From_Account,omitempty"`
	ToAccount        string           `json:"To_Account"`
	MsgRandom        uint32           `json:"MsgRandom"`
	MsgBody          []messageElement `json:"MsgBody"`
}

type messageElement struct {
	MsgType    string         `json:"MsgType"`
	MsgContent textMsgContent `json:"MsgContent"`
}

type textMsgContent struct {
	Text string `json:"Text"`
}

type restResponse struct {
	ActionStatus string `json:"ActionStatus"`
	ErrorInfo    string `json:"ErrorInfo"`
	ErrorCode    int    `json:"ErrorCode"`
	MsgKey       string `json:"MsgKey,omitempty"`
}

type resultLine struct {
	AtMS      int64  `json:"at_ms"`
	Mode      string `json:"mode"`
	ToAccount string `json:"to_account"`
	Status    string `json:"status"`
	ErrorCode int    `json:"error_code,omitempty"`
	ErrorInfo string `json:"error_info,omitempty"`
	MsgKey    string `json:"msg_key,omitempty"`
}

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

func run(ctx context.Context, args []string) error {
	var userConfigPath string
	var imConfigPath string
	var fromAccount string
	var appCode string
	var message string
	var outputPath string
	var limit int
	var sleepMS int
	var timeoutSeconds int
	var send bool
	var confirmSend string
	var includeSender bool

	flags := flag.NewFlagSet("im-c2c-broadcast", flag.ContinueOnError)
	flags.StringVar(&userConfigPath, "user-config", "services/user-service/configs/config.yaml", "user-service YAML config used for mysql_dsn")
	flags.StringVar(&imConfigPath, "im-config", "services/notice-service/configs/config.yaml", "service YAML config used for tencent_im")
	flags.StringVar(&fromAccount, "from-account", "123456", "Tencent IM From_Account")
	flags.StringVar(&appCode, "app-code", "lalu", "users.app_code filter")
	flags.StringVar(&message, "message", "", "plain C2C text message body")
	flags.StringVar(&outputPath, "output", "", "JSONL audit output path")
	flags.IntVar(&limit, "limit", 0, "maximum active users to process; 0 means all")
	flags.IntVar(&sleepMS, "sleep-ms", 25, "sleep between real sends")
	flags.IntVar(&timeoutSeconds, "timeout-seconds", 10, "Tencent REST request timeout")
	flags.BoolVar(&send, "send", false, "perform real Tencent IM sends; omitted means dry-run")
	flags.StringVar(&confirmSend, "confirm-send", "", "must equal "+confirmSendValue+" when --send is set")
	flags.BoolVar(&includeSender, "include-sender", false, "include from-account when it is also an active numeric user")
	if err := flags.Parse(args); err != nil {
		return err
	}

	fromAccount = strings.TrimSpace(fromAccount)
	appCode = strings.TrimSpace(appCode)
	message = strings.TrimSpace(message)
	if fromAccount == "" || appCode == "" || message == "" {
		return errors.New("from-account, app-code and message are required")
	}
	if send && confirmSend != confirmSendValue {
		return fmt.Errorf("--send requires --confirm-send=%s", confirmSendValue)
	}
	if sleepMS < 0 || timeoutSeconds <= 0 || limit < 0 {
		return errors.New("sleep-ms, timeout-seconds and limit are invalid")
	}

	userCfg, err := loadConfig(userConfigPath)
	if err != nil {
		return fmt.Errorf("load user config: %w", err)
	}
	imCfg, err := loadConfig(imConfigPath)
	if err != nil {
		return fmt.Errorf("load im config: %w", err)
	}
	if strings.TrimSpace(userCfg.MySQLDSN) == "" {
		return errors.New("user config mysql_dsn is empty")
	}
	if err := validateTencentIMConfig(imCfg.TencentIM); err != nil {
		return err
	}

	db, err := sql.Open("mysql", userCfg.MySQLDSN)
	if err != nil {
		return err
	}
	defer db.Close()
	if err := db.PingContext(ctx); err != nil {
		return fmt.Errorf("ping mysql: %w", err)
	}

	targets, totalActive, err := listActiveUserIDs(ctx, db, appCode, fromAccount, includeSender, limit)
	if err != nil {
		return err
	}
	mode := "dry-run"
	if send {
		mode = "send"
	}
	if outputPath == "" {
		outputPath = filepath.Join("run", fmt.Sprintf("im-c2c-broadcast-%s.jsonl", time.Now().UTC().Format("20060102-150405")))
	}
	if err := os.MkdirAll(filepath.Dir(outputPath), 0o755); err != nil {
		return err
	}
	output, err := os.Create(outputPath)
	if err != nil {
		return err
	}
	defer output.Close()

	fmt.Printf("mode=%s app_code=%s sdk_app_id=%d endpoint=%s from_account=%s total_active=%d targets=%d output=%s\n",
		mode, appCode, imCfg.TencentIM.SDKAppID, firstNonEmpty(imCfg.TencentIM.Endpoint, tencentim.DefaultEndpoint), fromAccount, totalActive, len(targets), outputPath)
	if len(targets) > 0 {
		fmt.Printf("first_target=%s last_target=%s\n", targets[0], targets[len(targets)-1])
	}
	if !send {
		for _, target := range targets {
			if err := writeResult(output, resultLine{AtMS: time.Now().UTC().UnixMilli(), Mode: mode, ToAccount: target, Status: "dry_run"}); err != nil {
				return err
			}
		}
		fmt.Println("dry_run_complete=true")
		return nil
	}

	client := &http.Client{Timeout: time.Duration(timeoutSeconds) * time.Second}
	for i, target := range targets {
		response, err := sendTencentText(ctx, client, imCfg.TencentIM, fromAccount, target, message)
		line := resultLine{AtMS: time.Now().UTC().UnixMilli(), Mode: mode, ToAccount: target, Status: "ok"}
		if err != nil {
			line.Status = "failed"
			line.ErrorInfo = err.Error()
		} else if response.ErrorCode != 0 || !strings.EqualFold(firstNonEmpty(response.ActionStatus, "OK"), "OK") {
			line.Status = "failed"
			line.ErrorCode = response.ErrorCode
			line.ErrorInfo = response.ErrorInfo
		} else {
			line.MsgKey = response.MsgKey
		}
		if err := writeResult(output, line); err != nil {
			return err
		}
		if line.Status != "ok" {
			fmt.Printf("target=%s status=failed error_code=%d error=%s\n", target, line.ErrorCode, line.ErrorInfo)
		}
		if sleepMS > 0 && i+1 < len(targets) {
			time.Sleep(time.Duration(sleepMS) * time.Millisecond)
		}
	}

	fmt.Println("send_complete=true")
	return nil
}

func loadConfig(path string) (serviceConfig, error) {
	var cfg serviceConfig
	err := configx.LoadYAML(strings.TrimSpace(path), &cfg)
	return cfg, err
}

func validateTencentIMConfig(cfg tencentIMConfig) error {
	if cfg.SDKAppID <= 0 || strings.TrimSpace(cfg.SecretKey) == "" || strings.TrimSpace(cfg.AdminIdentifier) == "" {
		return errors.New("tencent_im sdk_app_id, secret_key and admin_identifier are required")
	}
	return nil
}

func listActiveUserIDs(ctx context.Context, db *sql.DB, appCode string, fromAccount string, includeSender bool, limit int) ([]string, int64, error) {
	var totalActive int64
	if err := db.QueryRowContext(ctx, `SELECT COUNT(*) FROM users WHERE app_code = ? AND status = 'active'`, appCode).Scan(&totalActive); err != nil {
		return nil, 0, err
	}

	excludeUserID := int64(0)
	if !includeSender {
		if value, err := strconv.ParseInt(fromAccount, 10, 64); err == nil && value > 0 {
			excludeUserID = value
		}
	}
	query := `SELECT user_id FROM users WHERE app_code = ? AND status = 'active'`
	params := []any{appCode}
	if excludeUserID > 0 {
		query += ` AND user_id <> ?`
		params = append(params, excludeUserID)
	}
	query += ` ORDER BY user_id ASC`
	if limit > 0 {
		query += ` LIMIT ?`
		params = append(params, limit)
	}
	rows, err := db.QueryContext(ctx, query, params...)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()

	targets := make([]string, 0)
	for rows.Next() {
		var userID int64
		if err := rows.Scan(&userID); err != nil {
			return nil, 0, err
		}
		targets = append(targets, strconv.FormatInt(userID, 10))
	}
	return targets, totalActive, rows.Err()
}

func sendTencentText(ctx context.Context, client *http.Client, cfg tencentIMConfig, fromAccount string, toAccount string, text string) (restResponse, error) {
	requestBody := sendTextC2CRequest{
		SyncOtherMachine: 2,
		FromAccount:      fromAccount,
		ToAccount:        toAccount,
		MsgRandom:        randomUint32(),
		MsgBody: []messageElement{
			{MsgType: "TIMTextElem", MsgContent: textMsgContent{Text: text}},
		},
	}
	body, err := json.Marshal(requestBody)
	if err != nil {
		return restResponse{}, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint(cfg, sendC2CMsgCommand), bytes.NewReader(body))
	if err != nil {
		return restResponse{}, err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		return restResponse{}, err
	}
	defer resp.Body.Close()
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return restResponse{}, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return restResponse{}, fmt.Errorf("tencent im http status %d: %s", resp.StatusCode, string(respBody))
	}
	var parsed restResponse
	if err := json.Unmarshal(respBody, &parsed); err != nil {
		return restResponse{}, err
	}
	return parsed, nil
}

func endpoint(cfg tencentIMConfig, command string) string {
	adminSig, _ := tencentim.GenerateUserSig(tencentim.UserSigConfig{
		SDKAppID:  cfg.SDKAppID,
		SecretKey: cfg.SecretKey,
		TTL:       24 * time.Hour,
	}, strings.TrimSpace(cfg.AdminIdentifier), time.Now().UTC())

	values := url.Values{}
	values.Set("sdkappid", strconv.FormatInt(cfg.SDKAppID, 10))
	values.Set("identifier", strings.TrimSpace(cfg.AdminIdentifier))
	values.Set("usersig", adminSig.UserSig)
	values.Set("random", strconv.FormatUint(uint64(randomUint32()), 10))
	values.Set("contenttype", "json")

	return "https://" + strings.TrimSuffix(firstNonEmpty(cfg.Endpoint, tencentim.DefaultEndpoint), "/") + "/" + strings.TrimPrefix(command, "/") + "?" + values.Encode()
}

func writeResult(output *os.File, line resultLine) error {
	body, err := json.Marshal(line)
	if err != nil {
		return err
	}
	if _, err := output.Write(append(body, '\n')); err != nil {
		return err
	}
	return nil
}

func randomUint32() uint32 {
	var payload [4]byte
	if _, err := rand.Read(payload[:]); err != nil {
		return uint32(time.Now().UnixNano())
	}
	return binary.BigEndian.Uint32(payload[:])
}

func firstNonEmpty(values ...string) string {
	for _, value := range values {
		value = strings.TrimSpace(value)
		if value != "" {
			return value
		}
	}
	return ""
}