Added cleanup

lolzballs · May 30, 2015 · ebf90ed · ebf90ed
1 parent 06e6261
commit ebf90ed
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 1 deletion.
diff --git a/Client/src/tk/jackyliao123/proxy/client/socks/Socks5ClientTunnelChannelWrapper.java b/Client/src/tk/jackyliao123/proxy/client/socks/Socks5ClientTunnelChannelWrapper.java
@@ -9,6 +9,7 @@
 import java.nio.channels.SelectionKey;
 import java.nio.channels.SocketChannel;
 import java.util.HashMap;
+import java.util.Map;
 
 public class Socks5ClientTunnelChannelWrapper extends TunnelChannelWrapper {
     public final HashMap<Integer, Socks5ConnectionData> connections;
@@ -39,4 +40,14 @@ public void startReading(int connectionId) {
         connection.client.addInterest(SelectionKey.OP_READ);
         connection.client.stopReading = false;
     }
+
+    @Override
+    public void cleanup() {
+        for (Map.Entry<Integer, Socks5ConnectionData> entry : connections.entrySet()) {
+            ChannelWrapper wrapper = entry.getValue().client;
+            wrapper.disconnectListener = null;
+            wrapper.close();
+        }
+        connections.clear();
+    }
 }
diff --git a/Common/src/tk/jackyliao123/proxy/TunnelChannelWrapper.java b/Common/src/tk/jackyliao123/proxy/TunnelChannelWrapper.java
@@ -84,4 +84,7 @@ public ByteBuffer popWriteBuffer() throws IOException {
     public abstract void stopReading(int connectionId);
 
     public abstract void startReading(int connectionId);
+
+    public abstract void cleanup();
+
 }
diff --git a/Common/src/tk/jackyliao123/proxy/event/EventProcessor.java b/Common/src/tk/jackyliao123/proxy/event/EventProcessor.java
@@ -138,7 +138,12 @@ public void process(long timeout) throws IOException {
             } catch (Throwable t) {
                 Logger.error("Critical Error, Throwable caught");
                 Logger.error(t);
-                kill((ChannelWrapper) key.attachment());
+                ChannelWrapper wrapper = (ChannelWrapper) key.attachment();
+                if (wrapper instanceof TunnelChannelWrapper) {
+                    ((TunnelChannelWrapper) wrapper).cleanup();
+                }
+                kill(wrapper);
+                System.gc();
             }
 
             keys.remove();

diff --git a/Server/src/tk/jackyliao123/proxy/server/ServerTunnelChannelWrapper.java b/Server/src/tk/jackyliao123/proxy/server/ServerTunnelChannelWrapper.java
@@ -7,6 +7,7 @@
 import java.nio.channels.SelectionKey;
 import java.nio.channels.SocketChannel;
 import java.util.HashMap;
+import java.util.Map;
 
 public class ServerTunnelChannelWrapper extends TunnelChannelWrapper {
     public HashMap<Integer, ChannelWrapper> connections;
@@ -36,4 +37,14 @@ public void startReading(int connectionId) {
         connection.addInterest(SelectionKey.OP_READ);
         connection.stopReading = false;
     }
+
+    @Override
+    public void cleanup() {
+        for (Map.Entry<Integer, ChannelWrapper> entry : connections.entrySet()) {
+            ChannelWrapper wrapper = entry.getValue();
+            wrapper.disconnectListener = null;
+            wrapper.close();
+        }
+        connections.clear();
+    }
 }