Add proper shutdown mechanism for Nitro

This commit is contained in:
UnfamiliarLegacy 2021-11-26 18:51:31 +01:00
parent dd0312c020
commit ae4ccff20e
3 changed files with 80 additions and 15 deletions

View File

@ -6,6 +6,7 @@ import gearth.protocol.packethandler.PacketHandler;
import gearth.services.extension_handler.ExtensionHandler; import gearth.services.extension_handler.ExtensionHandler;
import gearth.services.extension_handler.OnHMessageHandled; import gearth.services.extension_handler.OnHMessageHandled;
import javax.websocket.Session;
import java.io.IOException; import java.io.IOException;
import java.nio.ByteBuffer; import java.nio.ByteBuffer;
@ -22,12 +23,18 @@ public class NitroPacketHandler extends PacketHandler {
@Override @Override
public boolean sendToStream(byte[] buffer) { public boolean sendToStream(byte[] buffer) {
final Session localSession = session.getSession();
if (localSession == null) {
return false;
}
// Required to prevent garbage buffer within the UI logger. // Required to prevent garbage buffer within the UI logger.
if (direction == HMessage.Direction.TOSERVER) { if (direction == HMessage.Direction.TOSERVER) {
buffer = buffer.clone(); buffer = buffer.clone();
} }
session.getSession().getAsyncRemote().sendBinary(ByteBuffer.wrap(buffer)); localSession.getAsyncRemote().sendBinary(ByteBuffer.wrap(buffer));
return true; return true;
} }

View File

@ -9,6 +9,7 @@ import gearth.protocol.connection.proxy.nitro.NitroProxyProvider;
import javax.websocket.*; import javax.websocket.*;
import javax.websocket.server.ServerEndpoint; import javax.websocket.server.ServerEndpoint;
import java.io.IOException; import java.io.IOException;
import java.util.concurrent.atomic.AtomicBoolean;
@ServerEndpoint(value = "/") @ServerEndpoint(value = "/")
public class NitroWebsocketClient implements NitroSession { public class NitroWebsocketClient implements NitroSession {
@ -19,9 +20,9 @@ public class NitroWebsocketClient implements NitroSession {
private final NitroProxyProvider proxyProvider; private final NitroProxyProvider proxyProvider;
private final NitroWebsocketServer server; private final NitroWebsocketServer server;
private final NitroPacketHandler packetHandler; private final NitroPacketHandler packetHandler;
private final AtomicBoolean shutdownLock;
private Session activeSession = null; private Session activeSession = null;
private HProxy proxy = null;
public NitroWebsocketClient(HProxySetter proxySetter, HStateSetter stateSetter, HConnection connection, NitroProxyProvider proxyProvider) { public NitroWebsocketClient(HProxySetter proxySetter, HStateSetter stateSetter, HConnection connection, NitroProxyProvider proxyProvider) {
this.proxySetter = proxySetter; this.proxySetter = proxySetter;
@ -30,6 +31,7 @@ public class NitroWebsocketClient implements NitroSession {
this.proxyProvider = proxyProvider; this.proxyProvider = proxyProvider;
this.server = new NitroWebsocketServer(connection, this); this.server = new NitroWebsocketServer(connection, this);
this.packetHandler = new NitroPacketHandler(HMessage.Direction.TOSERVER, server, connection.getExtensionHandler(), connection.getTrafficObservables()); this.packetHandler = new NitroPacketHandler(HMessage.Direction.TOSERVER, server, connection.getExtensionHandler(), connection.getTrafficObservables());
this.shutdownLock = new AtomicBoolean();
} }
@OnOpen @OnOpen
@ -39,7 +41,8 @@ public class NitroWebsocketClient implements NitroSession {
server.connect(proxyProvider.getOriginalWebsocketUrl()); server.connect(proxyProvider.getOriginalWebsocketUrl());
proxy = new HProxy(HClient.NITRO, "", "", -1, -1, ""); final HProxy proxy = new HProxy(HClient.NITRO, "", "", -1, -1, "");
proxy.verifyProxy( proxy.verifyProxy(
this.server.getPacketHandler(), this.server.getPacketHandler(),
this.packetHandler, this.packetHandler,
@ -53,25 +56,63 @@ public class NitroWebsocketClient implements NitroSession {
@OnMessage @OnMessage
public void onMessage(byte[] b, Session session) throws IOException { 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); packetHandler.act(b);
} }
@OnClose @OnClose
public void onClose(Session session) throws IOException { public void onClose(Session session) throws IOException {
activeSession = null; activeSession = null;
shutdownProxy();
} }
@OnError @OnError
public void onError(Session session, Throwable throwable) { public void onError(Session session, Throwable throwable) {
throwable.printStackTrace();
// Shutdown.
shutdownProxy();
} }
@Override @Override
public Session getSession() { public Session getSession() {
return activeSession; 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);
}
}
} }

View File

@ -13,9 +13,11 @@ import java.net.URI;
public class NitroWebsocketServer implements NitroSession { public class NitroWebsocketServer implements NitroSession {
private final PacketHandler packetHandler; private final PacketHandler packetHandler;
private final NitroWebsocketClient client;
private Session activeSession = null; private Session activeSession = null;
public NitroWebsocketServer(HConnection connection, NitroWebsocketClient client) { public NitroWebsocketServer(HConnection connection, NitroWebsocketClient client) {
this.client = client;
this.packetHandler = new NitroPacketHandler(HMessage.Direction.TOCLIENT, client, connection.getExtensionHandler(), connection.getTrafficObservables()); this.packetHandler = new NitroPacketHandler(HMessage.Direction.TOCLIENT, client, connection.getExtensionHandler(), connection.getTrafficObservables());
} }
@ -35,23 +37,21 @@ public class NitroWebsocketServer implements NitroSession {
@OnMessage @OnMessage
public void onMessage(byte[] b, Session session) throws IOException { 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); packetHandler.act(b);
} }
@OnClose @OnClose
public void onClose(Session userSession, CloseReason reason) { public void onClose(Session userSession, CloseReason reason) {
//System.out.println("closing websocket"); // Hotel closed connection.
client.shutdownProxy();
} }
@OnError @OnError
public void onError(Session session, Throwable throwable) { public void onError(Session session, Throwable throwable) {
//System.out.println("onError"); throwable.printStackTrace();
//System.out.println(session);
//System.out.println(throwable); // Shutdown.
client.shutdownProxy();
} }
@Override @Override
@ -62,4 +62,21 @@ public class NitroWebsocketServer implements NitroSession {
public PacketHandler getPacketHandler() { public PacketHandler getPacketHandler() {
return packetHandler; 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;
}
}
} }