From fb467074a62c2528fdb1ba0f5a0610e90ddf3928 Mon Sep 17 00:00:00 2001
From: h3mant-1 <hemant.aggarwal@insticator.com>
Date: Tue, 23 Jul 2024 16:04:20 +0530
Subject: [PATCH] Re:#926, Fix #925 Memory Leak while
 WebSocketServerHandshakeException or Channel failed

---
 .../socketio/handler/ClientsBox.java          |  4 +-
 .../transport/WebSocketTransport.java         | 44 +++++++++++++------
 2 files changed, 32 insertions(+), 16 deletions(-)

diff --git a/src/main/java/com/corundumstudio/socketio/handler/ClientsBox.java b/src/main/java/com/corundumstudio/socketio/handler/ClientsBox.java
index 96d4d6ca..beafd2ce 100644
--- a/src/main/java/com/corundumstudio/socketio/handler/ClientsBox.java
+++ b/src/main/java/com/corundumstudio/socketio/handler/ClientsBox.java
@@ -42,8 +42,8 @@ public void addClient(ClientHead clientHead) {
         uuid2clients.put(clientHead.getSessionId(), clientHead);
     }
 
-    public void removeClient(UUID sessionId) {
-        uuid2clients.remove(sessionId);
+    public ClientHead removeClient(UUID sessionId) {
+        return uuid2clients.remove(sessionId);
     }
 
     public ClientHead get(UUID sessionId) {
diff --git a/src/main/java/com/corundumstudio/socketio/transport/WebSocketTransport.java b/src/main/java/com/corundumstudio/socketio/transport/WebSocketTransport.java
index d5a7c127..ca7cc0a1 100644
--- a/src/main/java/com/corundumstudio/socketio/transport/WebSocketTransport.java
+++ b/src/main/java/com/corundumstudio/socketio/transport/WebSocketTransport.java
@@ -153,31 +153,47 @@ private void handshake(ChannelHandlerContext ctx, final UUID sessionId, String p
                 new WebSocketServerHandshakerFactory(getWebSocketLocation(req), null, true, configuration.getMaxFramePayloadLength());
         WebSocketServerHandshaker handshaker = factory.newHandshaker(req);
         if (handshaker != null) {
-            ChannelFuture f = handshaker.handshake(channel, req);
-            f.addListener(new ChannelFutureListener() {
-                @Override
-                public void operationComplete(ChannelFuture future) throws Exception {
-                    if (!future.isSuccess()) {
-                        log.error("Can't handshake " + sessionId, future.cause());
-                        return;
+            try {
+                ChannelFuture f = handshaker.handshake(channel, req);
+                f.addListener(new ChannelFutureListener() {
+                    @Override
+                    public void operationComplete(ChannelFuture future) throws Exception {
+                        if (!future.isSuccess()) {
+                            log.error("Can't handshake {}", sessionId, future.cause());
+                            closeClient(sessionId, channel);
+                            return;
+                        }
+                        channel.pipeline().addBefore(SocketIOChannelInitializer.WEB_SOCKET_TRANSPORT, SocketIOChannelInitializer.WEB_SOCKET_AGGREGATOR,
+                                new WebSocketFrameAggregator(configuration.getMaxFramePayloadLength()));
+                        connectClient(channel, sessionId);
                     }
-
-                    channel.pipeline().addBefore(SocketIOChannelInitializer.WEB_SOCKET_TRANSPORT, SocketIOChannelInitializer.WEB_SOCKET_AGGREGATOR,
-                            new WebSocketFrameAggregator(configuration.getMaxFramePayloadLength()));
-                    connectClient(channel, sessionId);
-                }
-            });
+                });
+            } catch (Throwable e) {
+                log.warn("Can't handshake {}, {}", sessionId, e.getMessage(), e);
+                closeClient(sessionId, channel);
+            }
         } else {
             WebSocketServerHandshakerFactory.sendUnsupportedVersionResponse(ctx.channel());
         }
     }
 
+    private void closeClient(UUID sessionId, Channel channel) {
+        try {
+            channel.close();
+        } catch (Throwable t) {
+            log.warn("Can't close channel for sessionId: {}", sessionId, t);
+        }
+        ClientHead clientHead = clientsBox.removeClient(sessionId);
+        clientHead.disconnect();
+        log.info("Client with sessionId: {} was disconnected", sessionId);
+    }
+
     private void connectClient(final Channel channel, final UUID sessionId) {
         ClientHead client = clientsBox.get(sessionId);
         if (client == null) {
             log.warn("Unauthorized client with sessionId: {} with ip: {}. Channel closed!",
                         sessionId, channel.remoteAddress());
-            channel.close();
+            closeClient(sessionId, channel);
             return;
         }