Revamp flag system

example/src/main/java/org/geysermc/mcprotocollib/protocol/example/MinecraftProtocolTest.java

@@ -63,7 +63,7 @@ public static void main(String[] args) {
             Server server = new TcpServer(HOST, PORT, MinecraftProtocol::new);
             server.setGlobalFlag(MinecraftConstants.SESSION_SERVICE_KEY, sessionService);
             server.setGlobalFlag(MinecraftConstants.VERIFY_USERS_KEY, VERIFY_USERS);
-            server.setGlobalFlag(MinecraftConstants.SERVER_INFO_BUILDER_KEY, (ServerInfoBuilder) session ->
+            server.setGlobalFlag(MinecraftConstants.SERVER_INFO_BUILDER_KEY, session ->
                     new ServerStatusInfo(
                             new VersionInfo(MinecraftCodec.CODEC.getMinecraftVersion(), MinecraftCodec.CODEC.getProtocolVersion()),
                             new PlayerInfo(100, 0, new ArrayList<>()),
@@ -73,7 +73,7 @@ public static void main(String[] args) {
                     )
             );
 
-            server.setGlobalFlag(MinecraftConstants.SERVER_LOGIN_HANDLER_KEY, (ServerLoginHandler) session ->
+            server.setGlobalFlag(MinecraftConstants.SERVER_LOGIN_HANDLER_KEY, session ->
                     session.send(new ClientboundLoginPacket(
                             0,
                             false,
@@ -152,7 +152,7 @@ private static void status() {
         MinecraftProtocol protocol = new MinecraftProtocol();
         Session client = new TcpClientSession(HOST, PORT, protocol, PROXY);
         client.setFlag(MinecraftConstants.SESSION_SERVICE_KEY, sessionService);
-        client.setFlag(MinecraftConstants.SERVER_INFO_HANDLER_KEY, (ServerInfoHandler) (session, info) -> {
+        client.setFlag(MinecraftConstants.SERVER_INFO_HANDLER_KEY, (session, info) -> {
             System.out.println("Version: " + info.getVersionInfo().getVersionName()
                     + ", " + info.getVersionInfo().getProtocolVersion());
             System.out.println("Player Count: " + info.getPlayerInfo().getOnlinePlayers()
@@ -162,7 +162,7 @@ private static void status() {
             System.out.println("Icon: " + new String(Base64.getEncoder().encode(info.getIconPng()), StandardCharsets.UTF_8));
         });
 
-        client.setFlag(MinecraftConstants.SERVER_PING_TIME_HANDLER_KEY, (ServerPingTimeHandler) (session, pingTime) ->
+        client.setFlag(MinecraftConstants.SERVER_PING_TIME_HANDLER_KEY, (session, pingTime) ->
                 System.out.println("Server ping took " + pingTime + "ms"));
 
         client.connect();

protocol/src/main/java/org/geysermc/mcprotocollib/network/AbstractServer.java

@@ -47,33 +47,32 @@ public Map<String, Object> getGlobalFlags() {
     }
 
     @Override
-    public boolean hasGlobalFlag(String key) {
-        return this.flags.containsKey(key);
+    public boolean hasGlobalFlag(Flag<?> flag) {
+        return this.flags.containsKey(flag.key());
     }
 
     @Override
-    public <T> T getGlobalFlag(String key) {
-        return this.getGlobalFlag(key, null);
+    public <T> T getGlobalFlag(Flag<T> flag) {
+        return this.getGlobalFlag(flag, null);
     }
 
-    @SuppressWarnings("unchecked")
     @Override
-    public <T> T getGlobalFlag(String key, T def) {
-        Object value = this.flags.get(key);
+    public <T> T getGlobalFlag(Flag<T> flag, T def) {
+        Object value = this.flags.get(flag.key());
         if(value == null) {
             return def;
         }
 
         try {
-            return (T) value;
+            return flag.cast(value);
         } catch(ClassCastException e) {
-            throw new IllegalStateException("Tried to get flag \"" + key + "\" as the wrong type. Actual type: " + value.getClass().getName());
+            throw new IllegalStateException("Tried to get flag \"" + flag.key() + "\" as the wrong type. Actual type: " + value.getClass().getName());
         }
     }
 
     @Override
-    public void setGlobalFlag(String key, Object value) {
-        this.flags.put(key, value);
+    public <T> void setGlobalFlag(Flag<T> flag, T value) {
+        this.flags.put(flag.key(), value);
     }
 
     @Override

protocol/src/main/java/org/geysermc/mcprotocollib/network/BuiltinFlags.java

@@ -1,22 +1,24 @@
 package org.geysermc.mcprotocollib.network;
 
+import java.net.InetSocketAddress;
+
 /**
  * Built-in PacketLib session flags.
  */
 public class BuiltinFlags {
     /**
      * When set to true, enables printing internal debug messages.
      */
-    public static final String PRINT_DEBUG = "print-packetlib-debug";
+    public static final Flag<Boolean> PRINT_DEBUG = new Flag<>("print-packetlib-debug", Boolean.class);
 
-    public static final String ENABLE_CLIENT_PROXY_PROTOCOL = "enable-client-proxy-protocol";
+    public static final Flag<Boolean> ENABLE_CLIENT_PROXY_PROTOCOL = new Flag<>("enable-client-proxy-protocol", Boolean.class);
 
-    public static final String CLIENT_PROXIED_ADDRESS = "client-proxied-address";
+    public static final Flag<InetSocketAddress> CLIENT_PROXIED_ADDRESS = new Flag<>("client-proxied-address", InetSocketAddress.class);
 
     /**
      * When set to false, an SRV record resolve is not attempted.
      */
-    public static final String ATTEMPT_SRV_RESOLVE = "attempt-srv-resolve";
+    public static final Flag<Boolean> ATTEMPT_SRV_RESOLVE = new Flag<>("attempt-srv-resolve", Boolean.class);
 
     private BuiltinFlags() {
     }

protocol/src/main/java/org/geysermc/mcprotocollib/network/Flag.java

@@ -0,0 +1,17 @@
+package org.geysermc.mcprotocollib.network;
+
+import lombok.RequiredArgsConstructor;
+
+@RequiredArgsConstructor
+public final class Flag<T> {
+    private final String key;
+    private final Class<T> type;
+
+    public T cast(Object obj) {
+        return type.cast(obj);
+    }
+
+    public String key() {
+        return key;
+    }
+}

protocol/src/main/java/org/geysermc/mcprotocollib/network/Server.java

@@ -49,41 +49,42 @@ public interface Server {
     /**
      * Checks whether this server has a flag set.
      *
-     * @param key Key of the flag to check for.
+     * @param flag Flag to check for.
      * @return Whether this server has a flag set.
      */
-    boolean hasGlobalFlag(String key);
+    boolean hasGlobalFlag(Flag<?> flag);
 
     /**
      * Gets the value of the given flag as an instance of the given type.
      *
      * @param <T> Type of the flag.
-     * @param key Key of the flag.
+     * @param flag Flag to check for.
      * @return Value of the flag.
      * @throws IllegalStateException If the flag's value isn't of the required type.
      */
-    <T> T getGlobalFlag(String key);
+    <T> T getGlobalFlag(Flag<T> flag);
 
     /**
      * Gets the value of the given flag as an instance of the given type.
      * If the flag is not set, the specified default value will be returned.
      *
      * @param <T> Type of the flag.
-     * @param key Key of the flag.
+     * @param flag Flag to check for.
      * @param def Default value of the flag.
      * @return Value of the flag.
      * @throws IllegalStateException If the flag's value isn't of the required type.
      */
-    <T> T getGlobalFlag(String key, T def);
+    <T> T getGlobalFlag(Flag<T> flag, T def);
 
     /**
      * Sets the value of a flag. The flag will be used in sessions if a session does
      * not contain a value for the flag.
      *
-     * @param key   Key of the flag.
+     * @param <T> Type of the flag.
+     * @param flag Flag to check for.
      * @param value Value to set the flag to.
      */
-    void setGlobalFlag(String key, Object value);
+    <T> void setGlobalFlag(Flag<T> flag, T value);
 
     /**
      * Gets the listeners listening on this session.

protocol/src/main/java/org/geysermc/mcprotocollib/network/Session.java

@@ -84,44 +84,45 @@ public interface Session {
      * Checks whether this session has a flag set. If this session belongs to a server,
      * the server's flags will also be checked.
      *
-     * @param key Key of the flag to check for.
+     * @param flag Flag to check for.
      * @return Whether this session has a flag set.
      */
-    boolean hasFlag(String key);
+    boolean hasFlag(Flag<?> flag);
 
     /**
      * Gets the value of the given flag as an instance of the given type. If this
      * session belongs to a server, the server's flags will be checked for the flag
      * as well.
      *
      * @param <T> Type of the flag.
-     * @param key Key of the flag.
+     * @param flag Flag to check for.
      * @return Value of the flag.
      * @throws IllegalStateException If the flag's value isn't of the required type.
      */
-    <T> T getFlag(String key);
+    <T> T getFlag(Flag<T> flag);
 
     /**
      * Gets the value of the given flag as an instance of the given type. If this
      * session belongs to a server, the server's flags will be checked for the flag
      * as well. If the flag is not set, the specified default value will be returned.
      *
      * @param <T> Type of the flag.
-     * @param key Key of the flag.
+     * @param flag Flag to check for.
      * @param def Default value of the flag.
      * @return Value of the flag.
      * @throws IllegalStateException If the flag's value isn't of the required type.
      */
-    <T> T getFlag(String key, T def);
+    <T> T getFlag(Flag<T> flag, T def);
 
     /**
      * Sets the value of a flag. This does not change a server's flags if this session
      * belongs to a server.
      *
-     * @param key   Key of the flag.
+     * @param <T> Type of the flag.
+     * @param flag Flag to check for.
      * @param value Value to set the flag to.
      */
-    void setFlag(String key, Object value);
+    <T> void setFlag(Flag<T> flag, T value);
 
     /**
      * Gets the listeners listening on this session.

protocol/src/main/java/org/geysermc/mcprotocollib/network/tcp/TcpServerSession.java

@@ -3,6 +3,7 @@
 import org.geysermc.mcprotocollib.network.codec.PacketCodecHelper;
 import org.geysermc.mcprotocollib.network.packet.PacketProtocol;
 import io.netty.channel.ChannelHandlerContext;
+import org.geysermc.mcprotocollib.network.Flag;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -31,22 +32,22 @@ public Map<String, Object> getFlags() {
     }
 
     @Override
-    public boolean hasFlag(String key) {
-        if(super.hasFlag(key)) {
+    public boolean hasFlag(Flag<?> flag) {
+        if(super.hasFlag(flag)) {
             return true;
         }
 
-        return this.server.hasGlobalFlag(key);
+        return this.server.hasGlobalFlag(flag);
     }
 
     @Override
-    public <T> T getFlag(String key, T def) {
-        T ret = super.getFlag(key, null);
+    public <T> T getFlag(Flag<T> flag, T def) {
+        T ret = super.getFlag(flag, null);
         if(ret != null) {
             return ret;
         }
 
-        return this.server.getGlobalFlag(key, def);
+        return this.server.getGlobalFlag(flag, def);
     }
 
     @Override

protocol/src/main/java/org/geysermc/mcprotocollib/network/tcp/TcpSession.java

@@ -18,6 +18,7 @@
 import io.netty.util.concurrent.DefaultThreadFactory;
 import net.kyori.adventure.text.Component;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.geysermc.mcprotocollib.network.Flag;
 
 import java.net.ConnectException;
 import java.net.SocketAddress;
@@ -99,33 +100,32 @@ public Map<String, Object> getFlags() {
     }
 
     @Override
-    public boolean hasFlag(String key) {
-        return this.flags.containsKey(key);
+    public boolean hasFlag(Flag<?> flag) {
+        return this.flags.containsKey(flag.key());
     }
 
     @Override
-    public <T> T getFlag(String key) {
-        return this.getFlag(key, null);
+    public <T> T getFlag(Flag<T> flag) {
+        return this.getFlag(flag, null);
     }
 
-    @SuppressWarnings("unchecked")
     @Override
-    public <T> T getFlag(String key, T def) {
-        Object value = this.flags.get(key);
+    public <T> T getFlag(Flag<T> flag, T def) {
+        Object value = this.flags.get(flag.key());
         if (value == null) {
             return def;
         }
 
         try {
-            return (T) value;
+            return flag.cast(value);
         } catch (ClassCastException e) {
-            throw new IllegalStateException("Tried to get flag \"" + key + "\" as the wrong type. Actual type: " + value.getClass().getName());
+            throw new IllegalStateException("Tried to get flag \"" + flag.key() + "\" as the wrong type. Actual type: " + value.getClass().getName());
         }
     }
 
     @Override
-    public void setFlag(String key, Object value) {
-        this.flags.put(key, value);
+    public <T> void setFlag(Flag<T> flag, T value) {
+        this.flags.put(flag.key(), value);
     }
 
     @Override