NixMsg/sdk/go/connect.go

package nixmsg

import (
	"context"
	"encoding/json"
	"time"
)

// Connect 连接服务器。credential 为密码或会话令牌。
func (c *Client) Connect(ctx context.Context, rawURL, endpointID string, credential Credential, opts Options) error {
	o := opts.withDefaults()
	if o.MaxReceiveBytes > 0 && o.MaxReceiveBytes < 1024 {
		return apiErr(CodeBadRequest, "max_receive_bytes 小于 1024")
	}

	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return apiErr(CodeClosed, "已关闭")
	}
	if c.transport != nil {
		c.mu.Unlock()
		return apiErr(CodeBadRequest, "已在连接中")
	}
	c.opts = o
	if c.store == nil || c.opts.DedupCapacity != c.store.cap {
		c.store = newLRU(o.DedupCapacity)
	}
	c.endpointID = endpointID
	c.url = rawURL
	c.stopReconnect = false
	c.lastStopCode = ""
	c.lastStopErr = nil
	c.handshook = false
	pass := credential.Password
	if credential.SessionToken != "" {
		pass = credential.SessionToken
		c.credKind = "token"
	} else {
		c.credKind = "password"
	}
	var tr transport
	if o.transport != nil {
		tr = o.transport
	} else {
		tr = newMQTTTransport()
	}
	c.transport = tr
	c.backoff = newReconnectBackoff()
	inner, cancel := context.WithCancel(context.Background())
	c.ctx = inner
	c.cancel = cancel
	c.incoming = make(chan queuedFrame, eventQueueCap)
	c.cbQ = make(chan func(), 256)
	c.setStateLocked(StateConnecting, "")
	c.mu.Unlock()

	go c.downLoop(inner)
	go c.cbLoop(inner)

	tr.SetCredential(pass)
	cfg := transportConfig{
		URL:            rawURL,
		EndpointID:     endpointID,
		ConnectTimeout: o.ConnectTimeout,
		AllowTCP:       o.AllowTCP,
		Backoff:        c.backoff,
		OnDown:         c.handleDown,
		OnOffline:      c.onTransportOffline,
		OnAuthFailed:   func(reason AuthReason) { c.failAuth(reason) },
		OnKicked:       func() { c.failKicked() },
		MQTTReady:      func(readyCtx context.Context) error { return c.doHello(readyCtx) },
	}
	if err := tr.Start(inner, cfg); err != nil {
		c.teardown(false, apiErr(CodeNotConnected, err.Error()), false)
		return err
	}

	deadline := time.Now().Add(o.ConnectTimeout)
	for time.Now().Before(deadline) {
		c.mu.Lock()
		ok := c.handshook
		failed := c.stopReconnect
		st := c.state
		code := c.lastStopCode
		c.mu.Unlock()
		if ok {
			return nil
		}
		if failed || st == StateAuthFailed || st == StateKicked {
			if code == "" {
				code = string(st)
			}
			return apiErr(code, "认证失败，停止重连")
		}
		select {
		case <-ctx.Done():
			c.teardown(false, apiErr(CodeNotConnected, "连接已取消"), true)
			return ctx.Err()
		case <-time.After(20 * time.Millisecond):
		}
	}
	c.teardown(false, apiErr(CodeNotConnected, "连接超时"), true)
	return apiErr(CodeNotConnected, "连接超时")
}

func (c *Client) failAuth(reason AuthReason) {
	code := string(reason)
	if code == "" {
		code = CodeBadCredentials
	}
	err := apiErr(code, "认证失败，停止重连")
	c.mu.Lock()
	if c.stopReconnect && c.lastStopCode != "" {
		c.mu.Unlock()
		return
	}
	c.stopReconnect = true
	c.handshook = false
	c.lastStopCode = code
	c.lastStopErr = err
	c.setStateLocked(StateAuthFailed, code)
	c.failQueuedLocked(err)
	c.failPendingLocked(err, true)
	c.closeConnLostLocked()
	tr := c.transport
	cancel := c.cancel
	c.transport = nil
	c.mu.Unlock()
	if cancel != nil {
		cancel()
	}
	if tr != nil {
		go func() {
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			_ = tr.Stop(ctx)
		}()
	}
}

func (c *Client) failKicked() {
	err := apiErr(CodeTakenOver, "被顶号，停止重连")
	c.mu.Lock()
	if c.stopReconnect && c.lastStopCode == CodeTakenOver {
		c.mu.Unlock()
		return
	}
	c.stopReconnect = true
	c.handshook = false
	c.lastStopCode = CodeTakenOver
	c.lastStopErr = err
	c.setStateLocked(StateKicked, CodeTakenOver)
	c.failQueuedLocked(err)
	c.failPendingLocked(err, true)
	c.closeConnLostLocked()
	tr := c.transport
	cancel := c.cancel
	c.transport = nil
	c.mu.Unlock()
	if cancel != nil {
		cancel()
	}
	if tr != nil {
		go func() {
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			_ = tr.Stop(ctx)
		}()
	}
}

func (c *Client) handleFatal(reason string) {
	if reason == "" {
		reason = "fatal"
	}
	err := apiErr(reason, "致命错误，停止重连")
	c.mu.Lock()
	if c.stopReconnect && c.lastStopCode != "" {
		c.mu.Unlock()
		return
	}
	c.stopReconnect = true
	c.handshook = false
	c.lastStopCode = reason
	c.lastStopErr = err
	c.setStateLocked(StateAuthFailed, reason)
	c.failQueuedLocked(err)
	c.failPendingLocked(err, true)
	c.closeConnLostLocked()
	tr := c.transport
	cancel := c.cancel
	c.transport = nil
	c.mu.Unlock()
	if cancel != nil {
		cancel()
	}
	if tr != nil {
		go func() {
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			_ = tr.Stop(ctx)
		}()
	}
}

func (c *Client) setStateLocked(st ConnectionState, reason string) {
	c.state = st
	h := c.onConnection
	ev := ConnectionEvent{State: st, Reason: reason}
	c.dispatchLocked(func() {
		if h != nil {
			h(ev)
		}
	})
}

func (c *Client) doHello(ctx context.Context) error {
	c.mu.Lock()
	c.helloSentAt = time.Now()
	sentAt := c.helloSentAt
	label := c.opts.ClientLabel
	maxRecv := c.opts.MaxReceiveBytes
	c.connLost = make(chan struct{})
	c.mu.Unlock()

	req := map[string]any{
		"v":      1,
		"type":   "hello",
		"rid":    c.nextRID(),
		"client": label,
	}
	if maxRecv > 0 {
		req["max_receive_bytes"] = maxRecv
	}
	data, err := c.request(ctx, req, true)
	if err != nil {
		return err
	}
	recvAt := time.Now()
	var hd struct {
		ServerTimeMs       int64  `json:"server_time_ms"`
		ServerVersion      string `json:"server_version"`
		MaxBodyBytes       int    `json:"max_body_bytes"`
		MaxMetaBytes       int    `json:"max_meta_bytes"`
		MaxFrameBytes      int    `json:"max_frame_bytes"`
		MaxTTLSeconds      int64  `json:"max_ttl_seconds"`
		MaxScheduleSeconds int64  `json:"max_schedule_seconds"`
		AckTimeoutSeconds  int64  `json:"ack_timeout_seconds"`
		SessionToken       string `json:"session_token"`
	}
	if err := json.Unmarshal(data, &hd); err != nil {
		return err
	}
	mid := (sentAt.UnixMilli() + recvAt.UnixMilli()) / 2
	skew := hd.ServerTimeMs - mid

	c.mu.Lock()
	c.limits = HandshakeLimits{
		ServerTimeMs:       hd.ServerTimeMs,
		ServerVersion:      hd.ServerVersion,
		MaxBodyBytes:       hd.MaxBodyBytes,
		MaxMetaBytes:       hd.MaxMetaBytes,
		MaxFrameBytes:      hd.MaxFrameBytes,
		MaxTTLSeconds:      hd.MaxTTLSeconds,
		MaxScheduleSeconds: hd.MaxScheduleSeconds,
		AckTimeoutSeconds:  hd.AckTimeoutSeconds,
	}
	c.clockSkew = skew
	c.handshook = true
	if c.backoff != nil {
		c.backoff.MarkOnline()
	}
	c.setStateLocked(StateOnline, "")
	token := hd.SessionToken
	if token != "" {
		c.session = token
		if c.transport != nil {
			c.transport.SetCredential(token)
		}
		c.credKind = "token"
	}
	watchSet := c.watchSet
	watchIDs := append([]string(nil), c.watchIDs...)
	watchAll := c.watchAll
	c.mu.Unlock()

	if token != "" && c.onSession != nil {
		c.onSession(token)
	}
	if watchSet {
		go func() {
			_ = c.WatchPresence(context.Background(), watchIDs, watchAll)
		}()
	}
	c.drainSendQueue()
	return nil
}

// ClockSkewMs 当前服务器时间偏差（毫秒）。
func (c *Client) ClockSkewMs() int64 {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.clockSkew
}

// Limits 握手上限。
func (c *Client) Limits() HandshakeLimits {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.limits
}

// Logout 作废会话并停止重连。请求失败仍返回给应用。
func (c *Client) Logout(ctx context.Context) error {
	req := map[string]any{"v": 1, "type": "self.logout", "rid": c.nextRID()}
	var err error
	_, err = c.request(ctx, req, false)
	stop := apiErr(CodeLoggedOut, "已退出")
	c.teardown(false, stop, true)
	return err
}

// Close 关闭连接并停止重连。
func (c *Client) Close() error {
	err := apiErr(CodeClosed, "已关闭")
	c.teardown(true, err, true)
	return nil
}

func (c *Client) teardown(setClosed bool, stopErr error, waitStop bool) {
	c.mu.Lock()
	if setClosed {
		c.closed = true
	}
	c.stopReconnect = true
	if stopErr != nil {
		c.lastStopErr = stopErr
		c.lastStopCode = errCode(stopErr)
	}
	c.failQueuedLocked(c.stopErrLocked())
	c.failPendingLocked(c.stopErrLocked(), true)
	c.closeConnLostLocked()
	c.handshook = false
	c.session = ""
	tr := c.transport
	c.transport = nil
	cancel := c.cancel
	c.setStateLocked(StateOffline, c.lastStopCode)
	c.mu.Unlock()
	if cancel != nil {
		cancel()
	}
	if tr == nil {
		return
	}
	fn := func() {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = tr.Stop(ctx)
	}
	if waitStop {
		fn()
		return
	}
	go fn()
}

func (c *Client) failQueuedLocked(err error) {
	for _, it := range c.sendQ {
		if it.result != nil {
			select {
			case it.result <- sendOutcome{err: err}:
			default:
			}
		}
	}
	c.sendQ = nil
	c.inflight = 0
}