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

View File

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