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 }