Fix chat race condition

SONO TROPPO FORTE
CoralRP · Aug 20, 2023 · b2d3fa2 · b2d3fa2
1 parent 19abb90
commit b2d3fa2
Showing 1 changed file with 14 additions and 15 deletions.
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/ChatQueue.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/ChatQueue.java
@@ -20,9 +20,10 @@
 import com.velocitypowered.proxy.connection.MinecraftConnection;
 import com.velocitypowered.proxy.connection.client.ConnectedPlayer;
 import com.velocitypowered.proxy.protocol.MinecraftPacket;
+
 import java.time.Instant;
 import java.util.concurrent.CompletableFuture;
-import java.util.function.BiConsumer;
+import java.util.function.Function;
 
 /**
  * A precisely ordered queue which allows for outside entries into the ordered queue through
@@ -58,9 +59,8 @@ public void queuePacket(CompletableFuture<MinecraftPacket> nextPacket, Instant t
       MinecraftConnection smc = player.ensureAndGetCurrentServer().ensureConnected();
 
       CompletableFuture<WrappedPacket> nextInLine = WrappedPacket.wrap(timestamp, nextPacket);
-      awaitChat(smc, this.packetFuture,
-          nextInLine); // we await chat, binding `this.packetFuture` -> `nextInLine`
-      this.packetFuture = nextInLine;
+      this.packetFuture = awaitChat(smc, this.packetFuture, nextInLine);
+      ;
     }
   }
 
@@ -84,21 +84,21 @@ public <K, V extends MinecraftPacket> void hijack(K packet,
     }
   }
 
-  private static BiConsumer<WrappedPacket, Throwable> writePacket(MinecraftConnection connection) {
-    return (wrappedPacket, throwable) -> {
-      if (wrappedPacket != null && !connection.isClosed()) {
+  private static Function<WrappedPacket, WrappedPacket> writePacket(MinecraftConnection connection) {
+    return wrappedPacket -> {
+      if (!connection.isClosed()) {
         wrappedPacket.write(connection);
       }
+
+      return wrappedPacket;
     };
   }
 
-  private static <T extends MinecraftPacket> void awaitChat(
+  private static <T extends MinecraftPacket> CompletableFuture<WrappedPacket> awaitChat(
       MinecraftConnection connection,
       CompletableFuture<WrappedPacket> binder,
-      CompletableFuture<WrappedPacket> future
-  ) {
-    // the binder will run -> then the future will get the `write packet` caller
-    binder.whenComplete((ignored1, ignored2) -> future.whenComplete(writePacket(connection)));
+      CompletableFuture<WrappedPacket> future) {
+    return binder.whenCompleteAsync((ignored1, ignored2) -> future.thenApply(writePacket(connection)).join());
   }
 
   private static <K, V extends MinecraftPacket> CompletableFuture<WrappedPacket> hijackCurrentPacket(
@@ -113,9 +113,8 @@ private static <K, V extends MinecraftPacket> CompletableFuture<WrappedPacket> h
       // map the new packet into a better "designed" packet with the hijacked packet's timestamp
       WrappedPacket.wrap(previous.timestamp,
               future.thenApply(item -> packetMapper.map(previous.timestamp, item)))
-          .whenCompleteAsync(writePacket(connection), connection.eventLoop())
-          .whenComplete(
-              (packet, throwable) -> awaitedFuture.complete(throwable != null ? null : packet));
+              .thenApplyAsync(writePacket(connection), connection.eventLoop())
+              .whenComplete((packet, throwable) -> awaitedFuture.complete(throwable != null ? null : packet));
     });
     return awaitedFuture;
   }