Properly abort when client closes connection

This commit is contained in:
UnfamiliarLegacy 2021-12-01 23:38:33 +01:00
parent f49906897c
commit e0cb283ab7
2 changed files with 17 additions and 1 deletions

View File

@ -13,6 +13,7 @@ import gearth.protocol.connection.proxy.nitro.websocket.NitroWebsocketProxy;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.concurrent.atomic.AtomicBoolean;
public class NitroProxyProvider implements ProxyProvider, NitroHttpProxyServerCallback, StateChangeListener {
@ -21,6 +22,7 @@ public class NitroProxyProvider implements ProxyProvider, NitroHttpProxyServerCa
private final HConnection connection;
private final NitroHttpProxy nitroHttpProxy;
private final NitroWebsocketProxy nitroWebsocketProxy;
private final AtomicBoolean abortLock;
private String originalWebsocketUrl;
private String originalOriginUrl;
@ -31,6 +33,7 @@ public class NitroProxyProvider implements ProxyProvider, NitroHttpProxyServerCa
this.connection = connection;
this.nitroHttpProxy = new NitroHttpProxy(this);
this.nitroWebsocketProxy = new NitroWebsocketProxy(proxySetter, stateSetter, connection, this);
this.abortLock = new AtomicBoolean();
}
public String getOriginalWebsocketUrl() {
@ -62,6 +65,14 @@ public class NitroProxyProvider implements ProxyProvider, NitroHttpProxyServerCa
@Override
public void abort() {
if (abortLock.get()) {
return;
}
if (abortLock.compareAndSet(true, true)) {
return;
}
stateSetter.setState(HState.ABORTING);
new Thread(() -> {
@ -98,6 +109,11 @@ public class NitroProxyProvider implements ProxyProvider, NitroHttpProxyServerCa
// We are not stopping the http proxy because some requests might still require it to be running.
nitroHttpProxy.pause();
}
// Catch setState ABORTING inside NitroWebsocketClient.
if (newState == HState.ABORTING) {
abort();
}
}
private static String extractOriginUrl(String url) {

View File

@ -112,7 +112,7 @@ public class NitroWebsocketClient implements NitroSession {
// Reset program state.
proxySetter.setProxy(null);
stateSetter.setState(HState.NOT_CONNECTED);
stateSetter.setState(HState.ABORTING);
}
}
}