Encrypt all packets except for first two packets

api/src/main/java/net/azisaba/azipluginmessaging/api/AziPluginMessaging.java

@@ -1,10 +1,12 @@
 package net.azisaba.azipluginmessaging.api;
 
+import net.azisaba.azipluginmessaging.api.entity.Player;
 import net.azisaba.azipluginmessaging.api.entity.PlayerAdapter;
 import net.azisaba.azipluginmessaging.api.server.PacketSender;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.Optional;
+import java.util.UUID;
 
 public interface AziPluginMessaging {
     /**
@@ -30,6 +32,9 @@ public interface AziPluginMessaging {
     @NotNull
     Server getServer();
 
+    @NotNull
+    Optional<Player> getPlayer(@NotNull UUID uuid);
+
     /**
      * Returns the player adapter for class.
      * <p>Generally, class should be one of these (depending on the environment):
@@ -44,15 +49,6 @@ public interface AziPluginMessaging {
     <T> PlayerAdapter<T> getPlayerAdapter(@NotNull Class<T> clazz);
 
     interface Proxy {
-        /**
-         * Returns the packet sender instance for provided server.
-         * @param serverName the server name
-         * @return the packet sender instance
-         */
-        @NotNull
-        default Optional<PacketSender> getPacketSenderForServer(@NotNull String serverName) {
-            throw new UnsupportedOperationException("Unsupported in current environment.");
-        }
     }
 
     interface Server {

api/src/main/java/net/azisaba/azipluginmessaging/api/Logger.java

@@ -49,7 +49,7 @@ public interface Logger {
             private String format(@NotNull String msg) {
                 int length = 0;
                 while (msg.contains("{}")) {
-                    msg = msg.replace("{}", "{" + length++ + "}");
+                    msg = msg.replaceFirst("\\{}", "{" + length++ + "}");
                 }
                 return msg;
             }

api/src/main/java/net/azisaba/azipluginmessaging/api/entity/Player.java

@@ -22,6 +22,12 @@ public interface Player {
     @NotNull
     UUID getUniqueId();
 
+    /**
+     * Sends a message to the player. This method might do nothing depending on the implementation.
+     * @param message the message
+     */
+    void sendMessage(@NotNull String message);
+
     /**
      * Returns the username if present, and fallbacks to unique id.
      * @return username or unique id

api/src/main/java/net/azisaba/azipluginmessaging/api/entity/SimplePlayer.java

@@ -1,5 +1,6 @@
 package net.azisaba.azipluginmessaging.api.entity;
 
+import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -31,6 +32,16 @@ public SimplePlayer(@NotNull UUID uuid, @Nullable String username) {
         return username;
     }
 
+    /**
+     * no-op method.
+     * @param message the message
+     */
+    @Override
+    public void sendMessage(@NotNull String message) {
+        // no-op
+    }
+
+    @Contract("_ -> new")
     @NotNull
     public static SimplePlayer read(@NotNull DataInputStream in) throws IOException {
         UUID uuid = UUID.fromString(in.readUTF());

api/src/main/java/net/azisaba/azipluginmessaging/api/protocol/PacketFlow.java

@@ -0,0 +1,6 @@
+package net.azisaba.azipluginmessaging.api.protocol;
+
+public enum PacketFlow {
+    TO_PROXY,
+    TO_SERVER,
+}

api/src/main/java/net/azisaba/azipluginmessaging/api/protocol/Protocol.java

@@ -1,18 +1,29 @@
 package net.azisaba.azipluginmessaging.api.protocol;
 
 import net.azisaba.azipluginmessaging.api.Logger;
-import net.azisaba.azipluginmessaging.api.protocol.handler.GiveGamingSaraHandler;
 import net.azisaba.azipluginmessaging.api.protocol.handler.MessageHandler;
-import net.azisaba.azipluginmessaging.api.protocol.handler.SetRankHandler;
-import net.azisaba.azipluginmessaging.api.protocol.handler.ToggleGamingSaraHandler;
-import net.azisaba.azipluginmessaging.api.protocol.handler.ToggleSaraHideHandler;
-import net.azisaba.azipluginmessaging.api.protocol.handler.ToggleSaraShowHandler;
+import net.azisaba.azipluginmessaging.api.protocol.handler.ProxyMessageHandler;
+import net.azisaba.azipluginmessaging.api.protocol.handler.ProxyboundEncryptionPacket;
+import net.azisaba.azipluginmessaging.api.protocol.handler.ProxyboundGiveGamingSaraPacket;
+import net.azisaba.azipluginmessaging.api.protocol.handler.ProxyboundGiveSaraPacket;
+import net.azisaba.azipluginmessaging.api.protocol.handler.ProxyboundSetRankPacket;
+import net.azisaba.azipluginmessaging.api.protocol.handler.ProxyboundToggleGamingSaraPacket;
+import net.azisaba.azipluginmessaging.api.protocol.handler.ProxyboundToggleSaraHidePacket;
+import net.azisaba.azipluginmessaging.api.protocol.handler.ProxyboundToggleSaraShowPacket;
+import net.azisaba.azipluginmessaging.api.protocol.handler.ServerMessageHandler;
+import net.azisaba.azipluginmessaging.api.protocol.handler.ServerboundActionResponsePacket;
+import net.azisaba.azipluginmessaging.api.protocol.handler.ServerboundEncryptionPacket;
 import net.azisaba.azipluginmessaging.api.protocol.message.Message;
 import net.azisaba.azipluginmessaging.api.protocol.message.PlayerMessage;
 import net.azisaba.azipluginmessaging.api.protocol.message.PlayerWithServerMessage;
-import net.azisaba.azipluginmessaging.api.protocol.message.SetRankMessage;
+import net.azisaba.azipluginmessaging.api.protocol.message.ProxyboundGiveSaraMessage;
+import net.azisaba.azipluginmessaging.api.protocol.message.ProxyboundSetRankMessage;
+import net.azisaba.azipluginmessaging.api.protocol.message.PublicKeyMessage;
+import net.azisaba.azipluginmessaging.api.protocol.message.ServerboundActionResponseMessage;
 import net.azisaba.azipluginmessaging.api.server.PacketSender;
 import net.azisaba.azipluginmessaging.api.server.ServerConnection;
+import net.azisaba.azipluginmessaging.api.util.Constants;
+import net.azisaba.azipluginmessaging.api.util.EncryptionUtil;
 import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -27,26 +38,47 @@
 import java.util.concurrent.ConcurrentHashMap;
 
 public final class Protocol<T extends MessageHandler<M>, M extends Message> {
-    private static final Map<Byte, Protocol<?, ?>> BY_ID = new ConcurrentHashMap<>();
+    private static final Map<Byte, Protocol<?, ?>> TO_PROXY_BY_ID = new ConcurrentHashMap<>();
+    private static final Map<Byte, Protocol<?, ?>> TO_SERVER_BY_ID = new ConcurrentHashMap<>();
     public static final String LEGACY_CHANNEL_ID = "AziPluginMessaging";
     public static final String CHANNEL_ID = "azipm:main";
 
-    public static final Protocol<SetRankHandler, SetRankMessage> SET_RANK = new Protocol<>(0x00, new SetRankHandler());
-    public static final Protocol<GiveGamingSaraHandler, PlayerMessage> GIVE_GAMING_SARA = new Protocol<>(0x01, new GiveGamingSaraHandler());
-    public static final Protocol<ToggleGamingSaraHandler, PlayerMessage> TOGGLE_GAMING_SARA = new Protocol<>(0x02, new ToggleGamingSaraHandler());
-    public static final Protocol<ToggleSaraHideHandler, PlayerMessage> TOGGLE_SARA_HIDE = new Protocol<>(0x03, new ToggleSaraHideHandler()); // Note that this is non-contextual
-    public static final Protocol<ToggleSaraShowHandler, PlayerWithServerMessage> TOGGLE_SARA_SHOW = new Protocol<>(0x04, new ToggleSaraShowHandler()); // Note that this is contextual
+    public static final Protocol<ProxyboundEncryptionPacket, PublicKeyMessage> P_ENCRYPTION = new Protocol<>(PacketFlow.TO_PROXY, 0x00, new ProxyboundEncryptionPacket());
+    public static final Protocol<ProxyboundSetRankPacket, ProxyboundSetRankMessage> P_SET_RANK = new Protocol<>(PacketFlow.TO_PROXY, 0x01, new ProxyboundSetRankPacket());
+    public static final Protocol<ProxyboundGiveGamingSaraPacket, PlayerMessage> P_GIVE_GAMING_SARA = new Protocol<>(PacketFlow.TO_PROXY, 0x02, new ProxyboundGiveGamingSaraPacket());
+    public static final Protocol<ProxyboundGiveSaraPacket, ProxyboundGiveSaraMessage> P_GIVE_SARA = new Protocol<>(PacketFlow.TO_PROXY, 0x03, new ProxyboundGiveSaraPacket());
+    public static final Protocol<ProxyboundToggleGamingSaraPacket, PlayerMessage> P_TOGGLE_GAMING_SARA = new Protocol<>(PacketFlow.TO_PROXY, 0x04, new ProxyboundToggleGamingSaraPacket());
+    public static final Protocol<ProxyboundToggleSaraHidePacket, PlayerMessage> P_TOGGLE_SARA_HIDE = new Protocol<>(PacketFlow.TO_PROXY, 0x05, new ProxyboundToggleSaraHidePacket()); // Note that this is non-contextual
+    public static final Protocol<ProxyboundToggleSaraShowPacket, PlayerWithServerMessage> P_TOGGLE_SARA_SHOW = new Protocol<>(PacketFlow.TO_PROXY, 0x06, new ProxyboundToggleSaraShowPacket()); // Note that this is contextual
 
+    public static final Protocol<ServerboundEncryptionPacket, PublicKeyMessage> S_ENCRYPTION = new Protocol<>(PacketFlow.TO_SERVER, 0x00, new ServerboundEncryptionPacket());
+    public static final Protocol<ServerboundActionResponsePacket, ServerboundActionResponseMessage> S_ACTION_RESPONSE = new Protocol<>(PacketFlow.TO_SERVER, 0x01, new ServerboundActionResponsePacket());
+
+    private final PacketFlow packetFlow;
     private final byte id;
     private final T handler;
 
-    private Protocol(int id, @NotNull T handler) {
+    private Protocol(@NotNull PacketFlow packetFlow, int id, @NotNull T handler) {
+        this.packetFlow = packetFlow;
         this.id = (byte) (id & 0xFF);
         this.handler = handler;
-        if (BY_ID.containsKey(this.id)) {
-            throw new AssertionError("Duplicate protocol id: " + this.id);
+        if (packetFlow == PacketFlow.TO_PROXY) {
+            if (!(handler instanceof ProxyMessageHandler)) {
+                throw new IllegalArgumentException("Handler must be instance of ProxyMessageHandler");
+            }
+            if (TO_PROXY_BY_ID.containsKey(this.id)) {
+                throw new AssertionError("Duplicate protocol id: " + this.id);
+            }
+            TO_PROXY_BY_ID.put(this.id, this);
+        } else {
+            if (!(handler instanceof ServerMessageHandler)) {
+                throw new IllegalArgumentException("Handler must be instance of ServerMessageHandler");
+            }
+            if (TO_SERVER_BY_ID.containsKey(this.id)) {
+                throw new AssertionError("Duplicate protocol id: " + this.id);
+            }
+            TO_SERVER_BY_ID.put(this.id, this);
         }
-        BY_ID.put(this.id, this);
     }
 
     /**
@@ -68,17 +100,29 @@ public T getHandler() {
 
     /**
      * Attempt to send a packet.
-     * @param packetSender the packet sender to send the packet from.
+     * @param sender the packet sender to send the packet from.
      * @param msg the message to send
      * @return true if the message was sent successfully, false otherwise.
      */
-    public boolean sendPacket(@NotNull PacketSender packetSender, @NotNull M msg) {
+    public boolean sendPacket(@NotNull PacketSender sender, @NotNull M msg) {
         try (ByteArrayOutputStream bout = new ByteArrayOutputStream();
              DataOutputStream out = new DataOutputStream(bout)) {
             out.writeByte(id);
             msg.write(out);
             byte[] bytes = bout.toByteArray();
-            return packetSender.sendPacket(bytes);
+            if (sender.isEncrypted()) {
+                try {
+                    bytes = EncryptionUtil.encrypt(bytes, sender.getRemotePublicKey());
+                } catch (Exception e) {
+                    throw new RuntimeException("Could not encrypt the packet (sender: " + sender + ")", e);
+                }
+            }
+            if (Constants.DEBUG) {
+                String hex = Integer.toString(id, 16);
+                if (hex.length() == 1) hex = '0' + hex;
+                Logger.getCurrentLogger().info("Sending packet {} (0x{}) to {} (encrypted: {})", id, hex, sender, sender.isEncrypted());
+            }
+            return sender.sendPacket(bytes);
         } catch (IOException e) {
             Logger.getCurrentLogger().warn("Failed to send packet", e);
             return false;
@@ -88,42 +132,106 @@ public boolean sendPacket(@NotNull PacketSender packetSender, @NotNull M msg) {
     /**
      * This method is called when a packet is received (proxy-side).
      * @param server the server connection
-     * @param data the data of the packet
+     * @param rawData the data of the packet
      */
-    public static void handleProxySide(ServerConnection server, byte[] data) {
+    public static void handleProxySide(ServerConnection server, byte[] rawData) {
+        byte[] data;
+        if (server.isEncrypted()) {
+            try {
+                data = EncryptionUtil.decrypt(rawData, server.getKeyPair().getPrivate());
+            } catch (Exception e) {
+                throw new RuntimeException("Could not decrypt the packet (server: " + server + ")", e);
+            }
+        } else {
+            data = rawData;
+        }
         try (ByteArrayInputStream bin = new ByteArrayInputStream(data);
              DataInputStream in = new DataInputStream(bin)) {
             byte id = (byte) (in.readByte() & 0xFF);
-            Protocol<?, ?> protocol = Protocol.getById(id);
+            Protocol<?, ?> protocol = Protocol.getById(PacketFlow.TO_PROXY, id);
             if (protocol == null) {
                 Logger.getCurrentLogger().warn(
                         "Received unknown protocol id from server connection {}: {}",
                         server, id);
                 return;
             }
-            String hex = Integer.toString(id, 16);
-            if (hex.length() == 1) hex = '0' + hex;
-            Logger.getCurrentLogger().info("Received packet {} (0x{}) from server connection {}",
-                    id, hex, server);
+            if (Constants.DEBUG) {
+                String hex = Integer.toString(id, 16);
+                if (hex.length() == 1) hex = '0' + hex;
+                Logger.getCurrentLogger().info("Received packet {} (0x{}) from server connection {} (encrypted: {})", id, hex, server, server.isEncrypted());
+            }
+            if (protocol.packetFlow != PacketFlow.TO_PROXY) {
+                throw new IllegalArgumentException("Packet " + protocol + " is not proxybound");
+            }
             @SuppressWarnings("unchecked")
-            MessageHandler<Message> handler = (MessageHandler<Message>) protocol.getHandler();
+            ProxyMessageHandler<Message> handler = (ProxyMessageHandler<Message>) protocol.getHandler();
             Message message = handler.read(server, in);
-            handler.handle(message);
+            handler.handle(server, message);
         } catch (Exception | AssertionError e) {
             Logger.getCurrentLogger().warn(
                     "Failed to handle plugin message from server connection {} (player: {})",
                     server.getServerInfo().getName(), server.getPlayer().getUniqueId(), e);
         }
     }
 
+    /**
+     * This method is called when a packet is received (server-side).
+     * @param rawData the data of the packet
+     */
+    public static void handleServerSide(@NotNull PacketSender sender, byte[] rawData) {
+        byte[] data;
+        if (sender.isEncrypted()) {
+            try {
+                data = EncryptionUtil.decrypt(rawData, sender.getKeyPair().getPrivate());
+            } catch (Exception e) {
+                throw new RuntimeException("Could not decrypt the packet (sender: " + sender + ")", e);
+            }
+        } else {
+            data = rawData;
+        }
+        try (ByteArrayInputStream bin = new ByteArrayInputStream(data);
+             DataInputStream in = new DataInputStream(bin)) {
+            byte id = (byte) (in.readByte() & 0xFF);
+            Protocol<?, ?> protocol = Protocol.getById(PacketFlow.TO_SERVER, id);
+            if (protocol == null) {
+                Logger.getCurrentLogger().warn("Received unknown protocol id from {}: {}", sender, id);
+                return;
+            }
+            if (Constants.DEBUG) {
+                String hex = Integer.toString(id, 16);
+                if (hex.length() == 1) hex = '0' + hex;
+                Logger.getCurrentLogger().info("Received packet {} (0x{}) from {}", id, hex, sender);
+            }
+            if (protocol.packetFlow != PacketFlow.TO_SERVER) {
+                throw new IllegalArgumentException("Packet " + protocol + " is not serverbound");
+            }
+            @SuppressWarnings("unchecked")
+            ServerMessageHandler<Message> handler = (ServerMessageHandler<Message>) protocol.getHandler();
+            Message message = handler.read(in);
+            handler.handle(sender, message);
+        } catch (Exception | AssertionError e) {
+            Logger.getCurrentLogger().warn("Failed to handle plugin message from proxy", e);
+        }
+    }
+
+    @NotNull
+    @Contract(pure = true)
+    public PacketFlow getPacketFlow() {
+        return packetFlow;
+    }
+
     /**
      * Gets the protocol (packet) by its id.
      * @param id the id
      * @return the protocol, or null if not found
      */
     @Nullable
-    public static Protocol<?, ?> getById(byte id) {
-        return BY_ID.get(id);
+    public static Protocol<?, ?> getById(@NotNull PacketFlow packetFlow, byte id) {
+        if (packetFlow == PacketFlow.TO_PROXY) {
+            return TO_PROXY_BY_ID.get(id);
+        } else {
+            return TO_SERVER_BY_ID.get(id);
+        }
     }
 
     /**
@@ -132,8 +240,19 @@ public static void handleProxySide(ServerConnection server, byte[] data) {
      */
     @Contract(pure = true)
     @NotNull
-    public static Collection<Protocol<?, ?>> values() {
-        return BY_ID.values();
+    public static Collection<Protocol<?, ?>> values(@NotNull PacketFlow packetFlow) {
+        if (packetFlow == PacketFlow.TO_PROXY) {
+            return TO_PROXY_BY_ID.values();
+        } else {
+            return TO_SERVER_BY_ID.values();
+        }
     }
-}
 
+    @Override
+    public String toString() {
+        return "Protocol{" +
+                "id=" + id +
+                ", handler=" + handler +
+                '}';
+    }
+}

api/src/main/java/net/azisaba/azipluginmessaging/api/protocol/handler/MessageHandler.java

@@ -1,31 +1,16 @@
 package net.azisaba.azipluginmessaging.api.protocol.handler;
 
 import net.azisaba.azipluginmessaging.api.protocol.message.Message;
-import net.azisaba.azipluginmessaging.api.server.ServerConnection;
+import net.azisaba.azipluginmessaging.api.server.PacketSender;
 import org.jetbrains.annotations.NotNull;
 
-import java.io.DataInputStream;
-import java.io.IOException;
-
-/**
- * Represents a message handler which is responsible for decoding and handling the message.
- * @param <T> the message type
- */
 public interface MessageHandler<T extends Message> {
-    /**
-     * Decodes the message from the stream.
-     * @param server The server connection.
-     * @param in The stream.
-     * @return The decoded message.
-     * @throws IOException If an I/O error occurs.
-     */
-    @NotNull
-    T read(@NotNull ServerConnection server, @NotNull DataInputStream in) throws IOException;
-
     /**
      * Handles the (decoded) message.
-     * @param msg The message.
+     *
+     * @param sender The sender of the message.
+     * @param msg    The message.
      * @throws Exception If an exception occurs.
      */
-    void handle(@NotNull T msg) throws Exception;
+    void handle(@NotNull PacketSender sender, @NotNull T msg) throws Exception;
 }