Cyrene_v2/plugins/channels/kook/websocket.go

package kook

// This file wires the stateful KOOK gateway WebSocket protocol onto the shared
// wsutil client. KOOK is a signal-driven gateway: after connecting, the server
// sends s=1 (hello) or s=6 (resume ack); the client heartbeats every ~30s with
// s=2 and expects s=3; inbound events arrive as s=0; the server may force a
// restart with s=5 (reconnect). The reconnect loop, heartbeat scheduling, and
// read-loop skeleton are provided by wsutil; only the signaling lives here.

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/channel/adapter"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/channel/adapter/wsutil"
)

// KOOK gateway signal numbers.
const (
	signalEvent     = 0
	signalHello     = 1
	signalPing      = 2
	signalPong      = 3
	signalResume    = 4
	signalReconnect = 5
	signalResumeAck = 6
)

// kookSignal is the envelope of every gateway frame.
type kookSignal struct {
	S  int             `json:"s"`
	D  json.RawMessage `json:"d"`
	SN int64           `json:"sn"`
}

// helloData is the "d" payload of an s=1 hello frame.
type helloData struct {
	Code      int    `json:"code"`
	SessionID string `json:"session_id"`
}

// parseSignal decodes a gateway text frame.
func parseSignal(payload []byte) (kookSignal, error) {
	var sig kookSignal
	if err := json.Unmarshal(payload, &sig); err != nil {
		return kookSignal{}, fmt.Errorf("kook: decode signal: %w", err)
	}
	return sig, nil
}

// heartbeatFrame builds the s=2 ping payload carrying the current sn.
func heartbeatFrame(sn int64) ([]byte, error) {
	return json.Marshal(map[string]any{"s": signalPing, "sn": sn})
}

// runWS starts the gateway client goroutine. It is idempotent: a second call
// while one is running is a no-op. The client tears itself down when ctx ends.
func (a *Adapter) runWS(ctx context.Context) {
	a.wsMu.Lock()
	if a.wsc != nil {
		a.wsMu.Unlock()
		return
	}
	client := wsutil.New(
		a.gatewayURL,
		wsutil.Handler{
			OnConnect: a.onConnect,
			OnMessage: a.onMessage,
			Heartbeat: a.heartbeat,
		},
		wsutil.WithHeartbeat(heartbeatInterval),
		wsutil.WithPongTimeout(helloTimeout),
		wsutil.WithLogger(a.log),
	)
	a.wsc = client
	runCtx, cancel := context.WithCancel(ctx)
	client.BindCancel(cancel)
	a.wsMu.Unlock()

	_ = client.Run(runCtx)

	a.wsMu.Lock()
	a.wsc = nil
	a.wsMu.Unlock()
}

// gatewayURL returns the gateway address to dial. On a first connect (or after a
// session reset) it returns a fresh URL; on a drop with an established session it
// appends the resume parameters so the gateway can replay missed messages.
func (a *Adapter) gatewayURL(ctx context.Context) (string, error) {
	base, err := a.fetchGatewayURL(ctx)
	if err != nil {
		return "", err
	}
	if !a.canResume() {
		return base, nil
	}
	return wsutil.AppendQuery(base, map[string]string{
		"resume":     "1",
		"sn":         strconv.FormatInt(a.currentSN(), 10),
		"session_id": a.currentSession(),
	}), nil
}

// onConnect performs the application handshake. It accepts either an s=1 hello
// (fresh or restarted session) or an s=6 resume ack (reconnected with an intact
// session) as the first frame within helloTimeout.
func (a *Adapter) onConnect(ctx context.Context, c *wsutil.Conn) error {
	_ = c.SetReadDeadline(time.Now().Add(helloTimeout))
	payload, err := c.ReadMessage()
	if err != nil {
		return fmt.Errorf("kook: hello read: %w", err)
	}
	sig, err := parseSignal(payload)
	if err != nil {
		return err
	}
	switch sig.S {
	case signalHello:
		var hd helloData
		if err := json.Unmarshal(sig.D, &hd); err != nil {
			return fmt.Errorf("kook: hello decode: %w", err)
		}
		if hd.Code != 0 {
			// Rejected (e.g. invalid/expired token or expired session): drop the
			// whole session state so the next dial is fresh.
			a.resetSession()
			return fmt.Errorf("kook: hello rejected code=%d", hd.Code)
		}
		a.setSession(hd.SessionID, 0)
	case signalResumeAck:
		// Resume accepted; keep the existing session id and sn.
		a.markResumeOK()
	default:
		return fmt.Errorf("kook: expected hello or resume ack, got s=%d", sig.S)
	}
	_ = c.SetReadDeadline(time.Time{})
	return nil
}

// onMessage dispatches every inbound application frame. s=0 events are converted
// to UnifiedMessage and delivered; control signals update session state or force
// a reconnect.
func (a *Adapter) onMessage(ctx context.Context, payload []byte) error {
	sig, err := parseSignal(payload)
	if err != nil {
		return err
	}
	switch sig.S {
	case signalEvent:
		if err := a.handleEvent(ctx, sig.D, payload); err != nil {
			return err
		}
		a.advanceSN(sig.SN)
		return nil
	case signalPong, signalResumeAck:
		return nil
	case signalReconnect:
		// Server invalidated the session: clear sn/session and re-dial fresh.
		a.resetSession()
		return errors.New("kook: gateway session invalid, reconnecting")
	default:
		return nil
	}
}

// heartbeat sends the s=2 ping carrying the current sn.
func (a *Adapter) heartbeat(ctx context.Context, c *wsutil.Conn) error {
	data, err := heartbeatFrame(a.currentSN())
	if err != nil {
		return err
	}
	return c.WriteMessage(data)
}

// handleEvent converts a single s=0 event into a UnifiedMessage and dispatches
// it. Malformed or non-chat events are logged and skipped, never fatal.
func (a *Adapter) handleEvent(ctx context.Context, data json.RawMessage, frame []byte) error {
	msg, ok := a.toUnifiedMessage(data, frame)
	if !ok {
		return nil
	}
	if err := msg.Validate(); err != nil {
		a.log.Warn("invalid message dropped", "err", err)
		return nil
	}
	a.dispatch(ctx, msg)
	return nil
}

// dispatch routes an inbound message to the push handler, falling back to the
// ReadMessages sink.
func (a *Adapter) dispatch(ctx context.Context, msg adapter.UnifiedMessage) {
	a.handlerMu.RLock()
	fn := a.handler
	a.handlerMu.RUnlock()
	if fn != nil {
		fn(ctx, msg)
		return
	}
	a.emitMu.RLock()
	emit := a.emit
	a.emitMu.RUnlock()
	if emit != nil {
		emit(msg)
	}
}

// setSession records a fresh session (new id, sn reset) and marks it resumable.
func (a *Adapter) setSession(sessionID string, sn int64) {
	a.sessMu.Lock()
	defer a.sessMu.Unlock()
	a.sessionID = sessionID
	a.lastSN = sn
	a.resumeOK = sessionID != ""
}

// markResumeOK marks the current session as resumable after a resume ack.
func (a *Adapter) markResumeOK() {
	a.sessMu.Lock()
	defer a.sessMu.Unlock()
	a.resumeOK = a.sessionID != ""
}

// resetSession clears all session state so the next connect is a fresh one.
func (a *Adapter) resetSession() {
	a.sessMu.Lock()
	defer a.sessMu.Unlock()
	a.sessionID = ""
	a.lastSN = 0
	a.resumeOK = false
}

// advanceSN raises the tracked sn when a newer event is seen.
func (a *Adapter) advanceSN(sn int64) {
	a.sessMu.Lock()
	defer a.sessMu.Unlock()
	if sn > a.lastSN {
		a.lastSN = sn
	}
}

func (a *Adapter) currentSession() string {
	a.sessMu.RLock()
	defer a.sessMu.RUnlock()
	return a.sessionID
}

func (a *Adapter) currentSN() int64 {
	a.sessMu.RLock()
	defer a.sessMu.RUnlock()
	return a.lastSN
}

func (a *Adapter) canResume() bool {
	a.sessMu.RLock()
	defer a.sessMu.RUnlock()
	return a.sessionID != "" && a.resumeOK
}