From e0cb283ab756df5a95bbb6647656e9072835054e Mon Sep 17 00:00:00 2001
From: UnfamiliarLegacy <74633542+UnfamiliarLegacy@users.noreply.github.com>
Date: Wed, 1 Dec 2021 23:38:33 +0100
Subject: [PATCH] Properly abort when client closes connection

---
 .../proxy/nitro/NitroProxyProvider.java          | 16 ++++++++++++++++
 .../nitro/websocket/NitroWebsocketClient.java    |  2 +-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/NitroProxyProvider.java b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/NitroProxyProvider.java
index b8c0810..5b90922 100644
--- a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/NitroProxyProvider.java
+++ b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/NitroProxyProvider.java
@@ -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) {
diff --git a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java
index 59bf027..ca7e88b 100644
--- a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java
+++ b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java
@@ -112,7 +112,7 @@ public class NitroWebsocketClient implements NitroSession {
 
             // Reset program state.
             proxySetter.setProxy(null);
-            stateSetter.setState(HState.NOT_CONNECTED);
+            stateSetter.setState(HState.ABORTING);
         }
     }
 }