From 0a68dacccc8e57001ce9a645b23fd2242b32e73a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Noel=20N=C3=A9meth?= <noel@noelnemeth.com>
Date: Sun, 12 Jun 2022 15:15:07 +0200
Subject: [PATCH] Add util methods

---
 .../client/play/ClientCommandChatPacket.java  | 28 +++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/src/main/java/net/minestom/server/network/packet/client/play/ClientCommandChatPacket.java b/src/main/java/net/minestom/server/network/packet/client/play/ClientCommandChatPacket.java
index 600e3fd12..15bc40c38 100644
--- a/src/main/java/net/minestom/server/network/packet/client/play/ClientCommandChatPacket.java
+++ b/src/main/java/net/minestom/server/network/packet/client/play/ClientCommandChatPacket.java
@@ -1,12 +1,19 @@
 package net.minestom.server.network.packet.client.play;
 
+import net.minestom.server.command.builder.arguments.Argument;
+import net.minestom.server.command.builder.arguments.minecraft.SignableArgument;
+import net.minestom.server.crypto.MessageSignature;
 import net.minestom.server.network.packet.client.ClientPacket;
 import net.minestom.server.utils.binary.BinaryReader;
 import net.minestom.server.utils.binary.BinaryWriter;
+import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
+import java.time.Instant;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.UUID;
 
 public record ClientCommandChatPacket(@NotNull String message, long timestamp,
                                       long salt, Map<String, byte[]> signatures,
@@ -20,6 +27,27 @@ public record ClientCommandChatPacket(@NotNull String message, long timestamp,
                 reader.readLong(), readSignature(reader), reader.readBoolean());
     }
 
+    /**
+     * Constructs a new {@link MessageSignature} to use in validation
+     *
+     * @param signer player's uuid who sent this packet
+     * @return null if the client didn't sign the parameter
+     */
+    @Contract("_, _ -> new")
+    public @Nullable MessageSignature signatureOf(String parameterName, UUID signer) {
+        final byte[] signature = signatures.get(parameterName);
+        return signature == null ? null : new MessageSignature(signer, Instant.ofEpochMilli(timestamp), salt, signature);
+    }
+
+    /**
+     * @see #signatureOf(String, UUID)
+     */
+    @SuppressWarnings("ConstantConditions")
+    @Contract("_, _ -> new")
+    public <T extends Argument<?> & SignableArgument> @Nullable MessageSignature signatureOf(T argument, UUID signer) {
+        return signatureOf(argument.getId(), signer);
+    }
+
     @Override
     public void write(@NotNull BinaryWriter writer) {
         writer.writeSizedString(message);