Implement protocol encryption #2

Merged
kodi merged 6 commits from encryption into main 2024-10-05 12:07:04 +02:00
2 changed files with 42 additions and 22 deletions
Showing only changes of commit 7f3bb357db - Show all commits

View File

@ -107,7 +107,11 @@ defmodule Amethyst.ConnectionHandler do
send_packet(socket, connstate, packet, version, state) send_packet(socket, connstate, packet, version, state)
loop(socket, connstate, version, state) loop(socket, connstate, version, state)
after 0 -> after 0 ->
# Received stuff from the connection receiver is lower priority
receive do receive do
{:get_encryption, from} ->
send(from, Map.get(state, :decryption_state))
loop(socket, connstate, version, state)
{:packet, id, data} -> {:packet, id, data} ->
state = handle_packet(id, data, connstate, version, state) state = handle_packet(id, data, connstate, version, state)
loop(socket, 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 defp handle_packet(id, data, connstate, version, state) do
try do try do
packet = case Map.get(state, :decryption_state) do packet = connstate.deserialize(id, version, data)
nil ->
connstate.deserialize(id, version, data)
dstate ->
data = :crypto.crypto_update(dstate, data)
connstate.deserialize(id, version, data)
end
case connstate.handle(packet, version, state) do case connstate.handle(packet, version, state) do
:ok -> state :ok -> state
{:error, reason} -> {:error, reason} ->

View File

@ -35,7 +35,7 @@ defmodule Amethyst.ConnectionReceiver do
{:ok, spawn(fn -> {:ok, spawn(fn ->
Process.set_label("ConnectionReceiver for #{inspect(socket)}") Process.set_label("ConnectionReceiver for #{inspect(socket)}")
{:ok, pid} = Amethyst.ConnectionHandler.start_link(socket, Amethyst.ConnectionState.Handshake, 0) {:ok, pid} = Amethyst.ConnectionHandler.start_link(socket, Amethyst.ConnectionState.Handshake, 0)
receive(socket, pid) receive(socket, pid, nil)
end)} end)}
end end
@ -44,29 +44,46 @@ defmodule Amethyst.ConnectionReceiver do
{:ok, spawn_link(fn -> {:ok, spawn_link(fn ->
Process.set_label("ConnectionReceiver for #{inspect(socket)}") Process.set_label("ConnectionReceiver for #{inspect(socket)}")
{:ok, pid} = Amethyst.ConnectionHandler.start_link(socket, Amethyst.ConnectionState.Handshake, 0) {:ok, pid} = Amethyst.ConnectionHandler.start_link(socket, Amethyst.ConnectionState.Handshake, 0)
receive(socket, pid) receive(socket, pid, nil)
end)} end)}
end end
@spec receive(:gen_tcp.socket(), pid()) :: no_return() @spec receive(:gen_tcp.socket(), pid(), nil | :crypto.crypto_state()) :: no_return()
def receive(socket, sender) do def receive(socket, sender, cstate) do
case get_packet(socket) do case get_packet(socket, cstate) do
:closed -> send(sender, :closed) :closed -> send(sender, :closed)
Process.exit(self(), :normal) Process.exit(self(), :normal)
{:error, error} -> Logger.error("Error reading packet: #{error}") {:error, error} -> Logger.error("Error reading packet: #{error}")
{id, data} -> send(sender, {:packet, id, data}) {id, data} -> send(sender, {:packet, id, data})
end 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 end
def get_packet(client) do def get_packet(client, cstate) do
case get_varint(client, "") do case get_varint(client, "", cstate) do
:closed -> :closed :closed -> :closed
{:error, error} -> {:error, error} {:error, error} -> {:error, error}
{[length], ""} -> {[length], ""} ->
recv = :gen_tcp.recv(client, length) recv = :gen_tcp.recv(client, length)
case recv do case recv do
{:ok, full_packet} -> ({[id], data} = Read.start(full_packet) |> Read.varint() |> Read.stop() {: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}) {id, data})
{:error, :closed} -> :closed {:error, :closed} -> :closed
{:error, error} -> {:error, error} {:error, error} -> {:error, error}
@ -74,11 +91,16 @@ defmodule Amethyst.ConnectionReceiver do
end end
end end
defp get_varint(client, acc) do defp get_varint(client, acc, cstate) do
case :gen_tcp.recv(client, 1) do case :gen_tcp.recv(client, 1) do
{:ok, byte} -> case byte do {: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() <<0::1, _::7>> -> Read.start(acc <> byte) |> Read.varint() |> Read.stop()
<<1::1, _::7>> -> get_varint(client, acc <> byte) <<1::1, _::7>> -> get_varint(client, acc <> byte, cstate)
end end
{:error, :closed} -> :closed {:error, :closed} -> :closed
{:error, error} -> {:error, error} {:error, error} -> {:error, error}