package collector import ( "context" "encoding/json" "errors" "log/slog" "strings" "sync" "sync/atomic" "time" "github.com/gorilla/websocket" "github.com/sky22333/qqbot/config" "github.com/sky22333/qqbot/internal/qqapi" "github.com/sky22333/qqbot/internal/targets" "github.com/sky22333/qqbot/message" ) const ( intentPublicGuildMessages = 1 << 30 intentDirectMessage = 1 << 12 intentGroupAndC2C = 1 << 25 intentGuildMembers = 1 << 1 ) var intentLevels = []int{ intentPublicGuildMessages | intentDirectMessage | intentGroupAndC2C, intentPublicGuildMessages | intentGroupAndC2C, intentPublicGuildMessages | intentGuildMembers, } type wsPayload struct { Op int `json:"op"` D json.RawMessage `json:"d"` S *int64 `json:"s,omitempty"` T string `json:"t,omitempty"` } type helloData struct { HeartbeatInterval int64 `json:"heartbeat_interval"` } type c2cMessageEvent struct { ID string `json:"id"` Content string `json:"content"` Author struct { UserOpenID string `json:"user_openid"` } `json:"author"` } type groupMessageEvent struct { ID string `json:"id"` Content string `json:"content"` GroupOpenID string `json:"group_openid"` } type Collector struct { cfg config.Config logger *slog.Logger client *qqapi.Client store *targets.Store ctx context.Context cancel context.CancelFunc wg sync.WaitGroup } func New(cfg config.Config, logger *slog.Logger, store *targets.Store) (*Collector, error) { client, err := qqapi.NewClient(cfg) if err != nil { return nil, err } if logger == nil { logger = slog.Default() } ctx, cancel := context.WithCancel(context.Background()) return &Collector{ cfg: cfg, logger: logger, client: client, store: store, ctx: ctx, cancel: cancel, }, nil } func (c *Collector) Start() { c.wg.Add(1) go func() { defer c.wg.Done() c.runLoop() }() } func (c *Collector) Stop() { c.cancel() c.wg.Wait() } func (c *Collector) runLoop() { reconnectDelay, _ := c.cfg.CollectorReconnectDelay() if reconnectDelay <= 0 { reconnectDelay = 3 * time.Second } for { select { case <-c.ctx.Done(): return default: } if err := c.runOnce(); err != nil { if c.ctx.Err() != nil { return } c.logger.Error("collector 连接中断", "error", err) } timer := time.NewTimer(reconnectDelay) select { case <-c.ctx.Done(): timer.Stop() return case <-timer.C: } } } func (c *Collector) runOnce() error { token, err := c.client.AccessToken(c.ctx) if err != nil { return err } gatewayURL, err := c.client.GatewayURL(c.ctx) if err != nil { return err } conn, _, err := websocket.DefaultDialer.DialContext(c.ctx, gatewayURL, nil) if err != nil { return err } defer conn.Close() stopClose := make(chan struct{}) go func() { select { case <-c.ctx.Done(): _ = conn.Close() case <-stopClose: } }() defer close(stopClose) c.logger.Info("collector 已连接 gateway") var lastSeq atomic.Int64 intentIndex := 0 heartbeatStop := make(chan struct{}) defer close(heartbeatStop) for { select { case <-c.ctx.Done(): return nil default: } _, data, err := conn.ReadMessage() if err != nil { return err } payload := wsPayload{} if err := json.Unmarshal(data, &payload); err != nil { continue } if payload.S != nil { lastSeq.Store(*payload.S) } switch payload.Op { case 10: hello := helloData{} if err := json.Unmarshal(payload.D, &hello); err == nil && hello.HeartbeatInterval > 0 { go startHeartbeat(c.ctx, conn, hello.HeartbeatInterval, &lastSeq, heartbeatStop) } if err := sendIdentify(conn, token, intentLevels[intentIndex]); err != nil { return err } case 9: if intentIndex < len(intentLevels)-1 { intentIndex++ } if err := sendIdentify(conn, token, intentLevels[intentIndex]); err != nil { return err } case 0: c.handleDispatch(payload.T, payload.D) case 7: return errors.New("gateway 要求重连") } } } func (c *Collector) handleDispatch(eventType string, raw json.RawMessage) { switch strings.TrimSpace(eventType) { case "C2C_MESSAGE_CREATE": event := c2cMessageEvent{} if err := json.Unmarshal(raw, &event); err != nil { return } if event.Author.UserOpenID == "" { return } _ = c.store.Upsert(message.TargetC2C, event.Author.UserOpenID, event.ID, event.Content) c.logger.Info("采集到 c2c 目标", "target_id", event.Author.UserOpenID, "message_id", event.ID) case "GROUP_AT_MESSAGE_CREATE": event := groupMessageEvent{} if err := json.Unmarshal(raw, &event); err != nil { return } if event.GroupOpenID == "" { return } _ = c.store.Upsert(message.TargetGroup, event.GroupOpenID, event.ID, event.Content) c.logger.Info("采集到 group 目标", "target_id", event.GroupOpenID, "message_id", event.ID) } } func startHeartbeat(ctx context.Context, conn *websocket.Conn, intervalMS int64, seq *atomic.Int64, stop <-chan struct{}) { ticker := time.NewTicker(time.Duration(intervalMS) * time.Millisecond) defer ticker.Stop() for { select { case <-ctx.Done(): return case <-stop: return case <-ticker.C: payload := map[string]any{ "op": 1, "d": seq.Load(), } _ = conn.WriteJSON(payload) } } } func sendIdentify(conn *websocket.Conn, token string, intents int) error { payload := map[string]any{ "op": 2, "d": map[string]any{ "token": "QQBot " + token, "intents": intents, "shard": []int{0, 1}, }, } return conn.WriteJSON(payload) }