From 7f3bb357db6f8e16f25daeb55b094f9e52350362 Mon Sep 17 00:00:00 2001
From: Kodi Craft <kodi@kdcf.me>
Date: Sat, 5 Oct 2024 11:55:27 +0200
Subject: [PATCH] Implement encryption

---
 apps/amethyst/lib/apps/connection_handler.ex  | 12 ++---
 apps/amethyst/lib/apps/connection_receiver.ex | 52 +++++++++++++------
 2 files changed, 42 insertions(+), 22 deletions(-)

diff --git a/apps/amethyst/lib/apps/connection_handler.ex b/apps/amethyst/lib/apps/connection_handler.ex
index 2d8d07c..2d8b880 100644
--- a/apps/amethyst/lib/apps/connection_handler.ex
+++ b/apps/amethyst/lib/apps/connection_handler.ex
@@ -107,7 +107,11 @@ defmodule Amethyst.ConnectionHandler do
         send_packet(socket, connstate, packet, version, state)
         loop(socket, connstate, version, state)
       after 0 ->
+        # Received stuff from the connection receiver is lower priority
         receive do
+          {:get_encryption, from} ->
+            send(from, Map.get(state, :decryption_state))
+            loop(socket, connstate, version, state)
           {:packet, id, data} ->
             state = handle_packet(id, data, connstate, version, state)
             loop(socket, connstate, version, state)
@@ -216,13 +220,7 @@ defmodule Amethyst.ConnectionHandler do
 
   defp handle_packet(id, data, connstate, version, state) do
     try do
-      packet = case Map.get(state, :decryption_state) do
-        nil ->
-          connstate.deserialize(id, version, data)
-        dstate ->
-          data = :crypto.crypto_update(dstate, data)
-          connstate.deserialize(id, version, data)
-      end
+      packet = connstate.deserialize(id, version, data)
       case connstate.handle(packet, version, state) do
         :ok -> state
         {:error, reason} ->
diff --git a/apps/amethyst/lib/apps/connection_receiver.ex b/apps/amethyst/lib/apps/connection_receiver.ex
index ea11a77..fd34386 100644
--- a/apps/amethyst/lib/apps/connection_receiver.ex
+++ b/apps/amethyst/lib/apps/connection_receiver.ex
@@ -35,7 +35,7 @@ defmodule Amethyst.ConnectionReceiver do
     {:ok, spawn(fn ->
       Process.set_label("ConnectionReceiver for #{inspect(socket)}")
       {:ok, pid} = Amethyst.ConnectionHandler.start_link(socket, Amethyst.ConnectionState.Handshake, 0)
-      receive(socket, pid)
+      receive(socket, pid, nil)
     end)}
   end
 
@@ -44,42 +44,64 @@ defmodule Amethyst.ConnectionReceiver do
     {:ok, spawn_link(fn ->
       Process.set_label("ConnectionReceiver for #{inspect(socket)}")
       {:ok, pid} = Amethyst.ConnectionHandler.start_link(socket, Amethyst.ConnectionState.Handshake, 0)
-      receive(socket, pid)
+      receive(socket, pid, nil)
     end)}
   end
 
-  @spec receive(:gen_tcp.socket(), pid()) :: no_return()
-  def receive(socket, sender) do
-    case get_packet(socket) do
+  @spec receive(:gen_tcp.socket(), pid(), nil | :crypto.crypto_state()) :: no_return()
+  def receive(socket, sender, cstate) do
+    case get_packet(socket, cstate) do
       :closed -> send(sender, :closed)
         Process.exit(self(), :normal)
       {:error, error} -> Logger.error("Error reading packet: #{error}")
       {id, data} -> send(sender, {:packet, id, data})
     end
-    receive(socket, sender)
+    if cstate == nil do
+      # Ask the handler if the encryption state has changed
+      send(sender, {:get_encryption, self()})
+      Logger.debug("Asking for news on encryption...")
+      receive do
+        nil -> receive(socket, sender, cstate)
+        some ->
+          Logger.debug("Enabling decryption!")
+          receive(socket, sender, some)
+      end
+    else
+     receive(socket, sender, cstate)
+    end
   end
 
-  def get_packet(client) do
-    case get_varint(client, "") do
+  def get_packet(client, cstate) do
+    case get_varint(client, "", cstate) do
       :closed -> :closed
       {:error, error} -> {:error, error}
       {[length], ""} ->
         recv = :gen_tcp.recv(client, length)
         case recv do
-          {:ok, full_packet} -> ({[id], data} = Read.start(full_packet) |> Read.varint() |> Read.stop()
-                                {id, data})
+          {:ok, full_packet} ->
+            full_packet = case cstate do
+              nil -> full_packet
+              ds -> :crypto.crypto_update(ds, full_packet)
+            end
+            ({[id], data} = Read.start(full_packet) |> Read.varint() |> Read.stop()
+              {id, data})
           {:error, :closed} -> :closed
           {:error, error} -> {:error, error}
         end
     end
   end
 
-  defp get_varint(client, acc) do
+  defp get_varint(client, acc, cstate) do
     case :gen_tcp.recv(client, 1) do
-      {:ok, byte} -> case byte do
-        <<0::1, _::7>> -> Read.start(acc <> byte) |> Read.varint() |> Read.stop()
-        <<1::1, _::7>> -> get_varint(client, acc <> byte)
-      end
+      {:ok, byte} ->
+        byte = case cstate do
+          nil -> byte
+          ds -> :crypto.crypto_update(ds, byte)
+        end
+        case byte do
+          <<0::1, _::7>> -> Read.start(acc <> byte) |> Read.varint() |> Read.stop()
+          <<1::1, _::7>> -> get_varint(client, acc <> byte, cstate)
+        end
       {:error, :closed} -> :closed
       {:error, error} -> {:error, error}
     end