diff --git a/Client/src/tk/jackyliao123/proxy/client/socks/Socks5ClientTunnelChannelWrapper.java b/Client/src/tk/jackyliao123/proxy/client/socks/Socks5ClientTunnelChannelWrapper.java index e2f6154..50adf77 100644 --- 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 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 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 index c9ba4c6..777f287 100644 --- 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 index 6da3526..9dcb25a 100644 --- 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 index 463e886..f076de8 100644 --- 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 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 entry : connections.entrySet()) { + ChannelWrapper wrapper = entry.getValue(); + wrapper.disconnectListener = null; + wrapper.close(); + } + connections.clear(); + } }