Change handshake intent from transfer when connecting first

ViaVersionAddons · Jan 5, 2025 · 862738d · 862738d
1 parent 478a1b8
commit 862738d
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 4 deletions.
diff --git a/src/main/java/net/lenni0451/miniconnect/Main.java b/src/main/java/net/lenni0451/miniconnect/Main.java
@@ -45,6 +45,7 @@ public void onEnable() {
         ViaProxy.EVENT_MANAGER.register(new HAProxyEnableHandler());
         ViaProxy.EVENT_MANAGER.register(new TargetOnlineModeHandler());
         ViaProxy.EVENT_MANAGER.register(new ProxyOnlineModeHandler());
+        ViaProxy.EVENT_MANAGER.register(new HandshakeIntentListener());
     }
 
 }
diff --git a/src/main/java/net/lenni0451/miniconnect/proxy/StateRegistry.java b/src/main/java/net/lenni0451/miniconnect/proxy/StateRegistry.java
@@ -16,6 +16,7 @@ public class StateRegistry {
     private final Map<InetAddress, ConnectionInfo> reconnectTargets = CacheBuilder.newBuilder().expireAfterWrite(1, TimeUnit.MINUTES).<InetAddress, ConnectionInfo>build().asMap();
     private final Map<InetAddress, ConnectionInfo> lobbyTargets = CacheBuilder.newBuilder().expireAfterWrite(1, TimeUnit.MINUTES).<InetAddress, ConnectionInfo>build().asMap();
     private final Set<UUID> verificationQueue = Collections.newSetFromMap(CacheBuilder.newBuilder().expireAfterWrite(1, TimeUnit.MINUTES).<UUID, Boolean>build().asMap());
+    private final Set<InetAddress> changeHandshakeIntent = Collections.newSetFromMap(CacheBuilder.newBuilder().expireAfterWrite(1, TimeUnit.MINUTES).<InetAddress, Boolean>build().asMap());
 
     public Map<InetAddress, ConnectionInfo> getConnectionTargets() {
         return this.connectionTargets;
@@ -33,4 +34,8 @@ public Set<UUID> getVerificationQueue() {
         return this.verificationQueue;
     }
 
+    public Set<InetAddress> getChangeHandshakeIntent() {
+        return this.changeHandshakeIntent;
+    }
+
 }
diff --git a/src/main/java/net/lenni0451/miniconnect/proxy/event/HandshakeIntentListener.java b/src/main/java/net/lenni0451/miniconnect/proxy/event/HandshakeIntentListener.java
@@ -0,0 +1,35 @@
+package net.lenni0451.miniconnect.proxy.event;
+
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.SimpleChannelInboundHandler;
+import net.lenni0451.lambdaevents.EventHandler;
+import net.lenni0451.miniconnect.Main;
+import net.lenni0451.miniconnect.utils.ChannelUtils;
+import net.raphimc.netminecraft.constants.IntendedState;
+import net.raphimc.netminecraft.constants.MCPipeline;
+import net.raphimc.netminecraft.packet.Packet;
+import net.raphimc.netminecraft.packet.impl.handshaking.C2SHandshakingClientIntentionPacket;
+import net.raphimc.viaproxy.plugins.events.Client2ProxyChannelInitializeEvent;
+import net.raphimc.viaproxy.plugins.events.types.ITyped;
+
+public class HandshakeIntentListener {
+
+    @EventHandler
+    public void onClient2ProxyChannelInitialize(final Client2ProxyChannelInitializeEvent event) {
+        if (!event.getType().equals(ITyped.Type.POST)) return;
+        event.getChannel().pipeline().addBefore(MCPipeline.HANDLER_HANDLER_NAME, "miniconnect-intent-changer", new SimpleChannelInboundHandler<Packet>() {
+            @Override
+            protected void channelRead0(ChannelHandlerContext channelHandlerContext, Packet packet) {
+                if (packet instanceof C2SHandshakingClientIntentionPacket clientIntention) {
+                    if (clientIntention.intendedState.equals(IntendedState.TRANSFER)) {
+                        if (Main.getInstance().getStateRegistry().getChangeHandshakeIntent().remove(ChannelUtils.getChannelAddress(event.getChannel()))) {
+                            clientIntention.intendedState = IntendedState.LOGIN;
+                        }
+                    }
+                }
+                channelHandlerContext.fireChannelRead(packet);
+            }
+        });
+    }
+
+}
diff --git a/src/main/java/net/lenni0451/miniconnect/server/states/play/screen/impl/MainScreen.java b/src/main/java/net/lenni0451/miniconnect/server/states/play/screen/impl/MainScreen.java
@@ -144,10 +144,9 @@ public void init(ScreenHandler screenHandler, ItemList itemList) {
                         Logger.LOGGER.error("Failed to save player config", e);
                     }
                 }
-                Main.getInstance().getStateRegistry().getConnectionTargets().put(
-                        ChannelUtils.getChannelAddress(screenHandler.getStateHandler().getChannel()),
-                        playerConfig.toConnectionInfo()
-                );
+                InetAddress channelAddress = ChannelUtils.getChannelAddress(screenHandler.getStateHandler().getChannel());
+                Main.getInstance().getStateRegistry().getConnectionTargets().put(channelAddress, playerConfig.toConnectionInfo());
+                Main.getInstance().getStateRegistry().getChangeHandshakeIntent().add(channelAddress);
                 screenHandler.getStateHandler().send(new S2CTransferPacket(playerConfig.handshakeAddress, playerConfig.handshakePort));
             } else {
                 screenHandler.getStateHandler().send(new S2CSystemChatPacket(new StringComponent("§cYou need to set all options before connecting"), false));