From 288ead3e9b03e8fef62eaa89efdf884559a8eedd Mon Sep 17 00:00:00 2001
From: Kodi Craft <kodi@kdcf.me>
Date: Fri, 4 Oct 2024 10:33:41 +0200
Subject: [PATCH] Implement rudimentary keepalive

---
 apps/amethyst/lib/states/configuration.ex |  7 ++++++-
 apps/amethyst/lib/states/play.ex          | 24 +++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/apps/amethyst/lib/states/configuration.ex b/apps/amethyst/lib/states/configuration.ex
index 47f19b2..a8bb380 100644
--- a/apps/amethyst/lib/states/configuration.ex
+++ b/apps/amethyst/lib/states/configuration.ex
@@ -269,6 +269,7 @@ defmodule Amethyst.ConnectionState.Configuration do
   def handle(%{packet_type: :acknowledge_finish_configuration}, 767, state) do
     Logger.debug("Received acknowledge finish configuration")
     send(self(), {:set_state, Amethyst.ConnectionState.Play})
+
     game = Application.fetch_env!(:amethyst, :default_game) |> Amethyst.GameCoordinator.find_or_create()
     state = state |> Map.put(:game, game)
     login = Amethyst.Game.login(game, state)
@@ -309,7 +310,11 @@ defmodule Amethyst.ConnectionState.Configuration do
           chunk_z: div(floor(z), 16)
         }})
         send(self(), {:set_position, {x, y, z}})
-        state
+        # Begin keepalive loop
+        # TODO: Put it under some supervisor
+        me = self()
+        pid = spawn(fn -> Amethyst.ConnectionState.Play.keepalive_loop(me) end)
+        state |> Map.put(:keepalive, pid)
     end
   end
 
diff --git a/apps/amethyst/lib/states/play.ex b/apps/amethyst/lib/states/play.ex
index cd4d824..69c4855 100644
--- a/apps/amethyst/lib/states/play.ex
+++ b/apps/amethyst/lib/states/play.ex
@@ -24,6 +24,7 @@ defmodule Amethyst.ConnectionState.Play do
   This module contains the packets and logic for the Play state.
   """
   Macros.defpacket_clientbound :disconnect, 0x1D, 767, [reason: :nbt]
+  Macros.defpacket_clientbound :keep_alive, 0x26, 767, [id: :long]
   Macros.defpacket_clientbound :chunk_data_and_update_light, 0x27, 767, [
     chunk_x: :int,
     chunk_z: :int,
@@ -157,6 +158,7 @@ defmodule Amethyst.ConnectionState.Play do
 
   Macros.defpacket_serverbound :confirm_teleportation, 0x00, 767, [teleport_id: :varint]
   Macros.defpacket_serverbound :serverbound_plugin_message, 0x12, 767, [channel: :string, data: :raw]
+  Macros.defpacket_serverbound :keep_alive, 0x18, 767, [id: :long]
   Macros.defpacket_serverbound :set_player_position, 0x1A, 767, [
     x: :double,
     feet_y: :double,
@@ -231,6 +233,28 @@ defmodule Amethyst.ConnectionState.Play do
     end
   end
 
+  def handle(%{packet_type: :keep_alive, id: id}, 767, state) do
+    ka = state |> Map.get(:keepalive)
+    send(ka, {:respond, id})
+    :ok
+  end
+  # This function should be started on a new task under the connection handler
+  # and is responsible for keepalive logic.
+  def keepalive_loop(player) do
+    Process.link(player) # Is it fine to do this on loop?
+    <<id::32>> = :rand.bytes(4)
+    Logger.debug("Sending keepalive...")
+    send(player, {:send_packet, %{packet_type: :keep_alive, id: id}})
+    receive do
+      {:respond, ^id} ->
+        :timer.sleep(250)
+        keepalive_loop(player)
+    after
+      15_000 ->
+        send(player, {:disconnect, "Timed out! Connection overloaded?"})
+    end
+  end
+
   def disconnect(reason) do
     %{
       packet_type: :disconnect,