From b4ddd3c06824db196ee2d742d66b6ceea287d659 Mon Sep 17 00:00:00 2001
From: layue13 <layue13@icloud.com>
Date: Fri, 25 Sep 2026 18:21:35 +0800
Subject: [PATCH] Close sessions accepted during listener shutdown

---
 .../core/session/ProxySessionListener.java    | 15 ++++++-
 .../session/ProxySessionListenerTest.java     | 45 +++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/proxy-core/src/main/java/dev/strataproxy/core/session/ProxySessionListener.java b/proxy-core/src/main/java/dev/strataproxy/core/session/ProxySessionListener.java
index 61fdab2..a5a32c3 100644
--- a/proxy-core/src/main/java/dev/strataproxy/core/session/ProxySessionListener.java
+++ b/proxy-core/src/main/java/dev/strataproxy/core/session/ProxySessionListener.java
@@ -160,8 +160,8 @@ public final class ProxySessionListener implements Players {
                 .childOption(io.netty.channel.ChannelOption.TCP_NODELAY, true)
                 .childHandler(new ChannelInitializer<SocketChannel>() {
                     @Override protected void initChannel(SocketChannel channel) {
-                        Session session = new Session(ProxySessionListener.this, channel);
-                        allSessions.add(session);
+                        Session session = registerSession(channel);
+                        if (session == null) return;
                         channel.pipeline().addLast("minecraft-frame-decoder", new dev.strataproxy.core.protocol.MinecraftFrameDecoder(
                                 dev.strataproxy.core.protocol.ProtocolProfile.minecraft1710(), true,
                                 dev.strataproxy.core.protocol.ProtocolProfile.MAX_LOGIN_FRAME_BYTES));
@@ -188,6 +188,17 @@ public final class ProxySessionListener implements Players {
         return result;
     }
 
+    /** Serializes admission with close() so every accepted session is either closed or rejected. */
+    private synchronized Session registerSession(SocketChannel channel) {
+        if (closed) {
+            channel.close();
+            return null;
+        }
+        Session session = new Session(this, channel);
+        allSessions.add(session);
+        return session;
+    }
+
     @Override public Optional<PlayerView> find(PlayerIdentity identity) {
         Objects.requireNonNull(identity, "identity");
         Session session = sessionsByPlayerId.get(identity.playerId());
diff --git a/proxy-core/src/test/java/dev/strataproxy/core/session/ProxySessionListenerTest.java b/proxy-core/src/test/java/dev/strataproxy/core/session/ProxySessionListenerTest.java
index 15daef4..fbeb4a4 100644
--- a/proxy-core/src/test/java/dev/strataproxy/core/session/ProxySessionListenerTest.java
+++ b/proxy-core/src/test/java/dev/strataproxy/core/session/ProxySessionListenerTest.java
@@ -27,16 +27,20 @@ import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.ServerSocket;
 import java.net.Socket;
+import java.net.SocketException;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
 import java.security.KeyFactory;
 import java.security.spec.X509EncodedKeySpec;
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import javax.crypto.Cipher;
@@ -48,6 +52,47 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 final class ProxySessionListenerTest {
+    @Test
+    void shutdownWhileAcceptingClientsClosesEveryRegisteredSession() throws Exception {
+        var listener = new ProxySessionListener(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0),
+                new InMemoryBackendCatalog());
+        listener.setPlacement(player -> CompletableFuture.completedFuture(Optional.empty()));
+        var clients = Collections.synchronizedList(new ArrayList<Socket>());
+        var firstConnected = new CountDownLatch(1);
+        CompletableFuture<Void> connector = null;
+        try {
+            int port = ((InetSocketAddress) listener.start().toCompletableFuture()
+                    .get(5, TimeUnit.SECONDS).localAddress()).getPort();
+            connector = CompletableFuture.runAsync(() -> {
+                for (int i = 0; i < 64; i++) {
+                    try {
+                        Socket client = new Socket(InetAddress.getLoopbackAddress(), port);
+                        clients.add(client);
+                        firstConnected.countDown();
+                    } catch (IOException listenerClosed) {
+                        return;
+                    }
+                }
+            });
+            assertTrue(firstConnected.await(5, TimeUnit.SECONDS));
+            listener.close().toCompletableFuture().get(15, TimeUnit.SECONDS);
+            connector.get(5, TimeUnit.SECONDS);
+            assertTrue(listener.allSessions().isEmpty(), "shutdown left an accepted session registered");
+            for (Socket client : clients) {
+                client.setSoTimeout(2000);
+                try {
+                    assertEquals(-1, client.getInputStream().read());
+                } catch (SocketException reset) {
+                    // A peer close with unread login bytes can surface as a TCP reset on Windows.
+                }
+            }
+        } finally {
+            listener.close().toCompletableFuture().get(15, TimeUnit.SECONDS);
+            if (connector != null) connector.get(5, TimeUnit.SECONDS);
+            for (Socket client : clients) client.close();
+        }
+    }
+
     @Test
     void rejectsOversizedFrontendFrameBeforeAuthentication() throws Exception {
         var listener = new ProxySessionListener(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0),