Introduce 'Connection' struct for managing client/server state

Author: thecodeboss

lib/minecraft/connection.ex

@@ -0,0 +1,149 @@
+defmodule Minecraft.Connection do
+  @moduledoc """
+  Maintains the state of a client's connection, and provides utilities for sending and receiving
+  data. It is designed to be chained in a fashion similar to [`Plug`](https://hexdocs.pm/plug/).
+  """
+  alias Minecraft.Packet
+  require Logger
+
+  @typedoc """
+  The possible states a client/server can be in.
+  """
+  @type state :: :handshake | :status | :login | :play
+
+  @typedoc """
+  Allowed ranch transport types.
+  """
+  @type transport :: :ranch_tcp
+
+  @type t :: %__MODULE__{
+          current_state: state,
+          socket: port | nil,
+          transport: transport | nil,
+          client_ip: String.t(),
+          data: binary | nil,
+          error: any,
+          protocol_version: integer | nil
+        }
+
+  defstruct current_state: nil,
+            socket: nil,
+            transport: nil,
+            client_ip: nil,
+            data: nil,
+            error: nil,
+            protocol_version: nil
+
+  @doc """
+  Closes the `Connection`.
+  """
+  @spec close(t) :: t
+  def close(conn) do
+    :ok = conn.transport.close(conn.socket)
+    %__MODULE__{conn | socket: nil, transport: nil}
+  end
+
+  @doc """
+  Continues receiving messages from the client.
+
+  To prevent a client from flooding our process mailbox, we only receive one message at a time,
+  and explicitly `continue` to receive messages once we finish processing the ones we have.
+  """
+  @spec continue(t) :: t
+  def continue(conn) do
+    :ok = conn.transport.setopts(conn.socket, active: :once)
+    conn
+  end
+
+  @doc """
+  Initializes a `Connection`.
+  """
+  @spec init(port(), transport()) :: t
+  def init(socket, transport) do
+    {:ok, {client_ip, _port}} = :inet.peername(socket)
+    client_ip = IO.iodata_to_binary(:inet.ntoa(client_ip))
+    :ok = transport.setopts(socket, active: :once)
+
+    Logger.info(fn -> "Client #{client_ip} connected." end)
+
+    %__MODULE__{
+      current_state: :handshake,
+      socket: socket,
+      transport: transport,
+      client_ip: client_ip,
+      data: ""
+    }
+  end
+
+  @doc """
+  Stores data received from the client in this `Connection`.
+  """
+  @spec put_data(t, binary) :: t
+  def put_data(conn, data) do
+    %__MODULE__{conn | data: conn.data <> data}
+  end
+
+  @doc """
+  Puts the `Connection` into the given `error` state.
+  """
+  @spec put_error(t, any) :: t
+  def put_error(conn, error) do
+    %__MODULE__{conn | error: error}
+  end
+
+  @doc """
+  Sets the protocol for the `Connection`.
+  """
+  @spec put_protocol(t, integer) :: t
+  def put_protocol(conn, protocol_version) do
+    %__MODULE__{conn | protocol_version: protocol_version}
+  end
+
+  @doc """
+  Replaces the `Connection`'s underlying socket.
+  """
+  @spec put_socket(t, port()) :: t
+  def put_socket(conn, socket) do
+    %__MODULE__{conn | socket: socket}
+  end
+
+  @doc """
+  Updates the `Connection` state.
+  """
+  @spec put_state(t, state) :: t
+  def put_state(conn, state) do
+    %__MODULE__{conn | current_state: state}
+  end
+
+  @doc """
+  Pops a packet from the `Connection`.
+  """
+  @spec read_packet(t) :: {:ok, struct, t} | {:error, t}
+  def read_packet(conn) do
+    case Packet.deserialize(conn.data, conn.current_state) do
+      {packet, rest} when is_binary(rest) ->
+        Logger.debug(fn -> "REQUEST: #{inspect(packet)}" end)
+        {:ok, packet, %__MODULE__{conn | data: rest}}
+
+      {:error, :invalid_packet} ->
+        Logger.error(fn ->
+          "Received an invalid packet from client, closing connection. #{inspect(conn.data)}"
+        end)
+
+        {:error, put_error(conn, :invalid_packet)}
+    end
+  end
+
+  @doc """
+  Sends a response to the client.
+  """
+  @spec send_response(t, struct) :: t
+  def send_response(conn, response) do
+    Logger.debug(fn -> "RESPONSE: #{inspect(response)}" end)
+
+    {:ok, response} = Packet.serialize(response)
+
+    :ok = conn.transport.send(conn.socket, response)
+    conn
+  end
+end

lib/minecraft/packet.ex

@@ -15,10 +15,10 @@ defmodule Minecraft.Packet do
   Given a raw binary packet, deserializes it into a `Packet` struct.
   """
   @spec deserialize(binary, state :: atom, type :: :client | :server) ::
-          {packet :: term, new_state :: atom, rest :: binary} | {:error, :invalid_packet}
+          {packet :: term, rest :: binary} | {:error, :invalid_packet}
   def deserialize(data, state, type \\ :client)
 
-  def deserialize(data, :handshaking, type) when is_binary(data) do
+  def deserialize(data, :handshake, type) when is_binary(data) do
     {_packet_size, data} = decode_varint(data)
     {packet_id, data} = decode_varint(data)
     Handshake.deserialize(packet_id, data, type)

lib/minecraft/packet/handshake.ex

@@ -3,7 +3,6 @@ defmodule Minecraft.Packet.Handshake do
   Serialization and deserialization routines for handshake packets.
   """
   alias Minecraft.Packet.Client
-  alias Minecraft.Protocol
   import Minecraft.Packet
 
   @type packet_id :: 0
@@ -12,8 +11,7 @@ defmodule Minecraft.Packet.Handshake do
   Deserializes a handshake packet.
   """
   @spec deserialize(packet_id, binary, type :: :client | :server) ::
-          {packet :: term, new_state :: Protocol.state(), rest :: binary}
-          | {:error, :invalid_packet}
+          {packet :: term, rest :: binary} | {:error, :invalid_packet}
   def deserialize(0 = _packet_id, data, :client = _type) do
     {protocol_version, rest} = decode_varint(data)
     {server_addr, rest} = decode_string(rest)
@@ -32,7 +30,7 @@ defmodule Minecraft.Packet.Handshake do
       next_state: next_state
     }
 
-    {packet, next_state, rest}
+    {packet, rest}
   end
 
   def deserialize(_, _, _) do

lib/minecraft/packet/status.ex

@@ -4,7 +4,6 @@ defmodule Minecraft.Packet.Status do
   """
   alias Minecraft.Packet.Client
   alias Minecraft.Packet.Server
-  alias Minecraft.Protocol
   import Minecraft.Packet
 
   @type packet_id :: 0..1
@@ -13,25 +12,24 @@ defmodule Minecraft.Packet.Status do
   Deserializes a status packet.
   """
   @spec deserialize(packet_id, binary, type :: :client | :server) ::
-          {packet :: term, new_state :: Protocol.state(), rest :: binary}
-          | {:error, :invalid_packet}
+          {packet :: term, rest :: binary} | {:error, :invalid_packet}
   def deserialize(0 = _packet_id, data, :client = _type) do
-    {%Client.Status.Request{}, :status, data}
+    {%Client.Status.Request{}, data}
   end
 
   def deserialize(1 = _packet_id, data, :client) do
     <<payload::64-signed, rest::binary>> = data
-    {%Client.Status.Ping{payload: payload}, :status, rest}
+    {%Client.Status.Ping{payload: payload}, rest}
   end
 
   def deserialize(0 = _packet_id, data, :server) do
     {json, rest} = decode_string(data)
-    {%Server.Status.Response{json: json}, :status, rest}
+    {%Server.Status.Response{json: json}, rest}
   end
 
   def deserialize(1 = _packet_id, data, :server) do
     <<payload::64-signed, rest::binary>> = data
-    {%Server.Status.Pong{payload: payload}, :status, rest}
+    {%Server.Status.Pong{payload: payload}, rest}
   end
 
   def deserialize(_, _, _) do

lib/minecraft/protocol.ex

@@ -5,21 +5,11 @@ defmodule Minecraft.Protocol do
   """
   use GenServer
   require Logger
+  alias Minecraft.Connection
   alias Minecraft.Protocol.Handler
-  alias Minecraft.Packet
 
   @behaviour :ranch_protocol
 
-  @typedoc """
-  The possible states a client/server can be in.
-  """
-  @type state :: :handshaking | :status | :login | :play
-
-  defmodule State do
-    @moduledoc false
-    defstruct [:current, :socket, :transport, :client_ip]
-  end
-
   @impl true
   def start_link(ref, socket, transport, protocol_opts) do
     pid = :proc_lib.spawn_link(__MODULE__, :init, [{ref, socket, transport, protocol_opts}])
@@ -33,66 +23,57 @@ defmodule Minecraft.Protocol do
   @impl true
   def init({ref, socket, transport, _protocol_opts}) do
     :ok = :ranch.accept_ack(ref)
-    {:ok, {client_ip, _port}} = :inet.peername(socket)
-    client_ip = :inet.ntoa(client_ip)
-
-    state = %State{
-      current: :handshaking,
-      socket: socket,
-      transport: transport,
-      client_ip: client_ip
-    }
-
-    :ok = transport.setopts(socket, active: :once)
-    Logger.info(fn -> "Client #{client_ip} connected." end)
-    :gen_server.enter_loop(__MODULE__, [], state)
+    conn = Connection.init(socket, transport)
+    :gen_server.enter_loop(__MODULE__, [], conn)
   end
 
   @impl true
-  def handle_info({:tcp, socket, packet}, state) do
-    case Packet.deserialize(packet, state.current) do
-      {packet, current, rest} when is_binary(rest) ->
-        Logger.debug(fn -> "REQUEST: #{inspect(packet)}" end)
-
-        if byte_size(rest) > 0 do
-          send(self(), {:tcp, socket, rest})
-        end
-
-        handle_packet(packet, socket, current, state)
-
-      {:error, :invalid_packet} ->
-        Logger.error(fn -> "Received an invalid packet from client, closing connection." end)
-        {:stop, :normal, state}
-    end
+  def handle_info({:tcp, socket, data}, conn) do
+    conn
+    |> Connection.put_socket(socket)
+    |> Connection.put_data(data)
+    |> handle_conn()
   end
 
-  def handle_info({:tcp_closed, socket}, state) do
-    Logger.info(fn -> "Client #{state.client_ip} disconnected." end)
-    :ok = state.transport.close(socket)
-    {:stop, :normal, state}
+  def handle_info({:tcp_closed, socket}, conn) do
+    Logger.info(fn -> "Client #{conn.client_ip} disconnected." end)
+    :ok = conn.transport.close(socket)
+    {:stop, :normal, conn}
   end
 
   #
   # Helpers
   #
+  defp handle_conn(%Connection{data: ""} = conn) do
+    conn = Connection.continue(conn)
+    {:noreply, conn}
+  end
+
+  defp handle_conn(%Connection{} = conn) do
+    case Connection.read_packet(conn) do
+      {:ok, packet, conn} ->
+        handle_packet(packet, conn)
+
+      {:error, conn} ->
+        conn = Connection.close(conn)
+        {:stop, :normal, conn}
+    end
+  end
 
-  defp handle_packet(packet, socket, current, state) do
-    case Handler.handle(packet) do
-      {:ok, :noreply} ->
-        :ok = state.transport.setopts(socket, active: :once)
-        {:noreply, %State{state | current: current}}
+  defp handle_packet(packet, conn) do
+    case Handler.handle(packet, conn) do
+      {:ok, :noreply, conn} ->
+        handle_conn(conn)
 
-      {:ok, response_packet} ->
-        Logger.debug(fn -> "RESPONSE: #{inspect(response_packet)}" end)
-        {:ok, response} = Packet.serialize(response_packet)
-        :ok = state.transport.setopts(socket, active: :once)
-        :ok = state.transport.send(socket, response)
-        {:noreply, %State{state | current: current}}
+      {:ok, response, conn} ->
+        conn
+        |> Connection.send_response(response)
+        |> handle_conn()
 
-      err ->
+      {:error, _, conn} = err ->
         Logger.error(fn -> "#{__MODULE__} error: #{inspect(err)}" end)
-        :ok = state.transport.close(socket)
-        {:stop, :normal, state}
+        conn = Connection.close(conn)
+        {:stop, :normal, conn}
     end
   end
 end