From ae4ccff20ec9cc12a0618177924b4d5186fac55e Mon Sep 17 00:00:00 2001 From: UnfamiliarLegacy <74633542+UnfamiliarLegacy@users.noreply.github.com> Date: Fri, 26 Nov 2021 18:51:31 +0100 Subject: [PATCH] Add proper shutdown mechanism for Nitro --- .../nitro/websocket/NitroPacketHandler.java | 9 +++- .../nitro/websocket/NitroWebsocketClient.java | 53 ++++++++++++++++--- .../nitro/websocket/NitroWebsocketServer.java | 33 +++++++++--- 3 files changed, 80 insertions(+), 15 deletions(-) diff --git a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroPacketHandler.java b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroPacketHandler.java index 8ea0ee8..d6ed58a 100644 --- a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroPacketHandler.java +++ b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroPacketHandler.java @@ -6,6 +6,7 @@ import gearth.protocol.packethandler.PacketHandler; import gearth.services.extension_handler.ExtensionHandler; import gearth.services.extension_handler.OnHMessageHandled; +import javax.websocket.Session; import java.io.IOException; import java.nio.ByteBuffer; @@ -22,12 +23,18 @@ public class NitroPacketHandler extends PacketHandler { @Override public boolean sendToStream(byte[] buffer) { + final Session localSession = session.getSession(); + + if (localSession == null) { + return false; + } + // Required to prevent garbage buffer within the UI logger. if (direction == HMessage.Direction.TOSERVER) { buffer = buffer.clone(); } - session.getSession().getAsyncRemote().sendBinary(ByteBuffer.wrap(buffer)); + localSession.getAsyncRemote().sendBinary(ByteBuffer.wrap(buffer)); return true; } diff --git a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java index ab890da..aabb946 100644 --- a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java +++ b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java @@ -9,6 +9,7 @@ import gearth.protocol.connection.proxy.nitro.NitroProxyProvider; import javax.websocket.*; import javax.websocket.server.ServerEndpoint; import java.io.IOException; +import java.util.concurrent.atomic.AtomicBoolean; @ServerEndpoint(value = "/") public class NitroWebsocketClient implements NitroSession { @@ -19,9 +20,9 @@ public class NitroWebsocketClient implements NitroSession { private final NitroProxyProvider proxyProvider; private final NitroWebsocketServer server; private final NitroPacketHandler packetHandler; + private final AtomicBoolean shutdownLock; private Session activeSession = null; - private HProxy proxy = null; public NitroWebsocketClient(HProxySetter proxySetter, HStateSetter stateSetter, HConnection connection, NitroProxyProvider proxyProvider) { this.proxySetter = proxySetter; @@ -30,6 +31,7 @@ public class NitroWebsocketClient implements NitroSession { this.proxyProvider = proxyProvider; this.server = new NitroWebsocketServer(connection, this); this.packetHandler = new NitroPacketHandler(HMessage.Direction.TOSERVER, server, connection.getExtensionHandler(), connection.getTrafficObservables()); + this.shutdownLock = new AtomicBoolean(); } @OnOpen @@ -39,7 +41,8 @@ public class NitroWebsocketClient implements NitroSession { server.connect(proxyProvider.getOriginalWebsocketUrl()); - proxy = new HProxy(HClient.NITRO, "", "", -1, -1, ""); + final HProxy proxy = new HProxy(HClient.NITRO, "", "", -1, -1, ""); + proxy.verifyProxy( this.server.getPacketHandler(), this.packetHandler, @@ -53,25 +56,63 @@ public class NitroWebsocketClient implements NitroSession { @OnMessage public void onMessage(byte[] b, Session session) throws IOException { - System.out.printf("onMessage (%d)%n", b.length); - System.out.println(session); - System.out.println(Hexdump.hexdump(b)); - packetHandler.act(b); } @OnClose public void onClose(Session session) throws IOException { activeSession = null; + shutdownProxy(); } @OnError public void onError(Session session, Throwable throwable) { + throwable.printStackTrace(); + // Shutdown. + shutdownProxy(); } @Override public Session getSession() { return activeSession; } + + /** + * Shutdown and clean up the client connection. + */ + private void shutdown() { + if (activeSession == null) { + return; + } + + try { + activeSession.close(); + } catch (IOException e) { + e.printStackTrace(); + } finally { + activeSession = null; + } + } + + /** + * Shutdown all connections and reset program state. + */ + public void shutdownProxy() { + if (shutdownLock.get()) { + return; + } + + if (shutdownLock.compareAndSet(false, true)) { + // Close client connection. + shutdown(); + + // Close server connection. + server.shutdown(); + + // Reset program state. + proxySetter.setProxy(null); + stateSetter.setState(HState.NOT_CONNECTED); + } + } } diff --git a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketServer.java b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketServer.java index f7928df..04b238c 100644 --- a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketServer.java +++ b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketServer.java @@ -13,9 +13,11 @@ import java.net.URI; public class NitroWebsocketServer implements NitroSession { private final PacketHandler packetHandler; + private final NitroWebsocketClient client; private Session activeSession = null; public NitroWebsocketServer(HConnection connection, NitroWebsocketClient client) { + this.client = client; this.packetHandler = new NitroPacketHandler(HMessage.Direction.TOCLIENT, client, connection.getExtensionHandler(), connection.getTrafficObservables()); } @@ -35,23 +37,21 @@ public class NitroWebsocketServer implements NitroSession { @OnMessage public void onMessage(byte[] b, Session session) throws IOException { - //System.out.printf("onMessage (%d)%n", b.length); - //System.out.println(session); - //System.out.println(Hexdump.hexdump(b)); - packetHandler.act(b); } @OnClose public void onClose(Session userSession, CloseReason reason) { - //System.out.println("closing websocket"); + // Hotel closed connection. + client.shutdownProxy(); } @OnError public void onError(Session session, Throwable throwable) { - //System.out.println("onError"); - //System.out.println(session); - //System.out.println(throwable); + throwable.printStackTrace(); + + // Shutdown. + client.shutdownProxy(); } @Override @@ -62,4 +62,21 @@ public class NitroWebsocketServer implements NitroSession { public PacketHandler getPacketHandler() { return packetHandler; } + + /** + * Shutdown and clean up the server connection. + */ + public void shutdown() { + if (activeSession == null) { + return; + } + + try { + activeSession.close(); + } catch (IOException e) { + e.printStackTrace(); + } finally { + activeSession = null; + } + } }