Better nitro connection state management

G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/NitroConnectionState.java

@@ -0,0 +1,58 @@
+package gearth.protocol.connection.proxy.nitro;
+
+import gearth.protocol.HMessage;
+import gearth.protocol.connection.HState;
+import gearth.protocol.connection.HStateSetter;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class NitroConnectionState {
+
+    private static final Logger logger = LoggerFactory.getLogger(NitroConnectionState.class);
+
+    private final HStateSetter stateSetter;
+
+    private boolean aborting;
+    private boolean toServer;
+    private boolean toClient;
+
+    public NitroConnectionState(HStateSetter stateSetter) {
+        this.stateSetter = stateSetter;
+    }
+
+    public void setConnected(HMessage.Direction direction) {
+        if (direction == HMessage.Direction.TOCLIENT) {
+            this.toClient = true;
+        } else if (direction == HMessage.Direction.TOSERVER) {
+            this.toServer = true;
+        }
+
+        this.checkConnected();
+    }
+
+    public void checkConnected() {
+        if (this.aborting) {
+            return;
+        }
+
+        if (!this.toClient) {
+            return;
+        }
+
+        if (!this.toServer) {
+            return;
+        }
+
+        this.stateSetter.setState(HState.CONNECTED);
+
+        logger.info("Connected");
+    }
+
+    public void setAborting() {
+        this.aborting = true;
+        this.stateSetter.setState(HState.ABORTING);
+
+        logger.info("Aborting");
+    }
+
+}

G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java

@@ -3,6 +3,7 @@ package gearth.protocol.connection.proxy.nitro.websocket;
 import gearth.protocol.HConnection;
 import gearth.protocol.HMessage;
 import gearth.protocol.connection.*;
+import gearth.protocol.connection.proxy.nitro.NitroConnectionState;
 import gearth.protocol.connection.proxy.nitro.NitroConstants;
 import gearth.protocol.connection.proxy.nitro.NitroProxyProvider;
 import gearth.protocol.packethandler.nitro.NitroPacketHandler;
@@ -25,8 +26,7 @@ public class NitroWebsocketClient implements NitroSession {
     private static final Logger logger = LoggerFactory.getLogger(NitroWebsocketClient.class);
 
     private final HProxySetter proxySetter;
-    private final HStateSetter stateSetter;
+    private final NitroConnectionState state;
-    private final HConnection connection;
     private final NitroProxyProvider proxyProvider;
     private final NitroWebsocketServer server;
     private final NitroPacketHandler packetHandler;
@@ -36,10 +36,9 @@ public class NitroWebsocketClient implements NitroSession {
 
     public NitroWebsocketClient(HProxySetter proxySetter, HStateSetter stateSetter, HConnection connection, NitroProxyProvider proxyProvider) {
         this.proxySetter = proxySetter;
-        this.stateSetter = stateSetter;
+        this.state = new NitroConnectionState(stateSetter);
-        this.connection = connection;
         this.proxyProvider = proxyProvider;
-        this.server = new NitroWebsocketServer(connection, this);
+        this.server = new NitroWebsocketServer(connection, this, this.state);
         this.packetHandler = new NitroPacketHandler(HMessage.Direction.TOSERVER, server, connection.getExtensionHandler(), connection.getTrafficObservables());
         this.shutdownLock = new AtomicBoolean();
     }
@@ -72,11 +71,13 @@ public class NitroWebsocketClient implements NitroSession {
         );
 
         proxySetter.setProxy(proxy);
-        stateSetter.setState(HState.CONNECTED);
+        state.setConnected(HMessage.Direction.TOSERVER);
     }
 
     @OnMessage
     public void onMessage(byte[] b, Session session) throws IOException {
+        logger.debug("Received packet from browser");
+
         packetHandler.act(b);
     }
 
@@ -133,7 +134,7 @@ public class NitroWebsocketClient implements NitroSession {
 
             // Reset program state.
             proxySetter.setProxy(null);
-            stateSetter.setState(HState.ABORTING);
+            state.setAborting();
         }
     }
 }

G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketServer.java

@@ -2,6 +2,7 @@ package gearth.protocol.connection.proxy.nitro.websocket;
 
 import gearth.protocol.HConnection;
 import gearth.protocol.HMessage;
+import gearth.protocol.connection.proxy.nitro.NitroConnectionState;
 import gearth.protocol.connection.proxy.nitro.NitroConstants;
 import gearth.protocol.packethandler.PacketHandler;
 import gearth.protocol.packethandler.nitro.NitroPacketHandler;
@@ -47,10 +48,12 @@ public class NitroWebsocketServer implements WebSocketListener, NitroSession {
 
     private final PacketHandler packetHandler;
     private final NitroWebsocketClient client;
+    private final NitroConnectionState state;
     private Session activeSession = null;
 
-    public NitroWebsocketServer(HConnection connection, NitroWebsocketClient client) {
+    public NitroWebsocketServer(HConnection connection, NitroWebsocketClient client, NitroConnectionState state) {
         this.client = client;
+        this.state = state;
         this.packetHandler = new NitroPacketHandler(HMessage.Direction.TOCLIENT, client, connection.getExtensionHandler(), connection.getTrafficObservables());
     }
 
@@ -83,8 +86,6 @@ public class NitroWebsocketServer implements WebSocketListener, NitroSession {
 
             client.start();
             client.connect(this, URI.create(websocketUrl), request);
-
-            logger.info("Connected to origin websocket");
         } catch (Exception e) {
             throw new IOException("Failed to start websocket client to origin " + websocketUrl, e);
         }
@@ -139,6 +140,7 @@ public class NitroWebsocketServer implements WebSocketListener, NitroSession {
     @Override
     public void onWebSocketConnect(org.eclipse.jetty.websocket.api.Session session) {
         activeSession = session;
+        state.setConnected(HMessage.Direction.TOCLIENT);
     }
 
     @Override