snowluma: SnowLuma (OneBot 11) Go SDK, drop-in compatible with napcat

connect_ws.go

@@ -0,0 +1,240 @@
+package snowluma
+
+import (
+	"context"
+	"errors"
+	"net"
+	"net/http"
+	neturl "net/url"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+// Session 是一条已建立的 OneBot WebSocket 会话，同时实现了 Caller。
+// 事件经 evCh 交给独立的分发协程串行处理（保序），读循环因此永远不会
+// 被 handler 里的 API 调用阻塞（否则 API 响应帧将无人读取而死锁）。
+type Session struct {
+	bot       *Client
+	conn      *websocket.Conn
+	caller    *wsCaller
+	selfID    atomic.Int64
+	closeOnce sync.Once
+	done      chan struct{}
+	evCh      chan sessionEvent
+}
+
+// sessionEvent 是会话分发给 worker 的一项。
+type sessionEvent struct {
+	ev    *Event
+	quick *quickBox
+}
+
+func newSession(bot *Client, conn *websocket.Conn) *Session {
+	s := &Session{bot: bot, conn: conn, done: make(chan struct{}), evCh: make(chan sessionEvent, 256)}
+	s.caller = newWSCaller(conn, s.done)
+	return s
+}
+
+// Call 实现 Caller：在当前连接上执行 action。
+func (s *Session) Call(ctx context.Context, action string, params any, out any) error {
+	return s.caller.Call(ctx, action, params, out)
+}
+
+// SelfID 返回该会话绑定的机器人 QQ 号（收到事件/生命周期后有效）。
+func (s *Session) SelfID() int64 { return s.selfID.Load() }
+
+// Close 主动断开会话。
+func (s *Session) Close() error {
+	var err error
+	s.closeOnce.Do(func() {
+		close(s.done)
+		err = s.conn.Close()
+	})
+	return err
+}
+
+// serve 阻塞读取连接上的数据，直到连接关闭。
+func (s *Session) serve() error {
+	defer s.Close()
+	go s.dispatchWorker()
+	for {
+		_, data, err := s.conn.ReadMessage()
+		if err != nil {
+			return err
+		}
+		s.bot.handleRaw(s, data)
+	}
+}
+
+// dispatchWorker 串行分发本会话的事件，直到会话关闭。
+func (s *Session) dispatchWorker() {
+	for {
+		select {
+		case <-s.done:
+			return
+		case job := <-s.evCh:
+			s.bot.dispatch(job.ev, s, job.quick)
+		}
+	}
+}
+
+// enqueue 把事件放入分发队列；队列满时丢弃并告警，避免阻塞读循环。
+func (s *Session) enqueue(ev *Event, quick *quickBox) {
+	select {
+	case s.evCh <- sessionEvent{ev: ev, quick: quick}:
+	case <-s.done:
+	default:
+		s.bot.log.Warn("snowluma: event queue full, event dropped", "post_type", ev.PostType, "self_id", ev.SelfID)
+	}
+}
+
+// ForwardWS 正向 WebSocket 连接器：框架作为客户端连到 SnowLuma 的 WS 服务。
+type ForwardWS struct {
+	url string
+}
+
+// NewForwardWS 创建正向 WS 连接器，url 形如 ws://127.0.0.1:3001。
+func NewForwardWS(url string) *ForwardWS { return &ForwardWS{url: url} }
+
+func (f *ForwardWS) Name() string { return "forward-ws(" + f.url + ")" }
+
+func (f *ForwardWS) Serve(ctx context.Context, bot *Client) error {
+	backoff := time.Second
+	for {
+		if ctx.Err() != nil {
+			return nil
+		}
+		conn, err := f.dial(bot)
+		if err != nil {
+			bot.log.Warn("forward ws dial failed", "err", err, "retry_in", backoff.String())
+			if !sleepBackoff(ctx, backoff) {
+				return nil
+			}
+			backoff = nextBackoff(backoff)
+			continue
+		}
+		backoff = time.Second
+		sess := newSession(bot, conn)
+		bot.bindSession(sess)
+		bot.log.Info("forward ws connected", "url", f.url)
+		err = sess.serve()
+		bot.unbindSession(sess)
+		bot.log.Warn("forward ws disconnected", "err", err)
+		// 连接建立后立刻被断开时同样退避，避免在服务端立刻拒绝的场景下
+		// 形成 dial→断开→dial 的紧循环（例如 access_token 错误）。
+		if !sleepBackoff(ctx, backoff) {
+			return nil
+		}
+		backoff = nextBackoff(backoff)
+	}
+}
+
+// nextBackoff 按 1s 起步、30s 封顶的指数退避递增。
+func nextBackoff(b time.Duration) time.Duration {
+	if b < 30*time.Second {
+		return b * 2
+	}
+	return 30 * time.Second
+}
+
+// sleepBackoff 等待退避时长或在 ctx 取消时提前返回 false。
+func sleepBackoff(ctx context.Context, d time.Duration) bool {
+	select {
+	case <-ctx.Done():
+		return false
+	case <-time.After(d):
+		return true
+	}
+}
+
+func (f *ForwardWS) dial(bot *Client) (*websocket.Conn, error) {
+	u, err := neturl.Parse(f.url)
+	if err != nil {
+		return nil, err
+	}
+	if bot.token != "" {
+		q := u.Query()
+		q.Set("access_token", bot.token)
+		u.RawQuery = q.Encode()
+	}
+	header := http.Header{}
+	if bot.token != "" {
+		header.Set("Authorization", "Bearer "+bot.token)
+	}
+	dialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second}
+	conn, _, err := dialer.Dial(u.String(), header)
+	return conn, err
+}
+
+// ReverseWS 反向 WebSocket 连接器：框架起一个 WS 服务，由 SnowLuma 主动连入。
+// 支持多个 SnowLuma 实例同时连入（每个连接一个独立会话）。
+type ReverseWS struct {
+	addr string
+	srv  *http.Server
+}
+
+// NewReverseWS 创建反向 WS 连接器，监听 addr（如 ":8080"），
+// SnowLuma 侧反向 WS 地址填 ws://<本机IP>:8080/（任意路径均可）。
+func NewReverseWS(addr string) *ReverseWS { return &ReverseWS{addr: addr} }
+
+func (r *ReverseWS) Name() string { return "reverse-ws(" + r.addr + ")" }
+
+func (r *ReverseWS) Serve(ctx context.Context, bot *Client) error {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
+		if !verifyToken(bot, req) {
+			http.Error(w, "invalid access_token", http.StatusUnauthorized)
+			return
+		}
+		up := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
+		conn, err := up.Upgrade(w, req, nil)
+		if err != nil {
+			return
+		}
+		sess := newSession(bot, conn)
+		bot.bindSession(sess)
+		bot.log.Info("reverse ws session connected", "remote", req.RemoteAddr)
+		defer bot.unbindSession(sess)
+		_ = sess.serve()
+	})
+	srv := &http.Server{
+		Addr:        r.addr,
+		Handler:     mux,
+		BaseContext: func(net.Listener) context.Context { return ctx },
+	}
+	r.srv = srv
+	err := srv.ListenAndServe()
+	if errors.Is(err, http.ErrServerClosed) {
+		return nil
+	}
+	return err
+}
+
+// Shutdown 关闭监听（Run 的退出流程会调用），已建立的会话由 Client 统一关闭。
+func (r *ReverseWS) Shutdown(ctx context.Context) error {
+	if r.srv == nil {
+		return nil
+	}
+	return r.srv.Shutdown(ctx)
+}
+
+// verifyToken 校验请求携带的 access_token（查询参数或 Authorization 头）。
+func verifyToken(bot *Client, r *http.Request) bool {
+	if bot.token == "" {
+		return true
+	}
+	if r.URL.Query().Get("access_token") == bot.token {
+		return true
+	}
+	auth := r.Header.Get("Authorization")
+	for _, prefix := range []string{"Bearer ", "Token "} {
+		if after, ok := strings.CutPrefix(auth, prefix); ok && after == bot.token {
+			return true
+		}
+	}
+	return false
+}