Add proper shutdown mechanism for Nitro

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;
     }
 

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);
+        }
+    }
 }

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;
+        }
+    }
 }