Selaa lähdekoodia

fix(discord): drop the gateway connection when heartbeats go unanswered (#6497)

The heartbeat goroutine wrote op 1 on its interval and ignored op 11, so a
connection that stopped being answered was never noticed. A half-open socket
is the case that matters: the kernel accepts the writes and the read loop
stays blocked, so the bot serves nothing for as long as the panel runs, and
nothing in the log says so. Discord asks clients to close and reconnect when
a heartbeat goes unacknowledged, which is what the ticker now does, letting
the existing reconnect loop take over.

The writeMu regression test's fake gateway answered no heartbeat at all,
which the new check reads as a dead socket; it now acknowledges them the way
Discord does and paces its op 1 flood, keeping its one-second window of
concurrent writes intact.
BlindMaster24 12 tuntia sitten
vanhempi
sitoutus
5c34baa8df

+ 12 - 1
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()

+ 64 - 0
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")
+	}
+}

+ 27 - 4
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: