377 lines
11 KiB
Go
377 lines
11 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"crypto/tls"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/url"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type ctrlPacket struct {
|
|
ID string `json:"id"`
|
|
Topic string `json:"topic"`
|
|
Code int `json:"code"`
|
|
Text string `json:"text"`
|
|
Params map[string]json.RawMessage `json:"params"`
|
|
}
|
|
|
|
type dataPacket struct {
|
|
Topic string `json:"topic"`
|
|
From string `json:"from"`
|
|
Timestamp time.Time `json:"ts"`
|
|
Seq int64 `json:"seq"`
|
|
Head map[string]any `json:"head"`
|
|
Content json.RawMessage `json:"content"`
|
|
}
|
|
|
|
type ctrlError struct {
|
|
Code int
|
|
Text string
|
|
}
|
|
|
|
func (e *ctrlError) Error() string {
|
|
return fmt.Sprintf("ctrl code=%d text=%s", e.Code, e.Text)
|
|
}
|
|
|
|
type WSClient struct {
|
|
account AccountRecord
|
|
conn *websocket.Conn
|
|
commandTimeout time.Duration
|
|
onData func(AccountRecord, dataPacket)
|
|
onDisconnect func(AccountRecord, error)
|
|
|
|
sendMu sync.Mutex
|
|
mu sync.Mutex
|
|
seq uint64
|
|
waits map[string]chan ctrlPacket
|
|
done chan struct{}
|
|
closed atomic.Bool
|
|
}
|
|
|
|
func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, apiKey, password string,
|
|
onData func(AccountRecord, dataPacket), onDisconnect func(AccountRecord, error)) (*WSClient, error) {
|
|
timeout, err := configDuration(cfg.CommandTimeout)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
endpoint, err := websocketEndpoint(account.Endpoint, apiKey)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
dialer := *websocket.DefaultDialer
|
|
if cfg.InsecureTLS {
|
|
dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} // #nosec G402: explicit load-test option.
|
|
}
|
|
conn, response, err := dialer.DialContext(ctx, endpoint, http.Header{})
|
|
if err != nil {
|
|
if response != nil {
|
|
return nil, fmt.Errorf("websocket dial status=%d: %w", response.StatusCode, err)
|
|
}
|
|
return nil, err
|
|
}
|
|
client := &WSClient{
|
|
account: account, conn: conn, commandTimeout: timeout, onData: onData, onDisconnect: onDisconnect,
|
|
waits: make(map[string]chan ctrlPacket), done: make(chan struct{}),
|
|
}
|
|
go client.readLoop()
|
|
if _, err = client.command(ctx, "hi", map[string]any{"hi": map[string]any{
|
|
"id": client.nextID("hi"), "ver": "0.25", "ua": "im-pressure/1.0",
|
|
"platf": "loadtest", "tenant": cfg.Tenant,
|
|
}}); err != nil {
|
|
client.Close()
|
|
return nil, fmt.Errorf("hi %s: %w", account.Username, err)
|
|
}
|
|
login := map[string]any{"id": client.nextID("login")}
|
|
if account.Token != "" {
|
|
login["scheme"] = "token"
|
|
login["secret"] = account.Token
|
|
} else {
|
|
login["scheme"] = "basic"
|
|
login["secret"] = base64.StdEncoding.EncodeToString([]byte(account.Username + ":" + password))
|
|
login["login_type"] = 0
|
|
login["tenant_code"] = cfg.Tenant
|
|
}
|
|
ctrl, err := client.command(ctx, "login", map[string]any{"login": login})
|
|
if err != nil && account.Token != "" && password != "" {
|
|
login["id"] = client.nextID("login-basic")
|
|
login["scheme"] = "basic"
|
|
login["secret"] = base64.StdEncoding.EncodeToString([]byte(account.Username + ":" + password))
|
|
login["login_type"] = 0
|
|
login["tenant_code"] = cfg.Tenant
|
|
ctrl, err = client.command(ctx, "login", map[string]any{"login": login})
|
|
}
|
|
if err != nil {
|
|
client.Close()
|
|
return nil, fmt.Errorf("login %s: %w", account.Username, err)
|
|
}
|
|
client.applyLogin(ctrl)
|
|
return client, nil
|
|
}
|
|
|
|
func connectForPrepare(ctx context.Context, cfg *Config, account AccountRecord, api *apiClient, password, verificationCode string) (*WSClient, bool, error) {
|
|
timeout, err := configDuration(cfg.CommandTimeout)
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
endpoint, err := websocketEndpoint(account.Endpoint, api.apiKey)
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
dialer := *websocket.DefaultDialer
|
|
if cfg.InsecureTLS {
|
|
dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} // #nosec G402: explicit load-test option.
|
|
}
|
|
conn, _, err := dialer.DialContext(ctx, endpoint, http.Header{})
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
client := &WSClient{account: account, conn: conn, commandTimeout: timeout, waits: make(map[string]chan ctrlPacket), done: make(chan struct{})}
|
|
go client.readLoop()
|
|
if _, err = client.command(ctx, "hi", map[string]any{"hi": map[string]any{
|
|
"id": client.nextID("hi"), "ver": "0.25", "ua": "im-pressure-prepare/1.0", "tenant": cfg.Tenant,
|
|
}}); err != nil {
|
|
client.Close()
|
|
return nil, false, err
|
|
}
|
|
secret := base64.StdEncoding.EncodeToString([]byte(account.Username + ":" + password))
|
|
ctrl, loginErr := client.command(ctx, "login", map[string]any{"login": map[string]any{
|
|
"id": client.nextID("login"), "scheme": "basic", "secret": secret,
|
|
"login_type": 0, "tenant_code": cfg.Tenant,
|
|
}})
|
|
created := false
|
|
if loginErr != nil {
|
|
var ce *ctrlError
|
|
if !errors.As(loginErr, &ce) || ce.Code != http.StatusUnauthorized {
|
|
client.Close()
|
|
return nil, false, loginErr
|
|
}
|
|
credential, credentialErr := api.createRegistrationCredential(ctx, account, verificationCode)
|
|
if credentialErr != nil {
|
|
client.Close()
|
|
return nil, false, credentialErr
|
|
}
|
|
acc := map[string]any{
|
|
"id": client.nextID("acc"), "user": "new", "scheme": "basic", "username": account.Username,
|
|
"registration_type": 0, "tenant_code": cfg.Tenant, "secret": secret, "login": true,
|
|
"desc": map[string]any{"public": map[string]any{"fn": strings.ToUpper(account.Username)}},
|
|
}
|
|
if credential != nil {
|
|
acc["phone"] = account.Phone
|
|
acc["cred"] = []map[string]any{{"meth": "tel", "val": account.Phone,
|
|
"resp": credential.VerificationToken, "params": map[string]any{"challenge_id": credential.ChallengeID}}}
|
|
}
|
|
ctrl, err = client.command(ctx, "acc", map[string]any{"acc": acc})
|
|
if err != nil {
|
|
client.Close()
|
|
return nil, false, err
|
|
}
|
|
created = true
|
|
}
|
|
client.applyLogin(ctrl)
|
|
return client, created, nil
|
|
}
|
|
|
|
func websocketEndpoint(raw, apiKey string) (string, error) {
|
|
u, err := url.Parse(strings.TrimSpace(raw))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if u.Scheme != "ws" && u.Scheme != "wss" {
|
|
return "", fmt.Errorf("unsupported websocket scheme %q", u.Scheme)
|
|
}
|
|
q := u.Query()
|
|
if q.Get("apikey") == "" {
|
|
q.Set("apikey", apiKey)
|
|
}
|
|
u.RawQuery = q.Encode()
|
|
return u.String(), nil
|
|
}
|
|
|
|
func (c *WSClient) applyLogin(ctrl ctrlPacket) {
|
|
if raw := ctrl.Params["user"]; len(raw) > 0 {
|
|
_ = json.Unmarshal(raw, &c.account.UserID)
|
|
}
|
|
if raw := ctrl.Params["token"]; len(raw) > 0 {
|
|
_ = json.Unmarshal(raw, &c.account.Token)
|
|
}
|
|
if raw := ctrl.Params["expires"]; len(raw) > 0 {
|
|
var value string
|
|
if json.Unmarshal(raw, &value) == nil {
|
|
c.account.TokenExpires, _ = time.Parse(time.RFC3339Nano, value)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *WSClient) Account() AccountRecord { return c.account }
|
|
|
|
func (c *WSClient) Subscribe(ctx context.Context, topic string) error {
|
|
if topic == "" {
|
|
return errors.New("empty topic")
|
|
}
|
|
id := c.nextID("sub")
|
|
_, err := c.command(ctx, "sub", map[string]any{"sub": map[string]any{
|
|
"id": id, "topic": topic, "get": map[string]any{"what": "desc"},
|
|
}})
|
|
return err
|
|
}
|
|
|
|
func (c *WSClient) GetData(ctx context.Context, topic string, since, before int64, limit int) error {
|
|
if topic == "" || since <= 0 || before <= since || limit <= 0 {
|
|
return errors.New("invalid get data range")
|
|
}
|
|
id := c.nextID("get-data")
|
|
_, err := c.command(ctx, "get", map[string]any{"get": map[string]any{
|
|
"id": id, "topic": topic, "what": "data",
|
|
"data": map[string]any{"since": since, "before": before, "limit": limit},
|
|
}})
|
|
return err
|
|
}
|
|
|
|
func (c *WSClient) Publish(ctx context.Context, topic, clientMid, runID, scenario string, messageIndex int) (AcceptedMessage, error) {
|
|
id := c.nextID("pub")
|
|
sentAt := time.Now().UTC()
|
|
ctrl, err := c.command(ctx, "pub", map[string]any{"pub": map[string]any{
|
|
"id": id, "topic": topic, "noecho": false,
|
|
"head": map[string]any{"mime": "application/json", "x-im-schema": 1, "x-im-type": "text", "x-im-client-mid": clientMid},
|
|
"content": map[string]any{"text": "loadtest", "run_id": runID, "scenario": scenario,
|
|
"sender_index": c.account.Index, "message_index": messageIndex, "sent_at_unix_nano": sentAt.UnixNano()},
|
|
}})
|
|
if err != nil {
|
|
return AcceptedMessage{}, err
|
|
}
|
|
seq, err := ctrlInt64(ctrl.Params["seq"])
|
|
if err != nil || seq <= 0 {
|
|
return AcceptedMessage{}, fmt.Errorf("publish ACK has invalid seq: %w", err)
|
|
}
|
|
kind := "p2p"
|
|
if strings.HasPrefix(topic, "grp") {
|
|
kind = "group"
|
|
}
|
|
return AcceptedMessage{RunID: runID, Scenario: scenario, Kind: kind, ClientMid: clientMid,
|
|
Topic: topic, Seq: seq, SenderMachine: c.account.Machine, SenderIndex: c.account.Index,
|
|
SentAt: sentAt, AckAt: time.Now().UTC()}, nil
|
|
}
|
|
|
|
func ctrlInt64(raw json.RawMessage) (int64, error) {
|
|
if len(raw) == 0 {
|
|
return 0, errors.New("missing value")
|
|
}
|
|
var number json.Number
|
|
if err := json.Unmarshal(raw, &number); err == nil {
|
|
return strconv.ParseInt(number.String(), 10, 64)
|
|
}
|
|
var text string
|
|
if err := json.Unmarshal(raw, &text); err != nil {
|
|
return 0, err
|
|
}
|
|
return strconv.ParseInt(text, 10, 64)
|
|
}
|
|
|
|
func (c *WSClient) command(ctx context.Context, _ string, message map[string]any) (ctrlPacket, error) {
|
|
id := nestedMessageID(message)
|
|
if id == "" {
|
|
return ctrlPacket{}, errors.New("command id is required")
|
|
}
|
|
wait := make(chan ctrlPacket, 1)
|
|
c.mu.Lock()
|
|
c.waits[id] = wait
|
|
c.mu.Unlock()
|
|
c.sendMu.Lock()
|
|
err := c.conn.WriteJSON(message)
|
|
c.sendMu.Unlock()
|
|
if err != nil {
|
|
c.removeWait(id)
|
|
return ctrlPacket{}, err
|
|
}
|
|
timer := time.NewTimer(c.commandTimeout)
|
|
defer timer.Stop()
|
|
select {
|
|
case ctrl := <-wait:
|
|
if ctrl.Code < 200 || ctrl.Code >= 300 {
|
|
return ctrl, &ctrlError{Code: ctrl.Code, Text: ctrl.Text}
|
|
}
|
|
return ctrl, nil
|
|
case <-timer.C:
|
|
c.removeWait(id)
|
|
return ctrlPacket{}, fmt.Errorf("ctrl timeout id=%s", id)
|
|
case <-ctx.Done():
|
|
c.removeWait(id)
|
|
return ctrlPacket{}, ctx.Err()
|
|
case <-c.done:
|
|
c.removeWait(id)
|
|
return ctrlPacket{}, errors.New("websocket closed")
|
|
}
|
|
}
|
|
|
|
func nestedMessageID(message map[string]any) string {
|
|
for _, value := range message {
|
|
if body, ok := value.(map[string]any); ok {
|
|
if id, ok := body["id"].(string); ok {
|
|
return id
|
|
}
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func (c *WSClient) nextID(prefix string) string {
|
|
n := atomic.AddUint64(&c.seq, 1)
|
|
return fmt.Sprintf("%s-%s-%04d-%d", prefix, c.account.Machine, c.account.Index, n)
|
|
}
|
|
|
|
func (c *WSClient) removeWait(id string) {
|
|
c.mu.Lock()
|
|
delete(c.waits, id)
|
|
c.mu.Unlock()
|
|
}
|
|
|
|
func (c *WSClient) readLoop() {
|
|
defer close(c.done)
|
|
for {
|
|
var envelope map[string]json.RawMessage
|
|
if err := c.conn.ReadJSON(&envelope); err != nil {
|
|
if !c.closed.Load() && c.onDisconnect != nil {
|
|
c.onDisconnect(c.account, err)
|
|
}
|
|
return
|
|
}
|
|
if raw := envelope["ctrl"]; len(raw) > 0 {
|
|
var ctrl ctrlPacket
|
|
if json.Unmarshal(raw, &ctrl) == nil && ctrl.ID != "" {
|
|
c.mu.Lock()
|
|
wait := c.waits[ctrl.ID]
|
|
delete(c.waits, ctrl.ID)
|
|
c.mu.Unlock()
|
|
if wait != nil {
|
|
wait <- ctrl
|
|
}
|
|
}
|
|
continue
|
|
}
|
|
if raw := envelope["data"]; len(raw) > 0 && c.onData != nil {
|
|
var data dataPacket
|
|
if json.Unmarshal(raw, &data) == nil {
|
|
c.onData(c.account, data)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *WSClient) Close() {
|
|
if c == nil || !c.closed.CompareAndSwap(false, true) {
|
|
return
|
|
}
|
|
_ = c.conn.Close()
|
|
}
|