Upstream

api/src/main/java/org/geysermc/floodgate/api/player/PropertyKey.java

@@ -78,6 +78,16 @@ public Result isAddAllowed(Object obj) {
         return Result.NOT_EQUALS;
     }
 
+    @Override
+    public boolean equals(Object obj) {
+        return obj instanceof PropertyKey && key.equals(((PropertyKey) obj).key);
+    }
+
+    @Override
+    public int hashCode() {
+        return key.hashCode();
+    }
+
     public enum Result {
         NOT_EQUALS,
         INVALID_TAGS,

bungee/src/main/java/org/geysermc/floodgate/listener/BungeeListener.java

@@ -47,6 +47,7 @@
 import org.geysermc.floodgate.api.ProxyFloodgateApi;
 import org.geysermc.floodgate.api.logger.FloodgateLogger;
 import org.geysermc.floodgate.api.player.FloodgatePlayer;
+import org.geysermc.floodgate.pluginmessage.channel.FormChannel;
 import org.geysermc.floodgate.skin.SkinApplier;
 import org.geysermc.floodgate.skin.SkinDataImpl;
 import org.geysermc.floodgate.util.LanguageManager;
@@ -82,6 +83,7 @@ public final class BungeeListener implements Listener {
     private AttributeKey<String> kickMessageAttribute;
 
     @Inject private MojangUtils mojangUtils;
+    @Inject private FormChannel formChannel;
 
     @EventHandler(priority = EventPriority.LOWEST)
     public void onPreLogin(PreLoginEvent event) {
@@ -159,6 +161,11 @@ public void onPostLogin(PostLoginEvent event) {
 
     @EventHandler(priority = EventPriority.HIGHEST)
     public void onPlayerDisconnect(PlayerDisconnectEvent event) {
+        FloodgatePlayer player = api.getPendingRemovePlayer(event.getPlayer().getUniqueId());
+        if (player != null) {
+            formChannel.disconnect(player);
+        }
+
         api.playerRemoved(event.getPlayer().getUniqueId());
     }
 }

bungee/src/main/java/org/geysermc/floodgate/module/BungeePlatformModule.java

@@ -111,8 +111,9 @@ public ListenerRegistration<Listener> listenerRegistration() {
     @Singleton
     public PluginMessageUtils pluginMessageUtils(
             PluginMessageManager manager,
+            FloodgateApi api,
             FloodgateLogger logger) {
-        return new BungeePluginMessageUtils(manager, logger);
+        return new BungeePluginMessageUtils(manager, api, logger);
     }
 
     @Provides

bungee/src/main/java/org/geysermc/floodgate/pluginmessage/BungeePluginMessageUtils.java

@@ -36,14 +36,17 @@
 import net.md_5.bungee.api.plugin.Listener;
 import net.md_5.bungee.event.EventHandler;
 import net.md_5.bungee.event.EventPriority;
+import org.geysermc.floodgate.api.FloodgateApi;
 import org.geysermc.floodgate.api.logger.FloodgateLogger;
+import org.geysermc.floodgate.api.player.FloodgatePlayer;
 import org.geysermc.floodgate.platform.pluginmessage.PluginMessageUtils;
 import org.geysermc.floodgate.pluginmessage.PluginMessageChannel.Identity;
 import org.geysermc.floodgate.pluginmessage.PluginMessageChannel.Result;
 
 @RequiredArgsConstructor
 public final class BungeePluginMessageUtils extends PluginMessageUtils implements Listener {
     private final PluginMessageManager pluginMessageManager;
+    private final FloodgateApi api;
     private final FloodgateLogger logger;
 
     @EventHandler(priority = EventPriority.LOW)
@@ -53,23 +56,26 @@ public void onPluginMessage(PluginMessageEvent event) {
             return;
         }
 
-        UUID sourceUuid = null;
-        String sourceUsername = null;
+        FloodgatePlayer fSource = null;
         Identity sourceIdentity = Identity.UNKNOWN;
 
         Connection source = event.getSender();
         if (source instanceof ProxiedPlayer) {
             ProxiedPlayer player = (ProxiedPlayer) source;
-            sourceUuid = player.getUniqueId();
-            sourceUsername = player.getName();
+            fSource = api.getPlayer(player.getUniqueId());
             sourceIdentity = Identity.PLAYER;
 
+            if (fSource == null) {
+                logKick(source, "Only Floodgate players can send floodgate messages!");
+                return;
+            }
+
         } else if (source instanceof ServerConnection) {
             sourceIdentity = Identity.SERVER;
         }
 
         Result result = channel.handleProxyCall(
-                event.getData(), sourceUuid, sourceUsername, sourceIdentity
+                event.getData(), fSource, sourceIdentity
         );
 
         event.setCancelled(!result.isAllowed());

core/src/main/java/org/geysermc/floodgate/api/SimpleFloodgateApi.java

@@ -117,7 +117,13 @@ public boolean isFloodgateId(UUID uuid) {
 
     @Override
     public boolean sendForm(UUID uuid, Form form) {
-        return pluginMessageManager.getChannel(FormChannel.class).sendForm(uuid, form);
+        FloodgatePlayer player = getPlayer(uuid);
+        // Before this check was added, we used to just send the form to the user no matter if they
+        // were a FloodgatePlayer or not. Keep this since the Floodgate API is deprecated anyway.
+        if (player == null) {
+            return true;
+        }
+        return pluginMessageManager.getChannel(FormChannel.class).sendForm(player, form);
     }
 
     @Override
@@ -127,7 +133,13 @@ public boolean sendForm(UUID uuid, FormBuilder<?, ?, ?> formBuilder) {
 
     @Override
     public boolean closeForm(UUID uuid) {
-        return pluginMessageManager.getChannel(FormChannel.class).closeForm(uuid);
+        FloodgatePlayer player = getPlayer(uuid);
+        // Before this check was added, we used to just send the form to the user no matter if they
+        // were a FloodgatePlayer or not. Keep this since the Floodgate API is deprecated anyway.
+        if (player == null) {
+            return true;
+        }
+        return pluginMessageManager.getChannel(FormChannel.class).closeForm(player);
     }
 
     @Override
@@ -217,7 +229,7 @@ public void playerRemoved(UUID correctUuid) {
         }
     }
 
-    private FloodgatePlayer getPendingRemovePlayer(UUID correctUuid) {
+    public FloodgatePlayer getPendingRemovePlayer(UUID correctUuid) {
         for (FloodgatePlayer player : pendingRemove.asMap().values()) {
             if (player.getCorrectUniqueId().equals(correctUuid)) {
                 return player;

core/src/main/java/org/geysermc/floodgate/player/FloodgatePlayerImpl.java

@@ -25,9 +25,10 @@
 
 package org.geysermc.floodgate.player;
 
-import java.util.HashMap;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Supplier;
 import lombok.AccessLevel;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
@@ -176,8 +177,8 @@ public <T> T removeProperty(PropertyKey key) {
     @Override
     public <T> T addProperty(PropertyKey key, Object value) {
         if (stringToPropertyKey == null) {
-            stringToPropertyKey = new HashMap<>();
-            propertyKeyToValue = new HashMap<>();
+            stringToPropertyKey = new ConcurrentHashMap<>();
+            propertyKeyToValue = new ConcurrentHashMap<>();
 
             stringToPropertyKey.put(key.getKey(), key);
             propertyKeyToValue.put(key, value);
@@ -202,8 +203,8 @@ public <T> T addProperty(String key, Object value) {
         PropertyKey propertyKey = new PropertyKey(key, true, true);
 
         if (stringToPropertyKey == null) {
-            stringToPropertyKey = new HashMap<>();
-            propertyKeyToValue = new HashMap<>();
+            stringToPropertyKey = new ConcurrentHashMap<>();
+            propertyKeyToValue = new ConcurrentHashMap<>();
 
             stringToPropertyKey.put(key, propertyKey);
             propertyKeyToValue.put(propertyKey, value);
@@ -223,4 +224,17 @@ public <T> T addProperty(String key, Object value) {
             return propertyKey;
         });
     }
+
+    public <T> T getOrAddProperty(PropertyKey key, Supplier<T> supplier) {
+        if (stringToPropertyKey == null) {
+            stringToPropertyKey = new ConcurrentHashMap<>();
+            propertyKeyToValue = new ConcurrentHashMap<>();
+        }
+
+        // The hashcode & equals of PropertyKey is based on the hashcode of key.
+        // stringToPropertyKey is solely for the updatable & removable checks, which we still handle
+        // correctly by using ifAbsent for both.
+        stringToPropertyKey.putIfAbsent(key.getKey(), key);
+        return (T) propertyKeyToValue.computeIfAbsent(key, (unused) -> supplier.get());
+    }
 }

core/src/main/java/org/geysermc/floodgate/pluginmessage/PluginMessageChannel.java

@@ -25,23 +25,22 @@
 
 package org.geysermc.floodgate.pluginmessage;
 
-import java.util.UUID;
 import lombok.AccessLevel;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
+import org.geysermc.floodgate.api.player.FloodgatePlayer;
 
 public interface PluginMessageChannel {
 
     String getIdentifier();
 
     Result handleProxyCall(
             byte[] data,
-            UUID sourceUuid,
-            String sourceUsername,
+            FloodgatePlayer source,
             Identity sourceIdentity
     );
 
-    Result handleServerCall(byte[] data, UUID playerUuid, String playerUsername);
+    Result handleServerCall(byte[] data, FloodgatePlayer source);
 
     @Getter
     @RequiredArgsConstructor(access = AccessLevel.PRIVATE)

core/src/main/java/org/geysermc/floodgate/pluginmessage/channel/FormChannel.java

@@ -27,24 +27,25 @@
 
 import com.google.common.base.Charsets;
 import com.google.inject.Inject;
-import it.unimi.dsi.fastutil.shorts.Short2ObjectMap;
-import it.unimi.dsi.fastutil.shorts.Short2ObjectMaps;
-import it.unimi.dsi.fastutil.shorts.Short2ObjectOpenHashMap;
-import java.util.UUID;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.geysermc.cumulus.form.Form;
 import org.geysermc.cumulus.form.impl.FormDefinition;
 import org.geysermc.cumulus.form.impl.FormDefinitions;
 import org.geysermc.floodgate.api.logger.FloodgateLogger;
+import org.geysermc.floodgate.api.player.FloodgatePlayer;
+import org.geysermc.floodgate.api.player.PropertyKey;
 import org.geysermc.floodgate.config.FloodgateConfig;
 import org.geysermc.floodgate.platform.pluginmessage.PluginMessageUtils;
+import org.geysermc.floodgate.player.FloodgatePlayerImpl;
 import org.geysermc.floodgate.pluginmessage.PluginMessageChannel;
 
 public class FormChannel implements PluginMessageChannel {
+    private static final PropertyKey PROPERTY_LAST_FORM_ID = new PropertyKey("floodgate:last_form_id", true, false);
+    private static final PropertyKey PROPERTY_ACTIVE_FORMS = new PropertyKey("floodgate:active_forms", true, true);
+
     private final FormDefinitions formDefinitions = FormDefinitions.instance();
-    private final Short2ObjectMap<Form> storedForms =
-            Short2ObjectMaps.synchronize(new Short2ObjectOpenHashMap<>());
-    private final AtomicInteger nextFormId = new AtomicInteger(0);
 
     @Inject private PluginMessageUtils pluginMessageUtils;
     @Inject private FloodgateConfig config;
@@ -58,8 +59,7 @@ public String getIdentifier() {
     @Override
     public Result handleProxyCall(
             byte[] data,
-            UUID sourceUuid,
-            String sourceUsername,
+            FloodgatePlayer source,
             Identity sourceIdentity
     ) {
         if (sourceIdentity == Identity.SERVER) {
@@ -79,35 +79,54 @@ public Result handleProxyCall(
                 return Result.forward();
             }
 
-            if (!callResponseConsumer(data)) {
+            if (!callResponseConsumer(source, data)) {
                 logger.error("Couldn't find stored form with id {} for player {}",
-                        formId, sourceUsername);
+                        formId, source.getCorrectUsername());
             }
         }
         return Result.handled();
     }
 
     @Override
-    public Result handleServerCall(byte[] data, UUID playerUuid, String playerUsername) {
-        callResponseConsumer(data);
+    public Result handleServerCall(byte[] data, FloodgatePlayer source) {
+        if (!callResponseConsumer(source, data)) {
+            logger.error("Couldn't find stored form for player {}", source.getCorrectUsername());
+        }
         return Result.handled();
     }
 
-    public boolean closeForm(UUID player) {
-        return pluginMessageUtils.sendMessage(player, getIdentifier(), new byte[0]);
+    public boolean closeForm(FloodgatePlayer player) {
+        closeForms0(player);
+        return pluginMessageUtils.sendMessage(player.getCorrectUniqueId(), getIdentifier(), new byte[0]);
     }
 
-    public boolean sendForm(UUID player, Form form) {
-        byte[] formData = createFormData(form);
-        return pluginMessageUtils.sendMessage(player, getIdentifier(), formData);
+    private void closeForms0(FloodgatePlayer player) {
+        Map<Short, Form> forms = player.removeProperty(PROPERTY_ACTIVE_FORMS);
+        if (forms != null && !forms.isEmpty()) {
+            for (Form form : forms.values()) {
+                try {
+                    formDefinitions.definitionFor(form).handleFormResponse(form, "");
+                } catch (Exception e) {
+                    logger.error("Error while closing form!", e);
+                }
+            }
+        }
     }
 
-    public byte[] createFormData(Form form) {
-        short formId = getNextFormId();
+    public boolean sendForm(FloodgatePlayer player, Form form) {
+        byte[] formData = createFormData(player, form);
+        return pluginMessageUtils.sendMessage(player.getCorrectUniqueId(), getIdentifier(), formData);
+    }
+
+    public byte[] createFormData(FloodgatePlayer player, Form form) {
+        short formId = getNextFormId(player);
         if (config.isProxy()) {
-            formId |= 0x8000;
+            formId |= (short) 0x8000;
         }
-        storedForms.put(formId, form);
+
+        ((FloodgatePlayerImpl) player)
+                .getOrAddProperty(PROPERTY_ACTIVE_FORMS, ConcurrentHashMap::new)
+                .put(formId, form);
 
         FormDefinition<Form, ?, ?> definition = formDefinitions.definitionFor(form);
 
@@ -124,8 +143,15 @@ public byte[] createFormData(Form form) {
         return data;
     }
 
-    protected boolean callResponseConsumer(byte[] data) {
-        Form storedForm = storedForms.remove(getFormId(data));
+    protected boolean callResponseConsumer(FloodgatePlayer player, byte[] data) {
+        short formId = getFormId(data);
+
+        Map<Short, Form> forms = player.getProperty(PROPERTY_ACTIVE_FORMS);
+        if (forms == null) {
+            return false;
+        }
+
+        Form storedForm = forms.remove(formId);
         if (storedForm != null) {
             String responseData = new String(data, 2, data.length - 2, Charsets.UTF_8);
             try {
@@ -139,11 +165,18 @@ protected boolean callResponseConsumer(byte[] data) {
         return false;
     }
 
-    protected short getFormId(byte[] data) {
+    public void disconnect(FloodgatePlayer player) {
+        closeForms0(player);
+    }
+
+    private short getFormId(byte[] data) {
         return (short) ((data[0] & 0xFF) << 8 | data[1] & 0xFF);
     }
 
-    protected short getNextFormId() {
+    private short getNextFormId(FloodgatePlayer player) {
+        AtomicInteger nextFormId =
+                ((FloodgatePlayerImpl) player).getOrAddProperty(PROPERTY_LAST_FORM_ID, AtomicInteger::new);
+
         // signed bit is used to check if the form is from a proxy or a server
         return (short) nextFormId.getAndUpdate(
                 (number) -> number == Short.MAX_VALUE ? 0 : number + 1);