A couple of fixes for handling of secure chat

proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java

@@ -54,6 +54,7 @@
 import com.velocitypowered.proxy.protocol.packet.TabCompleteRequestPacket;
 import com.velocitypowered.proxy.protocol.packet.TabCompleteResponsePacket;
 import com.velocitypowered.proxy.protocol.packet.TabCompleteResponsePacket.Offer;
+import com.velocitypowered.proxy.protocol.packet.chat.ChatAcknowledgementPacket;
 import com.velocitypowered.proxy.protocol.packet.chat.ChatHandler;
 import com.velocitypowered.proxy.protocol.packet.chat.ChatTimeKeeper;
 import com.velocitypowered.proxy.protocol.packet.chat.CommandHandler;
@@ -423,6 +424,15 @@ public boolean handle(FinishedUpdatePacket packet) {
     return true;
   }
 
+  @Override
+  public boolean handle(ChatAcknowledgementPacket packet) {
+    if (player.getCurrentServer().isEmpty()) {
+      return true;
+    }
+    player.getChatQueue().handleAcknowledgement(packet.offset());
+    return true;
+  }
+
   @Override
   public boolean handle(ServerboundCookieResponsePacket packet) {
     server.getEventManager()

proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java

@@ -83,6 +83,7 @@
 import com.velocitypowered.proxy.protocol.packet.chat.ComponentHolder;
 import com.velocitypowered.proxy.protocol.packet.chat.PlayerChatCompletionPacket;
 import com.velocitypowered.proxy.protocol.packet.chat.builder.ChatBuilderFactory;
+import com.velocitypowered.proxy.protocol.packet.chat.builder.ChatBuilderV2;
 import com.velocitypowered.proxy.protocol.packet.chat.legacy.LegacyChatPacket;
 import com.velocitypowered.proxy.protocol.packet.config.ClientboundServerLinksPacket;
 import com.velocitypowered.proxy.protocol.packet.config.StartUpdatePacket;
@@ -1108,11 +1109,12 @@ public void spoofChatInput(String input) {
         "input cannot be greater than " + LegacyChatPacket.MAX_SERVERBOUND_MESSAGE_LENGTH
             + " characters in length");
     if (getProtocolVersion().noLessThan(ProtocolVersion.MINECRAFT_1_19)) {
-      this.chatQueue.hijack(getChatBuilderFactory().builder().asPlayer(this).message(input),
-          (instant, item) -> {
-            item.setTimestamp(instant);
-            return item.toServer();
-          });
+      ChatBuilderV2 message = getChatBuilderFactory().builder().asPlayer(this).message(input);
+      this.chatQueue.queuePacket(chatState -> {
+        message.setTimestamp(chatState.lastTimestamp);
+        message.setLastSeenMessages(chatState.createLastSeen());
+        return message.toServer();
+      });
     } else {
       ensureBackendConnection().write(getChatBuilderFactory().builder()
           .asPlayer(this).message(input).toServer());

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/ChatAcknowledgementPacket.java

@@ -54,4 +54,8 @@ public String toString() {
                 "offset=" + offset +
                 '}';
     }
+
+    public int offset() {
+        return offset;
+    }
 }

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/ChatQueue.java

@@ -23,7 +23,9 @@
 import io.netty.channel.ChannelFuture;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import java.time.Instant;
+import java.util.BitSet;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 
 /**
@@ -32,9 +34,10 @@
  */
 public class ChatQueue {
 
-  private final Object internalLock;
+  private final Object internalLock = new Object();
   private final ConnectedPlayer player;
-  private CompletableFuture<WrappedPacket> packetFuture;
+  private final ChatState chatState = new ChatState();
+  private CompletableFuture<Void> head = CompletableFuture.completedFuture(null);
 
   /**
    * Instantiates a {@link ChatQueue} for a specific {@link ConnectedPlayer}.
@@ -43,130 +46,135 @@ public class ChatQueue {
    */
   public ChatQueue(ConnectedPlayer player) {
     this.player = player;
-    this.packetFuture = CompletableFuture.completedFuture(new WrappedPacket(Instant.EPOCH, null));
-    this.internalLock = new Object();
+  }
+
+  private void queueTask(Task task) {
+    synchronized (internalLock) {
+      MinecraftConnection smc = player.ensureAndGetCurrentServer().ensureConnected();
+      head = head.thenCompose(v -> {
+        try {
+          return task.update(chatState, smc).exceptionally(ignored -> null);
+        } catch (Throwable ignored) {
+          return CompletableFuture.completedFuture(null);
+        }
+      });
+    }
   }
 
   /**
    * Queues a packet sent from the player - all packets must wait until this processes to send their
    * packets. This maintains order on the server-level for the client insertions of commands
    * and messages. All entries are locked through an internal object lock.
    *
-   * @param nextPacket the {@link CompletableFuture} which will provide the next-processed packet.
-   * @param timestamp  the {@link Instant} timestamp of this packet so we can allow piggybacking.
+   * @param nextPacket       a function mapping {@link LastSeenMessages} state to a {@link CompletableFuture} that will
+   *                         provide the next-processed packet. This should include the fixed {@link LastSeenMessages}.
+   * @param timestamp        the new {@link Instant} timestamp of this packet to update the internal chat state.
+   * @param lastSeenMessages the new {@link LastSeenMessages} last seen messages to update the internal chat state.
    */
-  public void queuePacket(CompletableFuture<MinecraftPacket> nextPacket, Instant timestamp) {
-    synchronized (internalLock) { // wait for the lock to resolve - we don't want to drop packets
-      MinecraftConnection smc = player.ensureAndGetCurrentServer().ensureConnected();
-
-      CompletableFuture<WrappedPacket> nextInLine = WrappedPacket.wrap(timestamp, nextPacket);
-      this.packetFuture = awaitChat(smc, this.packetFuture,
-          nextInLine); // we await chat, binding `this.packetFuture` -> `nextInLine`
-    }
+  public void queuePacket(Function<LastSeenMessages, CompletableFuture<MinecraftPacket>> nextPacket, @Nullable Instant timestamp, @Nullable LastSeenMessages lastSeenMessages) {
+    queueTask((chatState, smc) -> {
+      LastSeenMessages newLastSeenMessages = chatState.updateFromMessage(timestamp, lastSeenMessages);
+      return nextPacket.apply(newLastSeenMessages).thenCompose(packet -> writePacket(packet, smc));
+    });
   }
 
   /**
-   * Hijacks the latest sent packet's timestamp to provide an in-order packet without polling the
+   * Hijacks the latest sent packet's chat state to provide an in-order packet without polling the
    * physical, or prior packets sent through the stream.
    *
-   * @param packet        the {@link MinecraftPacket} to send.
-   * @param instantMapper the {@link InstantPacketMapper} which maps the prior timestamp and current
-   *                      packet to a new packet.
-   * @param <K>           the type of base to expect when mapping the packet.
-   * @param <V>           the type of packet for instantMapper type-checking.
+   * @param packetFunction a function that maps the prior {@link ChatState} into a new packet.
+   * @param <T>            the type of packet to send.
    */
-  public <K, V extends MinecraftPacket> void hijack(K packet,
-      InstantPacketMapper<K, V> instantMapper) {
-    synchronized (internalLock) {
-      CompletableFuture<K> trueFuture = CompletableFuture.completedFuture(packet);
-      MinecraftConnection smc = player.ensureAndGetCurrentServer().ensureConnected();
+  public <T extends MinecraftPacket> void queuePacket(Function<ChatState, T> packetFunction) {
+    queueTask((chatState, smc) -> {
+      T packet = packetFunction.apply(chatState);
+      return writePacket(packet, smc);
+    });
+  }
 
-      this.packetFuture = hijackCurrentPacket(smc, this.packetFuture, trueFuture, instantMapper);
-    }
+  public void handleAcknowledgement(int offset) {
+    queueTask((chatState, smc) -> {
+      int ackCountToForward = chatState.accumulateAckCount(offset);
+      if (ackCountToForward > 0) {
+        return writePacket(new ChatAcknowledgementPacket(ackCountToForward), smc);
+      }
+      return CompletableFuture.completedFuture(null);
+    });
   }
 
-  private static Function<WrappedPacket, WrappedPacket> writePacket(MinecraftConnection connection) {
-    return wrappedPacket -> {
-      if (!connection.isClosed()) {
-        ChannelFuture future = wrappedPacket.write(connection);
+  private static <T extends MinecraftPacket> CompletableFuture<Void> writePacket(T packet, MinecraftConnection smc) {
+    return CompletableFuture.runAsync(() -> {
+      if (!smc.isClosed()) {
+        ChannelFuture future = smc.write(packet);
         if (future != null) {
           future.awaitUninterruptibly();
         }
       }
-
-      return wrappedPacket;
-    };
+    }, smc.eventLoop());
   }
 
-  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
-    return binder.thenCompose(ignored -> future.thenApply(writePacket(connection)));
-  }
-
-  private static <K, V extends MinecraftPacket> CompletableFuture<WrappedPacket> hijackCurrentPacket(
-      MinecraftConnection connection,
-      CompletableFuture<WrappedPacket> binder,
-      CompletableFuture<K> future,
-      InstantPacketMapper<K, V> packetMapper
-  ) {
-    CompletableFuture<WrappedPacket> awaitedFuture = new CompletableFuture<>();
-    // the binder will complete -> then the future will get the `write packet` caller
-    binder.whenComplete((previous, ignored) -> {
-      // 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)))
-          .thenApplyAsync(writePacket(connection), connection.eventLoop())
-          .whenComplete(
-              (packet, throwable) -> awaitedFuture.complete(throwable != null ? null : packet));
-    });
-    return awaitedFuture;
+  private interface Task {
+    CompletableFuture<Void> update(ChatState chatState, MinecraftConnection smc);
   }
 
   /**
-   * Provides an {@link Instant} based timestamp mapper from an existing object to create a packet.
+   * Tracks the last Secure Chat state that we received from the client. This is important to always have a valid 'last
+   * seen' state that is consistent with future and past updates from the client (which may be signed). This state is
+   * used to construct 'spoofed' command packets from the proxy to the server.
+   * <ul>
+   *     <li>If we last forwarded a chat or command packet from the client, we have a known 'last seen' that we can
+   *     reuse.</li>
+   *     <li>If we last forwarded a {@link ChatAcknowledgementPacket}, the previous 'last seen' cannot be reused. We
+   *     cannot predict an up-to-date 'last seen', as we do not know which messages the client actually saw.</li>
+   *     <li>Therefore, we need to hold back any acknowledgement packets so that we can continue to reuse the last valid
+   *     'last seen' state.</li>
+   *     <li>However, there is a limit to the number of messages that can remain unacknowledged on the server.</li>
+   *     <li>To address this, we know that if the client has moved its 'last seen' window far enough, we can fill in the
+   *     gap with dummy 'last seen', and it will never be checked.</li>
+   * </ul>
    *
-   * @param <K> The base object type to map.
-   * @param <V> The resulting packet type.
+   * Note that this is effectively unused for 1.20.5+ clients, as commands without any signature do not send 'last seen'
+   * updates.
    */
-  public interface InstantPacketMapper<K, V extends MinecraftPacket> {
-
-    /**
-     * Maps a value into a packet with it and a timestamp.
-     *
-     * @param nextInstant   the {@link Instant} timestamp to use for tracking.
-     * @param currentObject the current item to map to the packet.
-     * @return The resulting packet from the mapping.
-     */
-    V map(Instant nextInstant, K currentObject);
-  }
-
-  private static class WrappedPacket {
+  public static class ChatState {
+    private static final int MINIMUM_DELAYED_ACK_COUNT = LastSeenMessages.WINDOW_SIZE;
+    private static final BitSet DUMMY_LAST_SEEN_MESSAGES = new BitSet();
 
-    private final Instant timestamp;
-    private final MinecraftPacket packet;
+    public volatile Instant lastTimestamp = Instant.EPOCH;
+    private volatile BitSet lastSeenMessages = new BitSet();
+    private final AtomicInteger delayedAckCount = new AtomicInteger();
 
-    private WrappedPacket(Instant timestamp, MinecraftPacket packet) {
-      this.timestamp = timestamp;
-      this.packet = packet;
+    private ChatState() {
     }
 
     @Nullable
-    public ChannelFuture write(MinecraftConnection connection) {
-      if (packet != null) {
-        return connection.write(packet);
+    public LastSeenMessages updateFromMessage(@Nullable Instant timestamp, @Nullable LastSeenMessages lastSeenMessages) {
+      if (timestamp != null) {
+        this.lastTimestamp = timestamp;
+      }
+      if (lastSeenMessages != null) {
+        // We held back some acknowledged messages, so flush that out now that we have a known 'last seen' state again
+        int delayedAckCount = this.delayedAckCount.getAndSet(0);
+        this.lastSeenMessages = lastSeenMessages.getAcknowledged();
+        return lastSeenMessages.offset(delayedAckCount);
       }
       return null;
     }
 
-    private static CompletableFuture<WrappedPacket> wrap(Instant timestamp,
-        CompletableFuture<MinecraftPacket> nextPacket) {
-      return nextPacket
-          .thenApply(pkt -> new WrappedPacket(timestamp, pkt))
-          .exceptionally(ignored -> new WrappedPacket(timestamp, null));
+    public int accumulateAckCount(int ackCount) {
+      int delayedAckCount = this.delayedAckCount.addAndGet(ackCount);
+      int ackCountToForward = delayedAckCount - MINIMUM_DELAYED_ACK_COUNT;
+      if (ackCountToForward >= LastSeenMessages.WINDOW_SIZE) {
+        // Because we only forward acknowledgements above the window size, we don't have to shift the previous 'last seen' state
+        this.lastSeenMessages = DUMMY_LAST_SEEN_MESSAGES;
+        this.delayedAckCount.set(MINIMUM_DELAYED_ACK_COUNT);
+        return ackCountToForward;
+      }
+      return 0;
+    }
+
+    public LastSeenMessages createLastSeen() {
+      return new LastSeenMessages(0, lastSeenMessages);
     }
   }
 }

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/CommandHandler.java

@@ -23,11 +23,13 @@
 import com.velocitypowered.proxy.protocol.MinecraftPacket;
 import java.time.Instant;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.format.NamedTextColor;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 public interface CommandHandler<T extends MinecraftPacket> {
 
@@ -53,11 +55,12 @@ default CompletableFuture<MinecraftPacket> runCommand(VelocityServer server,
   }
 
   default void queueCommandResult(VelocityServer server, ConnectedPlayer player,
-      Function<CommandExecuteEvent, CompletableFuture<MinecraftPacket>> futurePacketCreator,
-      String message, Instant timestamp) {
-    player.getChatQueue().queuePacket(
-        server.getCommandManager().callCommandEvent(player, message)
-            .thenComposeAsync(futurePacketCreator)
+      BiFunction<CommandExecuteEvent, LastSeenMessages, CompletableFuture<MinecraftPacket>> futurePacketCreator,
+      String message, Instant timestamp, @Nullable LastSeenMessages lastSeenMessages) {
+      CompletableFuture<CommandExecuteEvent> eventFuture = server.getCommandManager().callCommandEvent(player, message);
+      player.getChatQueue().queuePacket(
+        newLastSeenMessages -> eventFuture
+            .thenComposeAsync(event -> futurePacketCreator.apply(event, newLastSeenMessages))
             .thenApply(pkt -> {
               if (server.getConfiguration().isLogCommandExecutions()) {
                 logger.info("{} -> executed command /{}", player, message);
@@ -68,6 +71,6 @@ default void queueCommandResult(VelocityServer server, ConnectedPlayer player,
               player.sendMessage(
                   Component.translatable("velocity.command.generic-error", NamedTextColor.RED));
               return null;
-            }), timestamp);
+            }), timestamp, lastSeenMessages);
   }
 }

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/LastSeenMessages.java

@@ -24,7 +24,8 @@
 
 public class LastSeenMessages {
 
-  private static final int DIV_FLOOR = -Math.floorDiv(-20, 8);
+  public static final int WINDOW_SIZE = 20;
+  private static final int DIV_FLOOR = -Math.floorDiv(-WINDOW_SIZE, 8);
   private int offset;
   private BitSet acknowledged;
 
@@ -33,6 +34,11 @@ public LastSeenMessages() {
     this.acknowledged = new BitSet();
   }
 
+  public LastSeenMessages(int offset, BitSet acknowledged) {
+    this.offset = offset;
+    this.acknowledged = acknowledged;
+  }
+
   public LastSeenMessages(ByteBuf buf) {
     this.offset = ProtocolUtils.readVarInt(buf);
 
@@ -46,14 +52,18 @@ public void encode(ByteBuf buf) {
     buf.writeBytes(Arrays.copyOf(acknowledged.toByteArray(), DIV_FLOOR));
   }
 
-  public boolean isEmpty() {
-    return acknowledged.isEmpty();
-  }
-
   public int getOffset() {
     return this.offset;
   }
 
+  public BitSet getAcknowledged() {
+    return acknowledged;
+  }
+
+  public LastSeenMessages offset(final int offset) {
+    return new LastSeenMessages(this.offset + offset, acknowledged);
+  }
+
   @Override
   public String toString() {
     return "LastSeenMessages{" +