[1.21.1] Screen Handler Message Handshake

src/main/java/io/wispforest/owo/client/screens/ScreenInternals.java

@@ -1,8 +1,11 @@
 package io.wispforest.owo.client.screens;
 
+import io.wispforest.endec.StructEndec;
 import io.wispforest.endec.impl.StructEndecBuilder;
 import io.wispforest.owo.Owo;
 import io.wispforest.endec.Endec;
+import io.wispforest.owo.network.OwoHandshake;
+import io.wispforest.owo.ops.TextOps;
 import io.wispforest.owo.serialization.CodecUtils;
 import io.wispforest.owo.serialization.endec.MinecraftEndecs;
 import io.wispforest.owo.util.pond.OwoScreenHandlerExtension;
@@ -11,16 +14,27 @@
 import net.fabricmc.fabric.api.client.networking.v1.ClientPlayNetworking;
 import net.fabricmc.fabric.api.client.screen.v1.ScreenEvents;
 import net.fabricmc.fabric.api.networking.v1.PayloadTypeRegistry;
+import net.fabricmc.fabric.api.networking.v1.ServerConfigurationNetworking;
 import net.fabricmc.fabric.api.networking.v1.ServerPlayNetworking;
 import net.minecraft.client.gui.screen.ingame.ScreenHandlerProvider;
 import net.minecraft.network.PacketByteBuf;
 import net.minecraft.network.packet.CustomPayload;
+import net.minecraft.registry.Registries;
+import net.minecraft.screen.ScreenHandler;
+import net.minecraft.screen.ScreenHandlerType;
+import net.minecraft.server.network.ServerPlayerEntity;
+import net.minecraft.text.Text;
+import net.minecraft.text.Texts;
 import net.minecraft.util.Identifier;
 import org.jetbrains.annotations.ApiStatus;
 
+import java.util.*;
+
 @ApiStatus.Internal
 public class ScreenInternals {
     public static final Identifier SYNC_PROPERTIES = Identifier.of("owo", "sync_screen_handler_properties");
+    public static final Identifier HANDSHAKE_REQUEST = Identifier.of("owo", "request_screen_handler_messages");
+    public static final Identifier HANDSHAKE_RESPONSE = Identifier.of("owo", "response_screen_handler_messages");
 
     public static void init() {
         var localPacketCodec = CodecUtils.toPacketCodec(LocalPacket.ENDEC);
@@ -29,6 +43,9 @@ public static void init() {
         PayloadTypeRegistry.playC2S().register(LocalPacket.ID, localPacketCodec);
         PayloadTypeRegistry.playS2C().register(SyncPropertiesPacket.ID, CodecUtils.toPacketCodec(SyncPropertiesPacket.ENDEC));
 
+        PayloadTypeRegistry.playS2C().register(HandshakeRequest.ID, CodecUtils.toPacketCodec(HandshakeRequest.ENDEC));
+        PayloadTypeRegistry.playC2S().register(HandshakeResponse.ID, CodecUtils.toPacketCodec(HandshakeResponse.ENDEC));
+
         ServerPlayNetworking.registerGlobalReceiver(LocalPacket.ID, (payload, context) -> {
             var screenHandler = context.player().currentScreenHandler;
 
@@ -39,6 +56,22 @@ public static void init() {
 
             ((OwoScreenHandlerExtension) screenHandler).owo$handlePacket(payload, false);
         });
+
+        ServerPlayNetworking.registerGlobalReceiver(HandshakeResponse.ID, (payload, context) -> {
+            var screenHandler = context.player().currentScreenHandler;
+
+            if (screenHandler == null) {
+                Owo.LOGGER.error("[ScreenHandlerHandshake] Received handshake response for null ScreenHandler");
+                return;
+            }
+
+            if (!payload.type().equals(screenHandler.getType())) {
+                Owo.LOGGER.error("[ScreenHandlerHandshake] Received handshake response packet for different ScreenHandler type: [Expected Type: {}, Current Type: {}]", payload.type(), screenHandler.getType());
+                return;
+            }
+
+            ((OwoScreenHandlerExtension) screenHandler).owo$verifyData(context.player(), payload.messageNames());
+        });
     }
 
     public record LocalPacket(int packetId, PacketByteBuf payload) implements CustomPayload {
@@ -68,6 +101,48 @@ public Id<? extends CustomPayload> getId() {
         }
     }
 
+    public static void attemptHandshake(ScreenHandlerType<?> type, ServerPlayerEntity player) {
+        if (type == null) return;
+
+        if (ServerPlayNetworking.canSend(player, OwoHandshake.OFF_CHANNEL_ID)) {
+            Owo.LOGGER.info("[ScreenHandlerHandshake] Handshake disabled by client, skipping");
+            return;
+        }
+
+        try {
+            ServerPlayNetworking.send(player, new HandshakeRequest(type));
+        } catch (Exception e) {
+            Owo.LOGGER.error("[ScreenHandlerHandshake] Unable to Handshake check handler as getting the type encountered an error: ", e);
+        }
+    }
+
+    private record HandshakeRequest(ScreenHandlerType<?> type) implements CustomPayload {
+        public static final Id<HandshakeRequest> ID = new Id<>(HANDSHAKE_REQUEST);
+        public static final Endec<HandshakeRequest> ENDEC = StructEndecBuilder.of(
+            MinecraftEndecs.ofRegistry(Registries.SCREEN_HANDLER).fieldOf("type", HandshakeRequest::type),
+            HandshakeRequest::new
+        );
+
+        @Override
+        public Id<? extends CustomPayload> getId() {
+            return ID;
+        }
+    }
+
+    private record HandshakeResponse(ScreenHandlerType<?> type, LinkedHashSet<String> messageNames) implements CustomPayload {
+        public static final CustomPayload.Id<HandshakeResponse> ID = new CustomPayload.Id<>(HANDSHAKE_RESPONSE);
+        public static final Endec<HandshakeResponse> ENDEC = StructEndecBuilder.of(
+            MinecraftEndecs.ofRegistry(Registries.SCREEN_HANDLER).fieldOf("type", HandshakeResponse::type),
+            Endec.STRING.listOf().xmap(LinkedHashSet::new, ArrayList::new).fieldOf("message_names", HandshakeResponse::messageNames),
+            HandshakeResponse::new
+        );
+
+        @Override
+        public CustomPayload.Id<? extends CustomPayload> getId() {
+            return ID;
+        }
+    }
+
     @Environment(EnvType.CLIENT)
     public static class Client {
         public static void init() {
@@ -97,6 +172,22 @@ public static void init() {
 
                 ((OwoScreenHandlerExtension) screenHandler).owo$readPropertySync(payload);
             });
+
+            ClientPlayNetworking.registerGlobalReceiver(HandshakeRequest.ID, (payload, context) -> {
+                var screenHandler = context.player().currentScreenHandler;
+
+                if (screenHandler == null) {
+                    Owo.LOGGER.error("[ScreenHandlerHandshake] Received handshake request packet for null ScreenHandler");
+                    return;
+                }
+
+                if (!payload.type().equals(screenHandler.getType())) {
+                    Owo.LOGGER.error("[ScreenHandlerHandshake] Received handshake request packet for different ScreenHandler type: [Expected Type: {}, Current Type: {}]", payload.type(), screenHandler.getType());
+                    return;
+                }
+
+                context.responseSender().sendPacket(new HandshakeResponse(screenHandler.getType(), ((OwoScreenHandlerExtension) screenHandler).owo$gatherMessageNames()));
+            });
         }
     }
 }

src/main/java/io/wispforest/owo/client/screens/ScreenhandlerMessageData.java

@@ -6,4 +6,8 @@
 import java.util.function.Consumer;
 
 @ApiStatus.Internal
-public record ScreenhandlerMessageData<T>(int id, boolean clientbound, Endec<T> endec, Consumer<T> handler) {}
+public record ScreenhandlerMessageData<T>(int id, Class<T> messageClass,  boolean clientbound, Endec<T> endec, Consumer<T> handler) {
+    public String messageName() {
+        return messageClass.getSimpleName();
+    }
+}

src/main/java/io/wispforest/owo/mixin/ScreenHandlerMixin.java

@@ -2,12 +2,15 @@
 
 import io.wispforest.endec.SerializationContext;
 import io.wispforest.endec.impl.ReflectiveEndecBuilder;
+import io.wispforest.owo.Owo;
 import io.wispforest.owo.client.screens.OwoScreenHandler;
 import io.wispforest.owo.client.screens.ScreenInternals;
 import io.wispforest.owo.client.screens.ScreenhandlerMessageData;
 import io.wispforest.owo.client.screens.SyncedProperty;
 import io.wispforest.owo.network.NetworkException;
 import io.wispforest.endec.Endec;
+import io.wispforest.owo.network.OwoHandshake;
+import io.wispforest.owo.network.OwoNetChannel;
 import io.wispforest.owo.serialization.RegistriesAttribute;
 import io.wispforest.owo.serialization.endec.MinecraftEndecs;
 import io.wispforest.owo.util.pond.OwoScreenHandlerExtension;
@@ -19,27 +22,37 @@
 import net.minecraft.entity.player.PlayerEntity;
 import net.minecraft.network.packet.CustomPayload;
 import net.minecraft.screen.ScreenHandler;
+import net.minecraft.screen.ScreenHandlerSyncHandler;
 import net.minecraft.screen.ScreenHandlerType;
 import net.minecraft.server.network.ServerPlayerEntity;
+import net.minecraft.text.Text;
+import net.minecraft.text.Texts;
+import org.apache.commons.lang3.stream.Streams;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.Unique;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
-import java.util.ArrayList;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
 
 @Mixin(ScreenHandler.class)
 public abstract class ScreenHandlerMixin implements OwoScreenHandler, OwoScreenHandlerExtension {
 
     @Shadow private boolean disableSync;
 
+    @Shadow
+    public abstract ScreenHandlerType<?> getType();
+
+    @Shadow
+    @Final
+    private @Nullable ScreenHandlerType<?> type;
     private final List<SyncedProperty<?>> owo$properties = new ArrayList<>();
 
     private final Map<Class<?>, ScreenhandlerMessageData<?>> owo$messages = new LinkedHashMap<>();
@@ -75,7 +88,7 @@ public PlayerEntity player() {
     public <R extends Record> void addServerboundMessage(Class<R> messageClass, Endec<R> endec, Consumer<R> handler) {
         int id = this.owo$serverboundMessages.size();
 
-        var messageData = new ScreenhandlerMessageData<>(id, false, endec, handler);
+        var messageData = new ScreenhandlerMessageData<>(id, messageClass, false, endec, handler);
         this.owo$serverboundMessages.add(messageData);
 
         if (this.owo$messages.put(messageClass, messageData) != null) {
@@ -87,7 +100,7 @@ public <R extends Record> void addServerboundMessage(Class<R> messageClass, Ende
     public <R extends Record> void addClientboundMessage(Class<R> messageClass, Endec<R> endec, Consumer<R> handler) {
         int id = this.owo$clientboundMessages.size();
 
-        var messageData = new ScreenhandlerMessageData<>(id, true, endec, handler);
+        var messageData = new ScreenhandlerMessageData<>(id, messageClass, true, endec, handler);
         this.owo$clientboundMessages.add(messageData);
 
         if (this.owo$messages.put(messageClass, messageData) != null) {
@@ -116,13 +129,13 @@ public <R extends Record> void sendMessage(@NotNull R message) {
 
         if (messageData.clientbound()) {
             if (!(this.owo$player instanceof ServerPlayerEntity serverPlayer)) {
-                throw new NetworkException("Tried to send clientbound message on the server");
+                throw new NetworkException("Tried to send clientbound message on the server: [Type: " + message.getClass().getSimpleName() + "]");
             }
 
             ServerPlayNetworking.send(serverPlayer, packet);
         } else {
             if (!this.owo$player.getWorld().isClient) {
-                throw new NetworkException("Tried to send serverbound message on the client");
+                throw new NetworkException("Tried to send serverbound message on the client: [Type: " + message.getClass().getSimpleName() + "]");
             }
 
             this.owo$sendToServer(packet);
@@ -138,12 +151,45 @@ public <R extends Record> void sendMessage(@NotNull R message) {
     @Override
     @SuppressWarnings({"rawtypes", "unchecked"})
     public void owo$handlePacket(ScreenInternals.LocalPacket packet, boolean clientbound) {
-        ScreenhandlerMessageData messageData = (clientbound ? this.owo$clientboundMessages : this.owo$serverboundMessages).get(packet.packetId());
+        var messages = (clientbound ? this.owo$clientboundMessages : this.owo$serverboundMessages);
+
+        if (packet.packetId() < 0 || packet.packetId() >= messages.size()) {
+            throw new NetworkException("Unable to handle packet as it was not properly registered on the [" + (clientbound ?  "CLIENT" : "SERVER") + "]");
+        }
+
+        ScreenhandlerMessageData messageData = messages.get(packet.packetId());
         var ctx = SerializationContext.attributes(RegistriesAttribute.of(this.owo$player.getRegistryManager()));
 
         messageData.handler().accept(packet.payload().read(ctx, messageData.endec()));
     }
 
+    @Inject(method = "updateSyncHandler", at = @At("HEAD"))
+    private void compareHandlersNetworking(ScreenHandlerSyncHandler handler, CallbackInfo ci) {
+        if (!(player() instanceof ServerPlayerEntity serverPlayer)) return;
+
+        ScreenInternals.attemptHandshake(this.type, serverPlayer);
+    }
+
+    @Override
+    public LinkedHashSet<String> owo$gatherMessageNames() {
+        return Streams.of(this.owo$clientboundMessages, this.owo$serverboundMessages)
+            .flatMap(Collection::stream)
+            .map(ScreenhandlerMessageData::messageName)
+            .collect(Collectors.toCollection(LinkedHashSet::new));
+    }
+
+    @Override
+    public void owo$verifyData(ServerPlayerEntity player, Set<String> clientMessageNames) {
+        var errorMessage = new StringBuilder();
+
+        if (OwoHandshake.checkForMismatchStrIds("screen_handler_messages", clientMessageNames, owo$gatherMessageNames(), errorMessage)) return;
+
+        player.closeHandledScreen();
+
+        player.sendMessage(Texts.join(List.of(Owo.PREFIX, Text.of("Unable to open screen as there was a message mismatch:")), Text.empty()));
+        player.sendMessage(Text.of(errorMessage.toString()));
+    }
+
     @Override
     public <T> SyncedProperty<T> createProperty(Class<T> clazz, Endec<T> endec, T initial) {
         var prop = new SyncedProperty<>(this.owo$properties.size(), endec, initial, (ScreenHandler)(Object) this);