diff --git a/internal/web/service/discord/gateway.go b/internal/web/service/discord/gateway.go index d115ba47b..3a9136142 100644 --- a/internal/web/service/discord/gateway.go +++ b/internal/web/service/discord/gateway.go @@ -11,6 +11,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/gorilla/websocket" @@ -278,6 +279,11 @@ func (g *GatewayClient) connectAndListen(ctx context.Context) error { hbStop := make(chan struct{}) defer close(hbStop) + // Discord answers every heartbeat with op 11; a half-open socket keeps taking + // writes and never answers, so a missing ACK means this one must be dropped. + var acked atomic.Bool + acked.Store(true) + go func() { interval := time.Duration(helloData.HeartbeatInterval) * time.Millisecond if interval <= 0 { @@ -293,6 +299,11 @@ func (g *GatewayClient) connectAndListen(ctx context.Context) error { case <-ctx.Done(): return case <-ticker.C: + if !acked.Swap(false) { + logger.Warning("Discord heartbeats went unanswered; dropping the zombied gateway connection") + _ = conn.Close() + return + } g.mu.Lock() seq := g.lastSeq c := g.conn @@ -334,7 +345,7 @@ func (g *GatewayClient) connectAndListen(ctx context.Context) error { switch payload.Op { case opHeartbeatACK: - // Heartbeat acknowledged + acked.Store(true) case opHeartbeat: // Discord requested immediate heartbeat g.mu.Lock() diff --git a/internal/web/service/discord/gateway_heartbeat_ack_test.go b/internal/web/service/discord/gateway_heartbeat_ack_test.go new file mode 100644 index 000000000..99a71c4ec --- /dev/null +++ b/internal/web/service/discord/gateway_heartbeat_ack_test.go @@ -0,0 +1,64 @@ +package discord + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/gorilla/websocket" +) + +func TestGatewayDropsConnectionWhenHeartbeatsGoUnanswered(t *testing.T) { + settingService := setupTestDB(t) + _ = settingService.SetDiscordBotEnable(true) + _ = settingService.SetDiscordBotToken("test-gw-token") + + upgrader := websocket.Upgrader{} + clientClosed := make(chan struct{}) + var once sync.Once + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + + // Hello with a short interval and not one op 11 after it: the socket stays + // open and reads fine, which is what a zombied connection looks like. + hello := GatewayPayload{Op: opHello} + hello.D, _ = json.Marshal(HelloData{HeartbeatInterval: 50}) + if err := conn.WriteJSON(hello); err != nil { + return + } + for { + var payload GatewayPayload + if err := conn.ReadJSON(&payload); err != nil { + once.Do(func() { close(clientClosed) }) + return + } + } + })) + defer wsServer.Close() + + discordSvc := NewDiscordService(settingService) + gw := NewGatewayClient(discordSvc, settingService, &mockServerProvider{}, &mockInboundProvider{}, &mockXrayRestart{}) + gw.SetGatewayURL("ws" + strings.TrimPrefix(wsServer.URL, "http")) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if err := gw.Start(ctx); err != nil { + t.Fatalf("gw.Start failed: %v", err) + } + defer gw.Stop() + + select { + case <-clientClosed: + case <-time.After(3 * time.Second): + t.Fatal("gateway held on to a connection whose heartbeats were never acknowledged") + } +} diff --git a/internal/web/service/discord/gateway_test.go b/internal/web/service/discord/gateway_test.go index c5340659c..3f5b52e35 100644 --- a/internal/web/service/discord/gateway_test.go +++ b/internal/web/service/discord/gateway_test.go @@ -302,22 +302,45 @@ func TestGatewayRequestedHeartbeatDoesNotRaceTicker(t *testing.T) { return } defer conn.Close() - _ = conn.WriteJSON(GatewayPayload{Op: opHello, D: []byte(`{"heartbeat_interval": 1}`)}) + // 10ms, not 1ms: Discord answers every heartbeat and the client now drops a + // socket it hears nothing back on, so the ACK needs room to arrive. + _ = conn.WriteJSON(GatewayPayload{Op: opHello, D: []byte(`{"heartbeat_interval": 10}`)}) + + // Two server goroutines write, so they share one writer: gorilla panics on + // concurrent writes, and this test is about the CLIENT's two writers. + var writeMu sync.Mutex + writeJSON := func(v any) error { + writeMu.Lock() + defer writeMu.Unlock() + return conn.WriteJSON(v) + } readErr := make(chan error, 1) go func() { for { - if _, _, err := conn.ReadMessage(); err != nil { + var payload GatewayPayload + if err := conn.ReadJSON(&payload); err != nil { readErr <- err return } + // Discord answers every heartbeat; without this the zombie check + // closes the socket a millisecond into the flood below. + if payload.Op == opHeartbeat { + if err := writeJSON(GatewayPayload{Op: opHeartbeatACK}); err != nil { + readErr <- err + return + } + } } }() - // Op 1 from the server makes the read loop write while the 1ms ticker writes too. + // Op 1 from the server makes the read loop write while the ticker writes too. for deadline := time.Now().Add(time.Second); time.Now().Before(deadline); { - if err := conn.WriteJSON(GatewayPayload{Op: opHeartbeat}); err != nil { + if err := writeJSON(GatewayPayload{Op: opHeartbeat}); err != nil { break } + // Leave the client room to drain the flood and answer: a saturated + // socket delays the ACK this test now depends on. + time.Sleep(time.Millisecond) } select { case err := <-readErr: