Capture a consistent and reusable 'last seen' state for use in command packets

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

@@ -54,6 +54,7 @@ import com.velocitypowered.proxy.protocol.packet.ServerboundCookieResponsePacket
 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 class ClientPlaySessionHandler implements MinecraftSessionHandler {
     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/protocol/packet/chat/ChatAcknowledgementPacket.java

@@ -54,4 +54,8 @@ public class ChatAcknowledgementPacket implements MinecraftPacket {
                 "offset=" + offset +
                 '}';
     }
+
+    public int offset() {
+        return offset;
+    }
 }

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

@@ -23,7 +23,9 @@ import com.velocitypowered.proxy.protocol.MinecraftPacket;
 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;
 
 /**
@@ -64,13 +66,15 @@ public class ChatQueue {
    * 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 new {@link Instant} timestamp of this packet to update the internal chat state.
+   * @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, @Nullable Instant timestamp) {
+  public void queuePacket(Function<LastSeenMessages, CompletableFuture<MinecraftPacket>> nextPacket, @Nullable Instant timestamp, @Nullable LastSeenMessages lastSeenMessages) {
     queueTask((chatState, smc) -> {
-      chatState.update(timestamp);
-      return nextPacket.thenCompose(packet -> writePacket(packet, smc));
+      LastSeenMessages newLastSeenMessages = chatState.updateFromMessage(timestamp, lastSeenMessages);
+      return nextPacket.apply(newLastSeenMessages).thenCompose(packet -> writePacket(packet, smc));
     });
   }
 
@@ -88,6 +92,16 @@ public class ChatQueue {
     });
   }
 
+  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 <T extends MinecraftPacket> CompletableFuture<Void> writePacket(T packet, MinecraftConnection smc) {
     return CompletableFuture.runAsync(() -> {
       if (!smc.isClosed()) {
@@ -103,16 +117,64 @@ public class ChatQueue {
     CompletableFuture<Void> update(ChatState chatState, MinecraftConnection smc);
   }
 
+  /**
+   * 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>
+   *
+   * Note that this is effectively unused for 1.20.5+ clients, as commands without any signature do not send 'last seen'
+   * updates.
+   */
   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();
+
     public volatile Instant lastTimestamp = Instant.EPOCH;
+    private volatile BitSet lastSeenMessages = new BitSet();
+    private final AtomicInteger delayedAckCount = new AtomicInteger();
 
     private ChatState() {
     }
 
-    public void update(@Nullable Instant timestamp) {
+    @Nullable
+    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;
+    }
+
+    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.connection.client.ConnectedPlayer;
 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 @@ public interface CommandHandler<T extends MinecraftPacket> {
   }
 
   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 @@ public interface CommandHandler<T extends MinecraftPacket> {
               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 @@ import java.util.BitSet;
 
 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 class 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);
 
@@ -54,6 +60,14 @@ public class LastSeenMessages {
     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{" +

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/keyed/KeyedChatHandler.java

@@ -91,11 +91,12 @@ public class KeyedChatHandler implements
       });
     }
     chatQueue.queuePacket(
-        chatFuture.exceptionally((ex) -> {
+        newLastSeen -> chatFuture.exceptionally((ex) -> {
           logger.error("Exception while handling player chat for {}", player, ex);
           return null;
         }),
-        packet.getExpiry()
+        packet.getExpiry(),
+        null
     );
   }
 

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/keyed/KeyedCommandHandler.java

@@ -43,7 +43,7 @@ public class KeyedCommandHandler implements CommandHandler<KeyedPlayerCommandPac
 
   @Override
   public void handlePlayerCommandInternal(KeyedPlayerCommandPacket packet) {
-    queueCommandResult(this.server, this.player, event -> {
+    queueCommandResult(this.server, this.player, (event, newLastSeenMessages) -> {
       CommandExecuteEvent.CommandResult result = event.getResult();
       IdentifiedKey playerKey = player.getIdentifiedKey();
       if (result == CommandExecuteEvent.CommandResult.denied()) {
@@ -111,6 +111,6 @@ public class KeyedCommandHandler implements CommandHandler<KeyedPlayerCommandPac
         }
         return null;
       });
-    }, packet.getCommand(), packet.getTimestamp());
+    }, packet.getCommand(), packet.getTimestamp(), null);
   }
 }

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/legacy/LegacyCommandHandler.java

@@ -42,7 +42,7 @@ public class LegacyCommandHandler implements CommandHandler<LegacyChatPacket> {
   @Override
   public void handlePlayerCommandInternal(LegacyChatPacket packet) {
     String command = packet.getMessage().substring(1);
-    queueCommandResult(this.server, this.player, event -> {
+    queueCommandResult(this.server, this.player, (event, newLastSeenMessages) -> {
       CommandExecuteEvent.CommandResult result = event.getResult();
       if (result == CommandExecuteEvent.CommandResult.denied()) {
         return CompletableFuture.completedFuture(null);
@@ -62,6 +62,6 @@ public class LegacyCommandHandler implements CommandHandler<LegacyChatPacket> {
         }
         return null;
       });
-    }, command, Instant.now());
+    }, command, Instant.now(), null);
   }
 }

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/session/SessionChatHandler.java

@@ -29,6 +29,8 @@ import com.velocitypowered.proxy.protocol.packet.chat.ChatQueue;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
+import java.util.concurrent.CompletableFuture;
+
 public class SessionChatHandler implements ChatHandler<SessionPlayerChatPacket> {
 
   private static final Logger logger = LogManager.getLogger(SessionChatHandler.class);
@@ -51,8 +53,9 @@ public class SessionChatHandler implements ChatHandler<SessionPlayerChatPacket>
     ChatQueue chatQueue = this.player.getChatQueue();
     EventManager eventManager = this.server.getEventManager();
     PlayerChatEvent toSend = new PlayerChatEvent(player, packet.getMessage());
+    CompletableFuture<PlayerChatEvent> eventFuture = eventManager.fire(toSend);
     chatQueue.queuePacket(
-        eventManager.fire(toSend)
+        newLastSeenMessages -> eventFuture
             .thenApply(pme -> {
               PlayerChatEvent.ChatResult chatResult = pme.getResult();
               if (!chatResult.isAllowed()) {
@@ -72,13 +75,14 @@ public class SessionChatHandler implements ChatHandler<SessionPlayerChatPacket>
                     .setTimestamp(packet.timestamp)
                     .toServer();
               }
-              return packet;
+              return packet.withLastSeenMessages(newLastSeenMessages);
             })
             .exceptionally((ex) -> {
               logger.error("Exception while handling player chat for {}", player, ex);
               return null;
             }),
-        packet.getTimestamp()
+        packet.getTimestamp(),
+        packet.getLastSeenMessages()
     );
   }
 }

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/session/SessionCommandHandler.java

@@ -93,23 +93,25 @@ public class SessionCommandHandler implements CommandHandler<SessionPlayerComman
 
   @Override
   public void handlePlayerCommandInternal(SessionPlayerCommandPacket packet) {
-    queueCommandResult(this.server, this.player, event -> {
+    queueCommandResult(this.server, this.player, (event, newLastSeenMessages) -> {
+      SessionPlayerCommandPacket fixedPacket = packet.withLastSeenMessages(newLastSeenMessages);
+
       CommandExecuteEvent.CommandResult result = event.getResult();
       if (result == CommandExecuteEvent.CommandResult.denied()) {
-        return CompletableFuture.completedFuture(consumeCommand(packet));
+        return CompletableFuture.completedFuture(consumeCommand(fixedPacket));
       }
 
-      String commandToRun = result.getCommand().orElse(packet.command);
+      String commandToRun = result.getCommand().orElse(fixedPacket.command);
       if (result.isForwardToServer()) {
-        return CompletableFuture.completedFuture(forwardCommand(packet, commandToRun));
+        return CompletableFuture.completedFuture(forwardCommand(fixedPacket, commandToRun));
       }
 
       return runCommand(this.server, this.player, commandToRun, hasRun -> {
         if (hasRun) {
-          return consumeCommand(packet);
+          return consumeCommand(fixedPacket);
         }
-        return forwardCommand(packet, commandToRun);
+        return forwardCommand(fixedPacket, commandToRun);
       });
-    }, packet.command, packet.timeStamp);
+    }, packet.command, packet.timeStamp, packet.lastSeenMessages);
   }
 }

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/session/SessionPlayerChatPacket.java

@@ -99,4 +99,15 @@ public class SessionPlayerChatPacket implements MinecraftPacket {
     buf.readBytes(signature);
     return signature;
   }
+
+  public SessionPlayerChatPacket withLastSeenMessages(LastSeenMessages lastSeenMessages) {
+    SessionPlayerChatPacket packet = new SessionPlayerChatPacket();
+    packet.message = message;
+    packet.timestamp = timestamp;
+    packet.salt = salt;
+    packet.signed = signed;
+    packet.signature = signature;
+    packet.lastSeenMessages = lastSeenMessages;
+    return packet;
+  }
 }

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/session/SessionPlayerCommandPacket.java

@@ -25,6 +25,8 @@ import com.velocitypowered.proxy.protocol.ProtocolUtils;
 import com.velocitypowered.proxy.protocol.packet.chat.LastSeenMessages;
 import com.velocitypowered.proxy.util.except.QuietDecoderException;
 import io.netty.buffer.ByteBuf;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
 import java.time.Instant;
 import java.util.List;
 
@@ -83,6 +85,21 @@ public class SessionPlayerCommandPacket implements MinecraftPacket {
             '}';
   }
 
+  public SessionPlayerCommandPacket withLastSeenMessages(@Nullable LastSeenMessages lastSeenMessages) {
+    if (lastSeenMessages == null) {
+      UnsignedPlayerCommandPacket packet = new UnsignedPlayerCommandPacket();
+      packet.command = command;
+      return packet;
+    }
+    SessionPlayerCommandPacket packet = new SessionPlayerCommandPacket();
+    packet.command = command;
+    packet.timeStamp = timeStamp;
+    packet.salt = salt;
+    packet.argumentSignatures = argumentSignatures;
+    packet.lastSeenMessages = lastSeenMessages;
+    return packet;
+  }
+
   public static class ArgumentSignatures {
 
     private final List<ArgumentSignature> entries;

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/chat/session/UnsignedPlayerCommandPacket.java

@@ -19,7 +19,9 @@ package com.velocitypowered.proxy.protocol.packet.chat.session;
 
 import com.velocitypowered.api.network.ProtocolVersion;
 import com.velocitypowered.proxy.protocol.ProtocolUtils;
+import com.velocitypowered.proxy.protocol.packet.chat.LastSeenMessages;
 import io.netty.buffer.ByteBuf;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 public class UnsignedPlayerCommandPacket extends SessionPlayerCommandPacket {
 
@@ -33,6 +35,11 @@ public class UnsignedPlayerCommandPacket extends SessionPlayerCommandPacket {
     ProtocolUtils.writeString(buf, this.command);
   }
 
+  @Override
+  public SessionPlayerCommandPacket withLastSeenMessages(@Nullable LastSeenMessages lastSeenMessages) {
+    return this;
+  }
+
   public boolean isSigned() {
     return false;
   }