From eece15d4159e9821e6ec144c2eaed27b00b93723 Mon Sep 17 00:00:00 2001
From: Luck <git@lucko.me>
Date: Sun, 20 Mar 2022 10:45:59 +0000
Subject: [PATCH] Close web editor sockets on plugin disable (#3326)

---
 .../common/plugin/AbstractLuckPermsPlugin.java        |  8 ++++++++
 .../common/webeditor/socket/WebEditorSocket.java      |  1 +
 .../common/webeditor/store/WebEditorSocketMap.java    | 11 +++++++++--
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/common/src/main/java/me/lucko/luckperms/common/plugin/AbstractLuckPermsPlugin.java b/common/src/main/java/me/lucko/luckperms/common/plugin/AbstractLuckPermsPlugin.java
index f6f2fd806..9bb3b7c1a 100644
--- a/common/src/main/java/me/lucko/luckperms/common/plugin/AbstractLuckPermsPlugin.java
+++ b/common/src/main/java/me/lucko/luckperms/common/plugin/AbstractLuckPermsPlugin.java
@@ -58,6 +58,7 @@ import me.lucko.luckperms.common.tasks.ExpireTemporaryTask;
 import me.lucko.luckperms.common.tasks.SyncTask;
 import me.lucko.luckperms.common.treeview.PermissionRegistry;
 import me.lucko.luckperms.common.verbose.VerboseHandler;
+import me.lucko.luckperms.common.webeditor.socket.WebEditorSocket;
 import me.lucko.luckperms.common.webeditor.store.WebEditorStore;
 
 import net.luckperms.api.LuckPerms;
@@ -234,6 +235,13 @@ public abstract class AbstractLuckPermsPlugin implements LuckPermsPlugin {
         // cancel delayed/repeating tasks
         getBootstrap().getScheduler().shutdownScheduler();
 
+        // close web editor sockets
+        for (WebEditorSocket socket : this.webEditorStore.sockets().getSockets()) {
+            if (!socket.isClosed()) {
+                socket.close();
+            }
+        }
+
         // shutdown permission vault and verbose handler tasks
         this.permissionRegistry.close();
         this.verboseHandler.close();
diff --git a/common/src/main/java/me/lucko/luckperms/common/webeditor/socket/WebEditorSocket.java b/common/src/main/java/me/lucko/luckperms/common/webeditor/socket/WebEditorSocket.java
index 74c82f17a..221b502e0 100644
--- a/common/src/main/java/me/lucko/luckperms/common/webeditor/socket/WebEditorSocket.java
+++ b/common/src/main/java/me/lucko/luckperms/common/webeditor/socket/WebEditorSocket.java
@@ -217,6 +217,7 @@ public class WebEditorSocket {
 
     private void closeSocket() {
         this.socket.socket().close(1000, "Normal");
+        this.plugin.getWebEditorStore().sockets().removeSocket(this);
         this.closed = true;
     }
 
diff --git a/common/src/main/java/me/lucko/luckperms/common/webeditor/store/WebEditorSocketMap.java b/common/src/main/java/me/lucko/luckperms/common/webeditor/store/WebEditorSocketMap.java
index 9e7fb2de7..9462cc6dc 100644
--- a/common/src/main/java/me/lucko/luckperms/common/webeditor/store/WebEditorSocketMap.java
+++ b/common/src/main/java/me/lucko/luckperms/common/webeditor/store/WebEditorSocketMap.java
@@ -31,13 +31,12 @@ import me.lucko.luckperms.common.sender.Sender;
 import me.lucko.luckperms.common.util.CaffeineFactory;
 import me.lucko.luckperms.common.webeditor.socket.WebEditorSocket;
 
+import java.util.Collection;
 import java.util.UUID;
-import java.util.concurrent.TimeUnit;
 
 public final class WebEditorSocketMap {
     private final Cache<UUID, WebEditorSocket> sockets = CaffeineFactory.newBuilder()
             .weakValues()
-            .expireAfterWrite(5, TimeUnit.MINUTES)
             .build();
 
     public WebEditorSocket getSocket(Sender sender) {
@@ -47,4 +46,12 @@ public final class WebEditorSocketMap {
     public void putSocket(Sender sender, WebEditorSocket socket) {
         this.sockets.put(sender.getUniqueId(), socket);
     }
+
+    public void removeSocket(WebEditorSocket socket) {
+        this.sockets.asMap().values().remove(socket);
+    }
+
+    public Collection<WebEditorSocket> getSockets() {
+        return this.sockets.asMap().values();
+    }
 }