From be449164613d95fadf3fbf540daedb718b295949 Mon Sep 17 00:00:00 2001
From: Kodi Craft <kodi@kdcf.me>
Date: Sun, 6 Oct 2024 15:20:39 +0200
Subject: [PATCH] Implement compression

---
 apps/amethyst/lib/apps/connection_handler.ex  | 37 +++++++++----------
 apps/amethyst/lib/apps/connection_receiver.ex | 17 +++------
 2 files changed, 24 insertions(+), 30 deletions(-)

diff --git a/apps/amethyst/lib/apps/connection_handler.ex b/apps/amethyst/lib/apps/connection_handler.ex
index f19f025..60968d8 100644
--- a/apps/amethyst/lib/apps/connection_handler.ex
+++ b/apps/amethyst/lib/apps/connection_handler.ex
@@ -258,27 +258,26 @@ defmodule Amethyst.ConnectionHandler do
   defp send_packet(socket, connstate, packet, version, state) do
     try do
       data = connstate.serialize(packet, version)
-      length = byte_size(data) |> Amethyst.Minecraft.Write.varint()
-      to_write = case Map.get(state, :compression_threshold) do
-        nil ->
-          # No compression
-          length <> data
-        threshold when byte_size(length <> data) < threshold ->
-          # Compression unnecessary
-          <<0>> <> length <> data
-        threshold when byte_size(length <> data) >= threshold ->
-          # Compression
-          cdata = :zlib.compress(data)
-          outer_length = byte_size(length <> cdata)
-          Write.varint(outer_length) <> length <> cdata
+      container = if Map.get(state, :compression) == nil do
+        # Packet ID is included in data
+        Write.varint(byte_size(data)) <> data
+      else
+        threshold = Map.get(state, :compression, 0)
+        data_length = byte_size(data)
+        if data_length >= threshold do
+          compressed = Write.varint(data_length) <> :zlib.zip(data)
+          Write.varint(byte_size(compressed)) <> compressed
+        else
+          compressed = Write.varint(0) <> data
+          Write.varint(byte_size(compressed)) <> compressed
+        end
       end
-      case Map.get(state, :encryption_state) do
-        nil ->
-          :gen_tcp.send(socket, to_write)
-        estate ->
-          encrypted = :crypto.crypto_update(estate, to_write)
-          :gen_tcp.send(socket, encrypted)
+      encrypted = if Map.get(state, :encryption_state) == nil do
+        container
+      else
+        Map.get(state, :encryption_state) |> :crypto.crypto_update(container)
       end
+      :gen_tcp.send(socket, encrypted)
     rescue
       e ->
         Logger.error("Error sending packet #{inspect(packet)} in state #{connstate}: #{Exception.format(:error, e, __STACKTRACE__)}")
diff --git a/apps/amethyst/lib/apps/connection_receiver.ex b/apps/amethyst/lib/apps/connection_receiver.ex
index d8bd08e..08c591a 100644
--- a/apps/amethyst/lib/apps/connection_receiver.ex
+++ b/apps/amethyst/lib/apps/connection_receiver.ex
@@ -60,25 +60,20 @@ defmodule Amethyst.ConnectionReceiver do
     estate = if estate == nil do
       # Ask the handler if we have encryption now
       send(sender, {:get_encryption, self()})
-      Logger.debug("Asking for news on encryption...")
       receive do
         nil -> nil
-        some ->
-          Logger.debug("Enabled decryption")
-          some
+        some -> some
       end
-    else nil end
+    else estate end
+
     cstate = if cstate == nil do
       # Ask the handler if we have encryption now
       send(sender, {:get_compression, self()})
-      Logger.debug("Asking for news on compression...")
       receive do
         nil -> nil
-        some ->
-          Logger.debug("Enabled decompression")
-          some
+        some -> some
       end
-    else nil end
+    else cstate end
 
     receive(socket, sender, estate, cstate)
   end
@@ -109,7 +104,7 @@ defmodule Amethyst.ConnectionReceiver do
                 rest
               else
                 # Compressed data
-                rest |> :zlib.uncompress()
+                rest |> :zlib.unzip()
               end
             end