Add schemas and functions to query signing keys from servers

lib/matrix_server/application.ex

@@ -13,7 +13,7 @@ defmodule MatrixServer.Application do
       MatrixServerWeb.Endpoint,
       {Registry, keys: :unique, name: MatrixServer.RoomServer.Registry},
       {DynamicSupervisor, name: MatrixServer.RoomServer.Supervisor, strategy: :one_for_one},
-      MatrixServer.SigningServer,
+      MatrixServer.KeyServer,
       {Finch, name: MatrixServerWeb.HTTPClient}
     ]
 

lib/matrix_server/key_server.ex

@@ -1,4 +1,4 @@
-defmodule MatrixServer.SigningServer do
+defmodule MatrixServer.KeyServer do
   use GenServer
 
   # TODO: only support one signing key for now.
@@ -14,15 +14,15 @@ defmodule MatrixServer.SigningServer do
     GenServer.call(__MODULE__, {:sign_object, object})
   end
 
-  def get_signing_keys(encoded \\ false) do
-    GenServer.call(__MODULE__, {:get_signing_keys, encoded})
+  def get_own_signing_keys() do
+    GenServer.call(__MODULE__, :get_own_signing_keys)
   end
 
   ## Implementation
 
   @impl true
   def init(_opts) do
-    {public_key, private_key} = get_keys()
+    {public_key, private_key} = read_keys()
     {:ok, %{public_key: public_key, private_key: private_key}}
   end
 
@@ -34,10 +34,10 @@ defmodule MatrixServer.SigningServer do
     end
   end
 
-  def handle_call({:get_signing_keys, encoded}, _from, %{public_key: public_key} = state) do
-    result = if encoded, do: MatrixServer.encode_unpadded_base64(public_key), else: public_key
+  def handle_call(:get_own_signing_keys, _from, %{public_key: public_key} = state) do
+    encoded_key = MatrixServer.encode_unpadded_base64(public_key)
 
-    {:reply, [{@signing_key_id, result}], state}
+    {:reply, [{@signing_key_id, encoded_key}], state}
   end
 
   # https://blog.swwomm.com/2020/09/elixir-ed25519-signatures-with-enacl.html
@@ -53,7 +53,7 @@ defmodule MatrixServer.SigningServer do
   end
 
   # TODO: not sure if there is a better way to do this...
-  defp get_keys do
+  def read_keys do
     raw_priv_key =
       Application.get_env(:matrix_server, :private_key_file)
       |> File.read!()

lib/matrix_server/schema/device.ex

@@ -42,9 +42,7 @@ defmodule MatrixServer.Device do
 
   def generate_device_id(localpart) do
     # TODO: use random string instead
-    time_string = System.os_time(:millisecond) |> Integer.to_string()
-
-    "#{localpart}_#{time_string}"
+    "#{localpart}_#{System.os_time(:millisecond)}"
   end
 
   def login(%Login{} = input, account) do

lib/matrix_server/schema/event.ex

@@ -3,7 +3,7 @@ defmodule MatrixServer.Event do
 
   import Ecto.Query
 
-  alias MatrixServer.{Repo, Room, Event, Account, OrderedMap, SigningServer}
+  alias MatrixServer.{Repo, Room, Event, Account, OrderedMap, KeyServer}
 
   @primary_key {:event_id, :string, []}
   schema "events" do
@@ -280,7 +280,7 @@ defmodule MatrixServer.Event do
     event
     |> Map.put(:hashes, %{"sha256" => content_hash})
     |> redact()
-    |> SigningServer.sign_object()
+    |> KeyServer.sign_object()
   end
 
   defp calculate_content_hash(event) do

lib/matrix_server/schema/server_key_info.ex

@@ -0,0 +1,80 @@
+defmodule MatrixServer.ServerKeyInfo do
+  use Ecto.Schema
+
+  import Ecto.Query
+
+  alias MatrixServer.{Repo, ServerKeyInfo, SigningKey}
+  alias MatrixServerWeb.FederationClient
+  alias MatrixServerWeb.Federation.Request.GetSigningKeys
+  alias Ecto.Multi
+
+  @primary_key {:server_name, :string, []}
+  schema "server_key_info" do
+    field :valid_until, :integer
+
+    has_many :signing_keys, SigningKey, foreign_key: :server_name
+  end
+
+  def with_fresh_signing_keys(server_name) do
+    current_time = System.os_time(:millisecond)
+
+    case with_signing_keys(server_name) do
+      nil ->
+        # We have not encountered this server before, always request keys.
+        refresh_signing_keys(server_name)
+
+      %ServerKeyInfo{valid_until: valid_until} when valid_until <= current_time ->
+        # Keys are expired; request fresh ones from server.
+        refresh_signing_keys(server_name)
+
+      ski ->
+        {:ok, ski}
+    end
+  end
+
+  defp refresh_signing_keys(server_name) do
+    in_a_week = System.os_time(:millisecond) + 1000 * 60 * 60 * 24 * 7
+    client = FederationClient.client(server_name)
+
+    with {:ok, %GetSigningKeys{verify_keys: verify_keys, valid_until_ts: valid_until}} <-
+           FederationClient.get_signing_keys(client) do
+      signing_keys =
+        Enum.map(verify_keys, fn {key_id, %{"key" => key}} ->
+          [server_name: server_name, signing_key_id: key_id, signing_key: key]
+        end)
+
+      # Always check every week to prevent misuse.
+      ski = %ServerKeyInfo{server_name: server_name, valid_until: min(valid_until, in_a_week)}
+
+      case upsert_multi(server_name, ski, signing_keys) |> Repo.transaction() do
+        {:ok, %{ski: ski}} -> {:ok, ski}
+        {:error, _} -> :error
+      end
+    else
+      :error ->
+        :error
+    end
+  end
+
+  defp upsert_multi(server_name, ski, signing_keys) do
+    Multi.new()
+    |> Multi.insert(:ski, ski,
+      on_conflict: {:replace, [:valid_until]},
+      conflict_target: [:server_name]
+    )
+    |> Multi.insert_all(:insert_keys, SigningKey, signing_keys, on_conflict: :nothing)
+    |> Multi.run(:ski, fn _, _ ->
+      case with_signing_keys(server_name) do
+        nil -> {:error, :ski}
+        ski -> {:ok, ski}
+      end
+    end)
+  end
+
+  defp with_signing_keys(server_name) do
+    ServerKeyInfo
+    |> where([ski], ski.server_name == ^server_name)
+    |> preload([ski], [:signing_keys])
+    |> Repo.one()
+  end
+end

lib/matrix_server/schema/signing_key.ex

@@ -0,0 +1,33 @@
+defmodule MatrixServer.SigningKey do
+  use Ecto.Schema
+
+  import Ecto.Changeset
+  import Ecto.Query
+
+  alias MatrixServer.{Repo, SigningKey, ServerKeyInfo}
+
+  @primary_key false
+  schema "signing_keys" do
+    field :signing_key_id, :string, primary_key: true
+    field :signing_key, :binary
+
+    belongs_to :server_key_info, ServerKeyInfo,
+      foreign_key: :server_name,
+      references: :server_name,
+      type: :string,
+      primary_key: true
+  end
+
+  def changeset(signing_key, params \\ %{}) do
+    signing_key
+    |> cast(params, [:server_name, :signing_key_id, :signing_key])
+    |> validate_required([:server_name, :signing_key_id, :signing_key])
+    |> unique_constraint([:server_name, :signing_key_id], name: :signing_keys_pkey)
+  end
+
+  def for_server(server_name) do
+    SigningKey
+    |> where([s], s.server_name == ^server_name)
+    |> Repo.all()
+  end
+end