Encrypt the entire packet stream (yes that means proxy packets are en…

…crypted twice)
Gaming32 · Aug 28, 2024 · d2453be · d2453be
1 parent addb8d7
commit d2453be
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 55 deletions.
diff --git a/src/main/java/io/github/gaming32/worldhost/protocol/ProtocolClient.java b/src/main/java/io/github/gaming32/worldhost/protocol/ProtocolClient.java
@@ -19,8 +19,9 @@
 import org.jetbrains.annotations.Nullable;
 
 import javax.crypto.Cipher;
+import javax.crypto.CipherInputStream;
+import javax.crypto.CipherOutputStream;
 import javax.crypto.SecretKey;
-import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
@@ -130,21 +131,20 @@ public ProtocolClient(String host, boolean successToast, boolean failureToast) {
 
             final Thread sendThread = SEND_THREAD_BUILDER.start(() -> {
                 try {
-                    final DataOutputStream dos = new DataOutputStream(fSocket.getOutputStream());
+                    final DataOutputStream dos = new DataOutputStream(
+                        new CipherOutputStream(fSocket.getOutputStream(), fEncryptCipher)
+                    );
                     final ByteArrayOutputStream baos = new ByteArrayOutputStream();
                     final DataOutputStream tempDos = new DataOutputStream(baos);
                     while (!closed) {
                         final var optionalMessage = sendQueue.take();
                         if (optionalMessage.isEmpty()) break;
                         final var message = optionalMessage.get();
                         message.encode(tempDos);
-                        final byte[] data = message.isEncrypted()
-                            ? fEncryptCipher.update(baos.toByteArray())
-                            : baos.toByteArray();
-                        baos.reset();
-                        dos.writeInt(data.length + 1);
+                        dos.writeInt(baos.size() + 1);
                         dos.writeByte(message.typeId() & 0xff);
-                        dos.write(data);
+                        dos.write(baos.toByteArray());
+                        baos.reset();
                         dos.flush();
                     }
                 } catch (IOException e) {
@@ -157,24 +157,19 @@ public ProtocolClient(String host, boolean successToast, boolean failureToast) {
 
             RECV_THREAD_BUILDER.start(() -> {
                 try {
-                    final DataInputStream dis = new DataInputStream(fSocket.getInputStream());
+                    final DataInputStream dis = new DataInputStream(
+                        new CipherInputStream(fSocket.getInputStream(), fDecryptCipher)
+                    );
                     while (!closed) {
                         final int length = dis.readInt() - 1;
                         if (length < 0) {
                             WorldHost.LOGGER.warn("Received invalid empty packet from WH server");
                             continue;
                         }
                         final int typeId = dis.readUnsignedByte();
-                        final CountingInputStream cis;
-                        if (WorldHostS2CMessage.isEncrypted(typeId)) {
-                            final byte[] data = dis.readNBytes(length);
-                            final byte[] decrypted = fDecryptCipher.update(data);
-                            cis = new CountingInputStream(new ByteArrayInputStream(decrypted));
-                        } else {
-                            final BoundedInputStream bis = new BoundedInputStream(dis, length);
-                            bis.setPropagateClose(false);
-                            cis = new CountingInputStream(bis);
-                        }
+                        final BoundedInputStream bis = new BoundedInputStream(dis, length);
+                        bis.setPropagateClose(false);
+                        final var cis = new CountingInputStream(bis);
                         WorldHostS2CMessage message = null;
                         try {
                             message = WorldHostS2CMessage.decode(typeId, new DataInputStream(cis));

diff --git a/src/main/java/io/github/gaming32/worldhost/protocol/WorldHostC2SMessage.java b/src/main/java/io/github/gaming32/worldhost/protocol/WorldHostC2SMessage.java
@@ -192,11 +192,6 @@ public void encode(DataOutputStream dos) throws IOException {
             writeString(dos, myHost);
             dos.writeShort(myPort);
         }
-
-        @Override
-        public boolean isEncrypted() {
-            return true;
-        }
     }
 
     record PunchFailed(long targetConnection, UUID punchId) implements WorldHostC2SMessage {
@@ -209,11 +204,6 @@ public byte typeId() {
         public void encode(DataOutputStream dos) throws IOException {
             writeUuid(dos, punchId);
         }
-
-        @Override
-        public boolean isEncrypted() {
-            return true;
-        }
     }
 
     record BeginPortLookup(UUID lookupId) implements WorldHostC2SMessage {
@@ -226,11 +216,6 @@ public byte typeId() {
         public void encode(DataOutputStream dos) throws IOException {
             writeUuid(dos, lookupId);
         }
-
-        @Override
-        public boolean isEncrypted() {
-            return true;
-        }
     }
 
     record PunchSuccess(long connectionId, UUID punchId, String host, int port) implements WorldHostC2SMessage {
@@ -246,22 +231,12 @@ public void encode(DataOutputStream dos) throws IOException {
             writeString(dos, host);
             dos.writeShort(port);
         }
-
-        @Override
-        public boolean isEncrypted() {
-            return true;
-        }
     }
 
     byte typeId();
 
     void encode(DataOutputStream dos) throws IOException;
 
-    // TODO: Rewrite encryption to encrypt whole stream
-    default boolean isEncrypted() {
-        return false;
-    }
-
     static void writeUuid(DataOutputStream dos, UUID uuid) throws IOException {
         dos.writeLong(uuid.getMostSignificantBits());
         dos.writeLong(uuid.getLeastSignificantBits());

diff --git a/src/main/java/io/github/gaming32/worldhost/protocol/WorldHostS2CMessage.java b/src/main/java/io/github/gaming32/worldhost/protocol/WorldHostS2CMessage.java
@@ -564,17 +564,6 @@ public void handle(ProtocolClient client) {
      */
     void handle(ProtocolClient client);
 
-    static boolean isEncrypted(int typeId) {
-        return switch (typeId) {
-            case PunchOpenRequest.ID,
-                 CancelPortLookup.ID,
-                 PortLookupSuccess.ID,
-                 PunchRequestCancelled.ID,
-                 PunchSuccess.ID -> true;
-            default -> false;
-        };
-    }
-
     static WorldHostS2CMessage decode(int typeId, DataInputStream dis) throws IOException {
         return switch (typeId) {
             case Error.ID -> Error.decode(dis);