snowluma/connect_ws.go

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
}