fabric-networking: fix outstanding TODOs

fabric-networking/src/main/java/net/fabricmc/fabric/impl/network/ClientSidePacketRegistryImpl.java

@@ -56,7 +56,7 @@ public class ClientSidePacketRegistryImpl extends PacketRegistryImpl implements
 				handler.getClientConnection().send(packet, completionListener);
 			}
 		} else {
-			// TODO: log warning
+			LOGGER.warn("Sending packet " + packet + " to server failed, not connected!");
 		}
 	}
 

fabric-networking/src/main/java/net/fabricmc/fabric/impl/network/PacketRegistryImpl.java

@@ -47,7 +47,8 @@ public abstract class PacketRegistryImpl implements PacketRegistry {
 	public void register(Identifier id, PacketConsumer consumer) {
 		boolean isNew = true;
 		if (consumerMap.containsKey(id)) {
-			// TODO: log warning
+			LOGGER.warn("Registered duplicate packet " + id + "!");
+			LOGGER.trace(new Throwable());
 			isNew = false;
 		}
 
@@ -59,8 +60,12 @@ public abstract class PacketRegistryImpl implements PacketRegistry {
 
 	@Override
 	public void unregister(Identifier id) {
-		consumerMap.remove(id);
-		onUnregister(id);
+		if (consumerMap.remove(id) != null) {
+			onUnregister(id);
+		} else {
+			LOGGER.warn("Tried to unregister non-registered packet " + id + "!");
+			LOGGER.trace(new Throwable());
+		}
 	}
 
 	protected abstract void onRegister(Identifier id);

fabric-networking/src/main/java/net/fabricmc/fabric/impl/network/ServerSidePacketRegistryImpl.java

@@ -24,22 +24,41 @@ import net.fabricmc.fabric.api.network.ServerSidePacketRegistry;
 import net.minecraft.client.network.packet.CustomPayloadS2CPacket;
 import net.minecraft.entity.player.PlayerEntity;
 import net.minecraft.network.Packet;
+import net.minecraft.server.MinecraftServer;
+import net.minecraft.server.network.ServerPlayNetworkHandler;
 import net.minecraft.server.network.ServerPlayerEntity;
 import net.minecraft.server.network.packet.CustomPayloadC2SPacket;
 import net.minecraft.server.network.packet.LoginQueryResponseC2SPacket;
 import net.minecraft.util.Identifier;
 import net.minecraft.util.PacketByteBuf;
 
-import java.util.Collection;
-import java.util.HashSet;
-import java.util.WeakHashMap;
+import java.lang.ref.WeakReference;
+import java.util.*;
+import java.util.function.Consumer;
 
 public class ServerSidePacketRegistryImpl extends PacketRegistryImpl implements ServerSidePacketRegistry {
 	private final WeakHashMap<PlayerEntity, Collection<Identifier>> playerPayloadIds = new WeakHashMap<>();
+	private final Set<WeakReference<ServerPlayNetworkHandler>> handlers = new HashSet<>();
 
 	public void onQueryResponse(LoginQueryResponseC2SPacket packet) {
 	}
 
+	public void addNetworkHandler(ServerPlayNetworkHandler handler) {
+		handlers.add(new WeakReference<>(handler));
+	}
+
+	protected void forEachHandler(Consumer<ServerPlayNetworkHandler> consumer) {
+		Iterator<WeakReference<ServerPlayNetworkHandler>> it = handlers.iterator();
+		while (it.hasNext()) {
+			ServerPlayNetworkHandler server = it.next().get();
+			if (server != null) {
+				consumer.accept(server);
+			} else {
+				it.remove();
+			}
+		}
+	}
+
 	@Override
 	public boolean canPlayerReceive(PlayerEntity player, Identifier id) {
 		Collection<Identifier> ids = playerPayloadIds.get(player);
@@ -66,22 +85,14 @@ public class ServerSidePacketRegistryImpl extends PacketRegistryImpl implements
 
 	@Override
 	protected void onRegister(Identifier id) {
-		/* MinecraftServer server = FabricLoader.INSTANCE.getEnvironmentHandler().getServerInstance();
-		if (server != null) {
-			Packet<?> packet = createRegisterTypePacket(PacketTypes.REGISTER, Collections.singleton(id));
-			PlayerStream.all(server).forEach((p) -> sendToPlayer(p, packet));
-		} */
-		// TODO
+		createRegisterTypePacket(PacketTypes.REGISTER, Collections.singleton(id))
+			.ifPresent((packet) -> forEachHandler((n) -> n.sendPacket(packet)));
 	}
 
 	@Override
 	protected void onUnregister(Identifier id) {
-		/* MinecraftServer server = FabricLoader.INSTANCE.getEnvironmentHandler().getServerInstance();
-		if (server != null) {
-			Packet<?> packet = createRegisterTypePacket(PacketTypes.UNREGISTER, Collections.singleton(id));
-			PlayerStream.all(server).forEach((p) -> sendToPlayer(p, packet));
-		} */
-		// TODO
+		createRegisterTypePacket(PacketTypes.UNREGISTER, Collections.singleton(id))
+			.ifPresent((packet) -> forEachHandler((n) -> n.sendPacket(packet)));
 	}
 
 	@Override

fabric-networking/src/main/java/net/fabricmc/fabric/mixin/network/MixinMinecraftClient.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.mixin.network;
+
+import net.fabricmc.fabric.impl.network.ClientSidePacketRegistryImpl;
+import net.minecraft.client.MinecraftClient;
+import net.minecraft.client.gui.Screen;
+import org.apache.logging.log4j.Logger;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+@Mixin(MinecraftClient.class)
+public class MixinMinecraftClient {
+	@Inject(at = @At("RETURN"), method = "disconnect(Lnet/minecraft/client/gui/Screen;)V")
+	public void disconnectAfter(Screen screen_1, CallbackInfo info) {
+		ClientSidePacketRegistryImpl.invalidateRegisteredIdList();
+	}
+}

fabric-networking/src/main/java/net/fabricmc/fabric/mixin/network/MixinPlayerManager.java

@@ -16,8 +16,10 @@
 
 package net.fabricmc.fabric.mixin.network;
 
+import net.fabricmc.fabric.api.network.PacketRegistry;
 import net.fabricmc.fabric.api.network.ServerSidePacketRegistry;
 import net.fabricmc.fabric.impl.network.PacketRegistryImpl;
+import net.fabricmc.fabric.impl.network.ServerSidePacketRegistryImpl;
 import net.minecraft.network.ClientConnection;
 import net.minecraft.network.Packet;
 import net.minecraft.server.PlayerManager;
@@ -32,11 +34,12 @@ import java.util.Optional;
 @Mixin(priority = 500, value = PlayerManager.class)
 public abstract class MixinPlayerManager {
 	@Inject(method = "onPlayerConnect", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/network/packet/DifficultyS2CPacket;<init>(Lnet/minecraft/world/Difficulty;Z)V"))
-	public void onPlayerConnect(ClientConnection lvt1, ServerPlayerEntity lvt2, CallbackInfo info) {
+	public void onPlayerConnect(ClientConnection connection, ServerPlayerEntity player, CallbackInfo info) {
 		Optional<Packet<?>> optionalPacket = PacketRegistryImpl.createInitialRegisterPacket(ServerSidePacketRegistry.INSTANCE);
 		//noinspection OptionalIsPresent
 		if (optionalPacket.isPresent()) {
-			lvt2.networkHandler.sendPacket(optionalPacket.get());
+			player.networkHandler.sendPacket(optionalPacket.get());
+			((ServerSidePacketRegistryImpl) ServerSidePacketRegistry.INSTANCE).addNetworkHandler(player.networkHandler);
 		}
 	}
 }

fabric-networking/src/main/resources/fabric-networking.mixins.json

@@ -10,7 +10,8 @@
     "MixinThreadedAnvilChunkStorage"
   ],
   "client": [
-    "MixinClientPlayNetworkHandler"
+    "MixinClientPlayNetworkHandler",
+    "MixinMinecraftClient"
   ],
   "injectors": {
     "defaultRequire": 1