fix: queue config plugin messages during backend login

Author: robinbraemer

pkg/edition/java/proxy/session_backend_login.go

@@ -309,6 +309,12 @@ func (b *backendLoginSessionHandler) handleServerLoginSuccess() {
 		}
 
 		ash := player.ActiveSessionHandler()
+		if csh, ok := ash.(*clientConfigSessionHandler); ok {
+			if err = csh.flushQueuedPluginMessagesTo(b.serverConn); err != nil {
+				fail(err)
+				return
+			}
+		}
 		csh, ok := ash.(*clientPlaySessionHandler)
 		if ok {
 			serverMc.SetAutoReading(false)

pkg/edition/java/proxy/session_backend_play_test.go

@@ -55,6 +55,8 @@ func TestBackendPlayRegisterForwardsToPlayer(t *testing.T) {
 type testMinecraftConn struct {
 	writtenPackets []proto.Packet
 	connType       phase.ConnectionType
+	activeHandler  netmc.SessionHandler
+	writer         netmc.Writer
 }
 
 func (t *testMinecraftConn) Context() context.Context { return context.Background() }
@@ -70,16 +72,19 @@ func (t *testMinecraftConn) Type() phase.ConnectionType {
 	return phase.Vanilla
 }
 func (t *testMinecraftConn) SetType(ct phase.ConnectionType)            { t.connType = ct }
-func (t *testMinecraftConn) ActiveSessionHandler() netmc.SessionHandler { return nil }
-func (t *testMinecraftConn) SetActiveSessionHandler(*state.Registry, netmc.SessionHandler) {
+func (t *testMinecraftConn) ActiveSessionHandler() netmc.SessionHandler { return t.activeHandler }
+func (t *testMinecraftConn) SetActiveSessionHandler(_ *state.Registry, handler netmc.SessionHandler) {
+	t.activeHandler = handler
 }
 func (t *testMinecraftConn) SwitchSessionHandler(*state.Registry) bool { return true }
 func (t *testMinecraftConn) AddSessionHandler(*state.Registry, netmc.SessionHandler) {
 }
-func (t *testMinecraftConn) SetAutoReading(bool)               {}
-func (t *testMinecraftConn) SetProtocol(proto.Protocol)        {}
-func (t *testMinecraftConn) SetState(*state.Registry)          {}
-func (t *testMinecraftConn) SetOutboundState(*state.Registry)  {}
+func (t *testMinecraftConn) SetAutoReading(bool)        {}
+func (t *testMinecraftConn) SetProtocol(proto.Protocol) {}
+func (t *testMinecraftConn) SetState(*state.Registry)   {}
+func (t *testMinecraftConn) SetOutboundState(s *state.Registry) {
+	t.Writer().SetState(s)
+}
 func (t *testMinecraftConn) SetCompressionThreshold(int) error { return nil }
 func (t *testMinecraftConn) EnableEncryption([]byte) error     { return nil }
 func (t *testMinecraftConn) WritePacket(packet proto.Packet) error {
@@ -94,7 +99,25 @@ func (t *testMinecraftConn) BufferPacket(packet proto.Packet) error {
 func (t *testMinecraftConn) BufferPayload([]byte) error { return nil }
 func (t *testMinecraftConn) Flush() error               { return nil }
 func (t *testMinecraftConn) Reader() netmc.Reader       { return nil }
-func (t *testMinecraftConn) Writer() netmc.Writer       { return nil }
-func (t *testMinecraftConn) EnablePlayPacketQueue()     {}
+func (t *testMinecraftConn) Writer() netmc.Writer {
+	if t.writer == nil {
+		t.writer = &testWriter{}
+	}
+	return t.writer
+}
+func (t *testMinecraftConn) EnablePlayPacketQueue() {}
 
 var _ netmc.MinecraftConn = (*testMinecraftConn)(nil)
+
+type testWriter struct {
+	state *state.Registry
+}
+
+func (t *testWriter) WritePacket(proto.Packet) (int, error) { return 0, nil }
+func (t *testWriter) Write([]byte) (int, error)             { return 0, nil }
+func (t *testWriter) Flush() error                          { return nil }
+func (t *testWriter) SetProtocol(proto.Protocol)            {}
+func (t *testWriter) SetState(s *state.Registry)            { t.state = s }
+func (t *testWriter) SetCompressionThreshold(int) error     { return nil }
+func (t *testWriter) EnableEncryption([]byte) error         { return nil }
+func (t *testWriter) Direction() proto.Direction            { return proto.ClientBound }

pkg/edition/java/proxy/session_client_config.go

@@ -2,9 +2,13 @@ package proxy
 
 import (
 	"bytes"
+	"fmt"
+	"sync"
 
+	"github.com/gammazero/deque"
 	"github.com/go-logr/logr"
 	"github.com/robinbraemer/event"
+	"go.minekube.com/common/minecraft/component"
 	"go.minekube.com/gate/pkg/edition/java/netmc"
 	"go.minekube.com/gate/pkg/edition/java/proto/packet"
 	"go.minekube.com/gate/pkg/edition/java/proto/packet/config"
@@ -22,10 +26,18 @@ type clientConfigSessionHandler struct {
 	player *connectedPlayer
 	log    logr.Logger
 
-	brandChannel string
-
 	configSwitchDone future.Future[any]
 
+	mu struct {
+		sync.Mutex
+		pluginMessages           deque.Deque[*plugin.Message]
+		pluginMessagesBytes      int
+		pluginMessagesOverflowed bool
+		brandChannel             string
+		brandForwardedServer     *serverConnection
+		readyServer              *serverConnection
+	}
+
 	nopSessionHandler
 }
 
@@ -80,20 +92,13 @@ func (h *clientConfigSessionHandler) HandlePacket(pc *proto.PacketContext) {
 
 // handleBackendFinishUpdate handles the backend finishing the config stage.
 func (h *clientConfigSessionHandler) handleBackendFinishUpdate(serverConn *serverConnection, p *config.FinishedUpdate) *future.Future[any] {
-	smc, ok := serverConn.ensureConnected()
+	_, ok := serverConn.ensureConnected()
 	if !ok {
 		return nil
 	}
 	brand := serverConn.player.ClientBrand()
-	if brand == "" && h.brandChannel != "" {
-		buf := new(bytes.Buffer)
-		_ = util.WriteString(buf, brand)
-
-		brandPacket := &plugin.Message{
-			Channel: h.brandChannel,
-			Data:    buf.Bytes(),
-		}
-		_ = smc.WritePacket(brandPacket)
+	if brand != "" {
+		h.writeBrandPacketTo(serverConn, brand)
 	}
 
 	if err := h.player.WritePacket(p); err != nil {
@@ -115,23 +120,29 @@ func handleResourcePackResponse(p *packet.ResourcePackResponse, handler resource
 }
 
 func (h *clientConfigSessionHandler) handlePluginMessage(p *plugin.Message) {
-	serverConn := h.player.connectionInFlight()
-	if serverConn == nil {
-		return
-	}
-
 	if plugin.McBrand(p) {
 		brand := plugin.ReadBrandMessage(p.Data)
-		h.brandChannel = p.Channel
+		h.player.setClientBrand(brand)
+		readyServer := h.setBrandChannel(p.Channel)
 		h.event().FireParallel(&PlayerClientBrandEvent{
 			player: h.player,
 			brand:  brand,
 		})
 		// Client sends `minecraft:brand` packet immediately after Login,
 		// but at this time the backend server may not be ready
+		if readyServer != nil {
+			h.writeBrandPacketTo(readyServer, brand)
+		}
+		return
 	} else if bungeecord.IsBungeeCordMessage(p) {
 		return
+	} else if h.enqueuePluginMessage(h.player.connectionInFlightOrConnectedServer(), p) {
+		return
 	} else {
+		serverConn := h.player.connectionInFlightOrConnectedServer()
+		if serverConn == nil {
+			return
+		}
 		id, ok := h.player.proxy.ChannelRegistrar().FromID(p.Channel)
 		if !ok {
 			smc, ok := serverConn.ensureConnected()
@@ -164,8 +175,127 @@ func (h *clientConfigSessionHandler) handlePluginMessage(p *plugin.Message) {
 	}
 }
 
+func (h *clientConfigSessionHandler) setBrandChannel(channel string) *serverConnection {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	h.mu.brandChannel = channel
+	return h.mu.readyServer
+}
+
+func (h *clientConfigSessionHandler) brandPacket(brand string) *plugin.Message {
+	h.mu.Lock()
+	channel := h.mu.brandChannel
+	h.mu.Unlock()
+	if channel == "" {
+		return nil
+	}
+	buf := new(bytes.Buffer)
+	_ = util.WriteString(buf, brand)
+	return &plugin.Message{
+		Channel: channel,
+		Data:    buf.Bytes(),
+	}
+}
+
+func (h *clientConfigSessionHandler) writeBrandPacketTo(serverConn *serverConnection, brand string) {
+	h.mu.Lock()
+	if h.mu.brandForwardedServer == serverConn {
+		h.mu.Unlock()
+		return
+	}
+	h.mu.Unlock()
+
+	brandPacket := h.brandPacket(brand)
+	if brandPacket == nil {
+		return
+	}
+	smc, ok := serverConn.ensureConnected()
+	if !ok {
+		return
+	}
+	if err := smc.WritePacket(brandPacket); err != nil {
+		return
+	}
+
+	h.mu.Lock()
+	h.mu.brandForwardedServer = serverConn
+	h.mu.Unlock()
+}
+
+// enqueuePluginMessage returns true when the message was handled by the queue
+// path, including overflow rejection. It returns false only once the backend is
+// ready for direct config plugin messages.
+func (h *clientConfigSessionHandler) enqueuePluginMessage(target *serverConnection, msg *plugin.Message) bool {
+	h.mu.Lock()
+	if target != nil && h.mu.readyServer == target {
+		h.mu.Unlock()
+		return false
+	}
+	if h.mu.pluginMessagesOverflowed {
+		h.mu.Unlock()
+		return true
+	}
+	newBytes := h.mu.pluginMessagesBytes + len(msg.Data)
+	newCount := h.mu.pluginMessages.Len() + 1
+	if newBytes > maxQueuedLoginPluginMessageBytes || newCount > maxQueuedLoginPluginMessages {
+		h.mu.pluginMessagesOverflowed = true
+		h.mu.pluginMessages.Clear()
+		h.mu.pluginMessagesBytes = 0
+		h.mu.Unlock()
+		h.log.Info("disconnecting player: pre-backend config plugin message queue exceeded its limits",
+			"messages", newCount, "bytes", newBytes)
+		h.player.Disconnect(&component.Text{
+			Content: "Too many plugin messages were sent before joining a server",
+		})
+		return true
+	}
+	h.mu.pluginMessages.PushBack(&plugin.Message{
+		Channel: msg.Channel,
+		Data:    append([]byte(nil), msg.Data...),
+	})
+	h.mu.pluginMessagesBytes = newBytes
+	h.mu.Unlock()
+	return true
+}
+
+func (h *clientConfigSessionHandler) flushQueuedPluginMessagesTo(serverConn *serverConnection) error {
+	smc, ok := serverConn.ensureConnected()
+	if !ok {
+		return netmc.ErrClosedConn
+	}
+
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	if h.mu.readyServer == serverConn {
+		return nil
+	}
+
+	h.mu.pluginMessagesBytes = 0
+	n := h.mu.pluginMessages.Len()
+	msgs := make([]*plugin.Message, 0, n)
+	for h.mu.pluginMessages.Len() != 0 {
+		msgs = append(msgs, h.mu.pluginMessages.PopFront())
+	}
+	for _, pm := range msgs {
+		if err := smc.BufferPacket(pm); err != nil {
+			return fmt.Errorf("error buffering queued config plugin message %q: %w", pm.Channel, err)
+		}
+	}
+	if n != 0 {
+		if err := smc.Flush(); err != nil {
+			return err
+		}
+	}
+	h.mu.readyServer = serverConn
+	return nil
+}
+
 func (h *clientConfigSessionHandler) handleKnownPacks(p *config.KnownPacks, pc *proto.PacketContext) {
-	smc, ok := h.player.connectionInFlightOrConnectedServer().ensureConnected()
+	serverConn := h.player.connectionInFlightOrConnectedServer()
+	if serverConn == nil {
+		return
+	}
+	smc, ok := serverConn.ensureConnected()
 	if ok {
 		_ = smc.WritePacket(p)
 	}