fix: 压测客户端改为 MQTT 5 并完成 hello 与收发统计

test/load/provision.go

@@ -0,0 +1,170 @@
+package load
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"time"
+
+	"git.asio.asia/nixevol/NixMsg/test/harness"
+)
+
+const maxImportRows = 1000
+
+type apiEnvelope struct {
+	OK    bool           `json:"ok"`
+	Data  map[string]any `json:"data"`
+	Error *struct {
+		Code    string `json:"code"`
+		Message string `json:"message"`
+	} `json:"error"`
+}
+
+// Provision 按配置准备端：管理 API 批量开通、自助注册，或复用已开通编号。
+func Provision(cfg Config, ids []string) error {
+	if cfg.Reuse && cfg.AdminUser == "" && cfg.AdminPass == "" && cfg.APIToken == "" && cfg.RegisterCode == "" {
+		return nil
+	}
+	if cfg.RegisterCode != "" && cfg.AdminPass == "" && cfg.APIToken == "" {
+		return registerAll(cfg, ids)
+	}
+	if cfg.AdminPass == "" && cfg.APIToken == "" {
+		return fmt.Errorf("开通端需要 -admin-pass / -admin-token，或 -register-code；复用已开通端请加 -reuse")
+	}
+	ac, err := adminLogin(cfg)
+	if err != nil {
+		return err
+	}
+	if cfg.Reuse {
+		return createOneByOne(ac, ids, cfg.Password, true)
+	}
+	return importCSV(ac, ids, cfg.Password)
+}
+
+func adminLogin(cfg Config) (*harness.AdminClient, error) {
+	base := cfg.adminBase()
+	ac, err := harness.NewAdminClient(base)
+	if err != nil {
+		return nil, err
+	}
+	ac.HTTP.Timeout = 10 * time.Minute
+	if cfg.APIToken != "" {
+		ac.APIToken = cfg.APIToken
+		return ac, nil
+	}
+	user := cfg.AdminUser
+	if user == "" {
+		user = "admin"
+	}
+	body, _ := json.Marshal(map[string]string{
+		"username": user,
+		"password": cfg.AdminPass,
+	})
+	resp, err := ac.PostJSON("/api/admin/login", body)
+	if err != nil {
+		return nil, fmt.Errorf("管理员登录: %w", err)
+	}
+	raw, _ := io.ReadAll(resp.Body)
+	_ = resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("管理员登录失败: %d %s", resp.StatusCode, raw)
+	}
+	return ac, nil
+}
+
+func importCSV(ac *harness.AdminClient, ids []string, password string) error {
+	for start := 0; start < len(ids); start += maxImportRows {
+		end := start + maxImportRows
+		if end > len(ids) {
+			end = len(ids)
+		}
+		var b strings.Builder
+		b.WriteString("id,name,login_password,talk_password,default_delay_seconds,remark\n")
+		for _, id := range ids[start:end] {
+			fmt.Fprintf(&b, "%s,%s,%s,,,\n", id, id, password)
+		}
+		resp, err := ac.Do(http.MethodPost, "/api/admin/endpoints/import", []byte(b.String()), "text/csv")
+		if err != nil {
+			return fmt.Errorf("批量开通: %w", err)
+		}
+		raw, _ := io.ReadAll(resp.Body)
+		_ = resp.Body.Close()
+		if resp.StatusCode != http.StatusOK {
+			return fmt.Errorf("批量开通失败: %d %s", resp.StatusCode, raw)
+		}
+	}
+	return nil
+}
+
+func createOneByOne(ac *harness.AdminClient, ids []string, password string, reuse bool) error {
+	for _, id := range ids {
+		body, _ := json.Marshal(map[string]string{
+			"id":             id,
+			"login_password": password,
+		})
+		resp, err := ac.PostJSON("/api/admin/endpoints", body)
+		if err != nil {
+			return fmt.Errorf("开通 %s: %w", id, err)
+		}
+		raw, _ := io.ReadAll(resp.Body)
+		_ = resp.Body.Close()
+		if resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusCreated {
+			continue
+		}
+		if reuse && isTaken(resp.StatusCode, raw) {
+			continue
+		}
+		return fmt.Errorf("开通 %s 失败: %d %s", id, resp.StatusCode, raw)
+	}
+	return nil
+}
+
+func registerAll(cfg Config, ids []string) error {
+	base := strings.TrimRight(cfg.HTTPBase, "/")
+	client := &http.Client{Timeout: 30 * time.Second}
+	for _, id := range ids {
+		payload, _ := json.Marshal(map[string]string{
+			"registration_code": cfg.RegisterCode,
+			"id":                id,
+			"login_password":    cfg.Password,
+		})
+		resp, err := client.Post(base+"/api/client/register", "application/json", bytes.NewReader(payload))
+		if err != nil {
+			return fmt.Errorf("注册 %s: %w", id, err)
+		}
+		raw, _ := io.ReadAll(resp.Body)
+		_ = resp.Body.Close()
+		if resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusCreated {
+			continue
+		}
+		if cfg.Reuse && isTaken(resp.StatusCode, raw) {
+			continue
+		}
+		return fmt.Errorf("注册 %s 失败: %d %s", id, resp.StatusCode, raw)
+	}
+	return nil
+}
+
+func isTaken(status int, raw []byte) bool {
+	if status == http.StatusConflict {
+		return true
+	}
+	var env apiEnvelope
+	if json.Unmarshal(raw, &env) != nil {
+		return false
+	}
+	if env.Error != nil && (env.Error.Code == "id_taken" || env.Error.Code == "conflict") {
+		return true
+	}
+	return false
+}
+
+func (c Config) adminBase() string {
+	if c.AdminHTTP != "" {
+		return strings.TrimRight(c.AdminHTTP, "/")
+	}
+	return strings.TrimRight(c.HTTPBase, "/")
+}