add ability to directly listen to response packets

Author: jumpingpxl

api/src/main/java/net/labymod/serverapi/api/Protocol.java

@@ -25,12 +25,14 @@
 package net.labymod.serverapi.api;
 
 import net.labymod.serverapi.api.packet.Direction;
+import net.labymod.serverapi.api.packet.IdentifiablePacket;
 import net.labymod.serverapi.api.packet.OnlyWriteConstructor;
 import net.labymod.serverapi.api.packet.Packet;
 import net.labymod.serverapi.api.packet.PacketHandler;
 import net.labymod.serverapi.api.payload.PayloadChannelIdentifier;
 import net.labymod.serverapi.api.payload.io.PayloadReader;
 import net.labymod.serverapi.api.payload.io.PayloadWriter;
+import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import sun.misc.Unsafe;
@@ -49,6 +51,7 @@ public class Protocol {
   private final Set<ProtocolPacket> packets;
   private final AbstractProtocolService protocolService;
   private final AbstractProtocolService.Side protocolSide;
+  private final Set<AwaitingResponse> awaitingResponses;
 
   /**
    * Creates a new protocol.
@@ -67,6 +70,7 @@ public Protocol(
     this.protocolService = protocolService;
     this.protocolSide = protocolService.getSide();
     this.packets = new HashSet<>();
+    this.awaitingResponses = new HashSet<>();
   }
 
   /**
@@ -233,6 +237,29 @@ public void handlePacket(@NotNull UUID sender, @NotNull Packet packet) {
     return packet;
   }
 
+  /**
+   * Sends a packet to the recipient and waits for the provided response
+   *
+   * @param recipient        the recipient of the packet
+   * @param packet           the packet to send
+   * @param responseClass    the class of the response packet
+   * @param responseCallback the callback for the response packet. Return {@code true} to wait
+   *                         for another packet or {@code false} remove the callback.
+   * @param <R>              the type of the response packet
+   * @throws NullPointerException If the recipient, packet, response class or response callback is
+   *                              null.
+   */
+  public <R extends IdentifiablePacket> void sendPacket(
+      @NotNull UUID recipient,
+      @NotNull IdentifiablePacket packet,
+      @NotNull Class<R> responseClass,
+      @NotNull Predicate<R> responseCallback
+  ) {
+    Objects.requireNonNull(responseClass, "Response class cannot be null");
+    Objects.requireNonNull(responseCallback, "Response callback cannot be null");
+    this.sendPacketInternal(recipient, packet, responseClass, responseCallback);
+  }
+
   /**
    * Sends a packet to the recipient.
    *
@@ -241,24 +268,19 @@ public void handlePacket(@NotNull UUID sender, @NotNull Packet packet) {
    * @throws NullPointerException If the recipient or packet is null.
    */
   public void sendPacket(@NotNull UUID recipient, @NotNull Packet packet) {
-    Objects.requireNonNull(recipient, "Recipient cannot be null");
-    Objects.requireNonNull(packet, "Packet cannot be null");
-    ProtocolPacket protocolPacket = this.getProtocolPacketByClass(packet.getClass());
-    if (protocolPacket == null) {
-      throw new IllegalArgumentException(
-          "No packet with the class " + packet.getClass().getSimpleName() + " in protocol "
-              + this.identifier + "found");
-    }
+    this.sendPacketInternal(recipient, packet, null, null);
+  }
 
-    if (this.protocolSide.isAcceptingDirection(protocolPacket.direction)) {
-      return;
+  /**
+   * Clears all awaiting responses for the recipient.
+   *
+   * @param recipient the recipient to clear the awaiting responses for
+   */
+  @ApiStatus.Internal
+  public void clearAwaitingResponsesFor(UUID recipient) {
+    synchronized (this.awaitingResponses) {
+      this.awaitingResponses.removeIf(response -> response.recipient.equals(recipient));
     }
-
-    PayloadWriter writer = new PayloadWriter();
-    writer.writeVarInt(protocolPacket.id);
-    packet.write(writer);
-    this.protocolService.send(this.identifier, recipient, writer);
-    this.protocolService.afterPacketSent(this, packet, recipient);
   }
 
   private ProtocolPacket getProtocolPacket(Predicate<ProtocolPacket> filter) {
@@ -280,13 +302,84 @@ private ProtocolPacket getProtocolPacketById(int id) {
   }
 
   private void handlePacket(ProtocolPacket protocolPacket, UUID sender, Packet packet) {
+    if (packet instanceof IdentifiablePacket identifiablePacket) {
+      AwaitingResponse responsePacket = null;
+      for (AwaitingResponse awaitingResponse : this.awaitingResponses) {
+        if (awaitingResponse.recipient.equals(sender)
+            && awaitingResponse.responseClass.equals(protocolPacket.packet)
+            && identifiablePacket.getIdentifier() == awaitingResponse.identifier) {
+          responsePacket = awaitingResponse;
+          break;
+        }
+      }
+
+      if (responsePacket != null) {
+        try {
+          if (!responsePacket.responseCallback.test(identifiablePacket)) {
+            this.awaitingResponses.remove(responsePacket);
+          }
+        } catch (Exception e) {
+          e.printStackTrace();
+        }
+      }
+    }
+
     for (PacketHandler handler : protocolPacket.handlers) {
       handler.handle(sender, packet);
     }
 
     this.protocolService.afterPacketHandled(this, packet, sender);
   }
 
+  private void sendPacketInternal(
+      @NotNull UUID recipient,
+      @NotNull Packet packet,
+      @Nullable Class<? extends IdentifiablePacket> responseClass,
+      @Nullable Predicate<? extends IdentifiablePacket> responseCallback
+  ) {
+    Objects.requireNonNull(recipient, "Recipient cannot be null");
+    Objects.requireNonNull(packet, "Packet cannot be null");
+    ProtocolPacket protocolPacket = this.getProtocolPacketByClass(packet.getClass());
+    if (protocolPacket == null) {
+      throw new IllegalArgumentException(
+          "No packet with the class " + packet.getClass().getSimpleName() + " in protocol "
+              + this.identifier + "found");
+    }
+
+    if (this.protocolSide.isAcceptingDirection(protocolPacket.direction)) {
+      return;
+    }
+
+    PayloadWriter writer = new PayloadWriter();
+    writer.writeVarInt(protocolPacket.id);
+    packet.write(writer);
+    this.protocolService.send(this.identifier, recipient, writer);
+    if (responseClass != null && responseCallback != null
+        && packet instanceof IdentifiablePacket identifiablePacket) {
+      synchronized (this.awaitingResponses) {
+        this.awaitingResponses.add(new AwaitingResponse(
+            recipient,
+            identifiablePacket,
+            identifiablePacket.getIdentifier(),
+            responseClass,
+            responseCallback
+        ));
+      }
+    }
+
+    this.protocolService.afterPacketSent(this, packet, recipient);
+  }
+
+  private record AwaitingResponse(
+      UUID recipient,
+      IdentifiablePacket initialPacket,
+      int identifier,
+      Class responseClass,
+      Predicate responseCallback
+  ) {
+
+  }
+
   private static class ProtocolPacket {
 
     private static final UnsafeProvider UNSAFE_PROVIDER = new UnsafeProvider();

core/src/main/java/net/labymod/serverapi/core/AbstractLabyModPlayer.java

@@ -26,13 +26,20 @@
 
 import net.labymod.serverapi.api.Protocol;
 import net.labymod.serverapi.api.model.component.ServerAPIComponent;
+import net.labymod.serverapi.api.packet.IdentifiablePacket;
 import net.labymod.serverapi.api.packet.Packet;
 import net.labymod.serverapi.core.integration.LabyModIntegrationPlayer;
 import net.labymod.serverapi.core.integration.LabyModProtocolIntegration;
 import net.labymod.serverapi.core.model.display.Subtitle;
 import net.labymod.serverapi.core.model.moderation.Permission;
+import net.labymod.serverapi.core.model.supplement.InputPrompt;
+import net.labymod.serverapi.core.model.supplement.ServerSwitchPrompt;
 import net.labymod.serverapi.core.packet.clientbound.game.display.SubtitlePacket;
 import net.labymod.serverapi.core.packet.clientbound.game.moderation.PermissionPacket;
+import net.labymod.serverapi.core.packet.clientbound.game.supplement.InputPromptPacket;
+import net.labymod.serverapi.core.packet.clientbound.game.supplement.ServerSwitchPromptPacket;
+import net.labymod.serverapi.core.packet.serverbound.game.supplement.InputPromptResponsePacket;
+import net.labymod.serverapi.core.packet.serverbound.game.supplement.ServerSwitchPromptResponsePacket;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -41,6 +48,7 @@
 import java.util.Objects;
 import java.util.UUID;
 import java.util.function.Consumer;
+import java.util.function.Predicate;
 
 public abstract class AbstractLabyModPlayer<P extends AbstractLabyModPlayer<?>> {
 
@@ -77,6 +85,13 @@ protected AbstractLabyModPlayer(
     }
   }
 
+  /**
+   * Gets the integration player of the provided class
+   *
+   * @param clazz The class of the integration player
+   * @param <T>   The type of the integration player
+   * @return The integration player or null if not found
+   */
   public <T extends LabyModIntegrationPlayer> @Nullable T getIntegrationPlayer(Class<T> clazz) {
     for (LabyModIntegrationPlayer integrationPlayer : this.integrationPlayers) {
       if (clazz == integrationPlayer.getClass()) {
@@ -87,30 +102,53 @@ protected AbstractLabyModPlayer(
     return null;
   }
 
+  /**
+   * @return The unique identifier of the player
+   */
   public @NotNull UUID getUniqueId() {
     return this.uniqueId;
   }
 
+  /**
+   * @return The LabyMod version of the player
+   */
   public @NotNull String getLabyModVersion() {
     return this.labyModVersion;
   }
 
+  /**
+   * @return the currently set subtitle
+   */
   public @Nullable Subtitle getSubtitle() {
     return this.subtitle;
   }
 
+  /**
+   * Sets the subtitle of the player and sends it to all other online players
+   *
+   * @param component The component to set as subtitle
+   */
   public void setSubtitle(@NotNull ServerAPIComponent component) {
     this.setSubtitleInternal(Subtitle.create(this.uniqueId, component));
   }
 
-  public void resetSubtitle(Consumer<P> consumer) {
-    this.setSubtitleInternal(null);
-  }
-
+  /**
+   * Sets the subtitle of the player and sends it to all other online players
+   *
+   * @param component The component to set as subtitle
+   * @param size      The size of the subtitle
+   */
   public void setSubtitle(@NotNull ServerAPIComponent component, double size) {
     this.setSubtitleInternal(Subtitle.create(this.uniqueId, component, size));
   }
 
+  /**
+   * Resets the subtitle of the player and sends it to all other online players
+   */
+  public void resetSubtitle() {
+    this.setSubtitleInternal(null);
+  }
+
   public void sendPacket(@NotNull Packet packet) {
     Class<? extends Packet> packetClass = packet.getClass();
     for (Protocol protocol : this.protocolService.registry().getProtocols()) {
@@ -120,15 +158,86 @@ public void sendPacket(@NotNull Packet packet) {
     }
   }
 
+  /**
+   * Sends the provided permissions to the player
+   *
+   * @param permissions The permissions to send
+   */
   public void sendPermissions(@NotNull Permission.StatedPermission... permissions) {
-    this.sendPacket(new PermissionPacket(permissions));
+    this.sendLabyModPacket(new PermissionPacket(permissions));
   }
 
+  /**
+   * Sends the provided permissions to the player
+   *
+   * @param permissions The permissions to send
+   */
   public void sendPermissions(@NotNull List<Permission.StatedPermission> permissions) {
-    this.sendPacket(new PermissionPacket(permissions));
+    this.sendLabyModPacket(new PermissionPacket(permissions));
+  }
+
+  /**
+   * Opens the provided input prompt for the player
+   *
+   * @param inputPrompt The input prompt to open
+   */
+  public void openInputPrompt(@NotNull InputPrompt inputPrompt) {
+    this.sendLabyModPacket(new InputPromptPacket(inputPrompt));
+  }
+
+  /**
+   * Opens the provided input prompt for the player and handle the response via the provided
+   * consumer
+   *
+   * @param inputPrompt      The input prompt to open
+   * @param responseConsumer The consumer for the response
+   */
+  public void openInputPrompt(
+      @NotNull InputPrompt inputPrompt,
+      @NotNull Consumer<String> responseConsumer
+  ) {
+    Objects.requireNonNull(responseConsumer, "Response consumer cannot be null");
+    this.sendLabyModPacket(
+        new InputPromptPacket(inputPrompt),
+        InputPromptResponsePacket.class,
+        response -> {
+          responseConsumer.accept(response.getValue());
+          return false;
+        }
+    );
   }
 
-  private void setSubtitleInternal(Subtitle subtitle) {
+  /**
+   * Opens the provided server switch prompt
+   *
+   * @param serverSwitchPrompt The server switch prompt to open
+   */
+  public void openServerSwitchPrompt(@NotNull ServerSwitchPrompt serverSwitchPrompt) {
+    this.sendLabyModPacket(new ServerSwitchPromptPacket(serverSwitchPrompt));
+  }
+
+  /**
+   * Opens the provided server switch prompt and handle the response via the provided consumer
+   *
+   * @param serverSwitchPrompt The server switch prompt to open
+   * @param responseConsumer   The consumer for the response
+   */
+  public void openServerSwitchPrompt(
+      @NotNull ServerSwitchPrompt serverSwitchPrompt,
+      @NotNull Consumer<ServerSwitchPromptResponsePacket> responseConsumer
+  ) {
+    Objects.requireNonNull(responseConsumer, "Response consumer cannot be null");
+    this.sendLabyModPacket(
+        new ServerSwitchPromptPacket(serverSwitchPrompt),
+        ServerSwitchPromptResponsePacket.class,
+        response -> {
+          responseConsumer.accept(response);
+          return false;
+        }
+    );
+  }
+
+  protected void setSubtitleInternal(@Nullable Subtitle subtitle) {
     if (Objects.equals(this.subtitle, subtitle)) {
       return;
     }
@@ -160,6 +269,23 @@ private void setSubtitleInternal(Subtitle subtitle) {
       return;
     }
 
-    this.sendPacket(new SubtitlePacket(subtitles));
+    this.sendLabyModPacket(new SubtitlePacket(subtitles));
+  }
+
+  private void sendLabyModPacket(@NotNull Packet packet) {
+    this.protocolService.labyModProtocol.sendPacket(this.uniqueId, packet);
+  }
+
+  private <T extends IdentifiablePacket> void sendLabyModPacket(
+      @NotNull IdentifiablePacket packet,
+      @NotNull Class<T> responseClass,
+      @NotNull Predicate<T> responseConsumer
+  ) {
+    this.protocolService.labyModProtocol.sendPacket(
+        this.uniqueId,
+        packet,
+        responseClass,
+        responseConsumer
+    );
   }
 }