Better nitro connection state management

This commit is contained in:
UnfamiliarLegacy 2023-06-07 19:35:21 +02:00
parent df3aabf518
commit c7037c0441
3 changed files with 71 additions and 10 deletions

View File

@ -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");
}
}

View File

@ -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 HConnection connection;
private final NitroConnectionState state;
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.connection = connection;
this.state = new NitroConnectionState(stateSetter);
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();
}
}
}

View File

@ -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