Add schemas and functions to query signing keys from servers

lib/matrix_server.ex

@@ -90,4 +90,23 @@ defmodule MatrixServer do
     |> Map.from_struct()
     |> Map.drop(waste_fields)
   end
+
+  def add_signature(object, key_id, sig) when not is_map_key(object, :signatures) do
+    Map.put(object, :signatures, %{MatrixServer.server_name() => %{key_id => sig}})
+  end
+
+  def add_signature(%{signatures: sigs} = object, key_id, sig) do
+    new_sigs =
+      Map.update(sigs, MatrixServer.server_name(), %{key_id => sig}, &Map.put(&1, key_id, sig))
+
+    %{object | signatures: new_sigs}
+  end
+
+  def validate_change_simple(changeset, field, func) do
+    augmented_func = fn _, val ->
+      if func.(val), do: [], else: [{field, "invalid"}]
+    end
+
+    Ecto.Changeset.validate_change(changeset, field, augmented_func)
+  end
 end

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

lib/matrix_server_web/authenticate_server.ex

@@ -1,19 +1,17 @@
 defmodule MatrixServerWeb.AuthenticateServer do
   import MatrixServerWeb.Plug.Error
 
-  alias MatrixServer.SigningServer
+  alias MatrixServer.SigningKey
 
   @auth_header_regex ~r/^X-Matrix origin=(?<origin>.*),key="(?<key>.*)",sig="(?<sig>.*)"$/
 
-  def authenticate(
-        %Plug.Conn{
+  def authenticate(%Plug.Conn{
         body_params: body_params,
         req_headers: headers,
         request_path: path,
         method: method,
         query_string: query_string
-        }
-      ) do
+      }) do
     object_to_sign = %{
       uri: path <> "?" <> URI.decode_www_form(query_string),
       method: method,
@@ -32,13 +30,16 @@ defmodule MatrixServerWeb.AuthenticateServer do
     headers
     |> parse_authorization_headers()
     |> Enum.find(:error, fn {origin, _, sig} ->
-      # TODO: fetch actual signing keys for origin from cache/key store.
-      {_, signing_key} = SigningServer.get_signing_keys() |> hd()
       object = object_fun.(origin)
 
       with {:ok, raw_sig} <- MatrixServer.decode_base64(sig),
            {:ok, encoded_object} <- MatrixServer.encode_canonical_json(object) do
+        # TODO: Only query once per origin.
+        # TODO: Handle expired keys.
+        SigningKey.for_server(origin)
+        |> Enum.find_value(false, fn %SigningKey{signing_key: signing_key} ->
           :enacl.sign_verify_detached(raw_sig, encoded_object, signing_key)
+        end)
       else
         _ -> false
       end
@@ -70,6 +71,7 @@ defmodule MatrixServerWeb.AuthenticateServer do
             {origin, _key, _sig} ->
               conn = Plug.Conn.assign(conn, :origin, origin)
               apply(__MODULE__, action, [conn, conn.params])
+
             :error ->
               put_error(conn, :unauthorized, "Signature verification failed.")
           end

lib/matrix_server_web/federation/key_controller.ex

@@ -3,13 +3,13 @@ defmodule MatrixServerWeb.Federation.KeyController do
 
   import MatrixServerWeb.Plug.Error
 
-  alias MatrixServer.SigningServer
+  alias MatrixServer.KeyServer
 
   @key_valid_time_ms 1000 * 60 * 24 * 30
 
   def get_signing_keys(conn, _params) do
     keys =
-      SigningServer.get_signing_keys(true)
+      KeyServer.get_own_signing_keys()
       |> Enum.into(%{}, fn {key_id, key} ->
         {key_id, %{"key" => key}}
       end)
@@ -21,13 +21,15 @@ defmodule MatrixServerWeb.Federation.KeyController do
       valid_until_ts: System.os_time(:millisecond) + @key_valid_time_ms
     }
 
-    case SigningServer.sign_object(data) do
-      {:ok, signed_data} ->
+    case KeyServer.sign_object(data) do
+      {:ok, sig, key_id} ->
+        signed_data = MatrixServer.add_signature(data, key_id, sig)
+
         conn
         |> put_status(200)
         |> json(signed_data)
 
-      {:error, _msg} ->
+      :error ->
         put_error(conn, :unknown, "Error signing object.")
     end
   end

lib/matrix_server_web/federation/request/get_signing_keys.ex

@@ -0,0 +1,36 @@
+defmodule MatrixServerWeb.Federation.Request.GetSigningKeys do
+  use Ecto.Schema
+
+  import Ecto.Changeset
+
+  @primary_key false
+  embedded_schema do
+    field :server_name, :string
+    field :verify_keys, {:map, {:map, :string}}
+    field :old_verify_keys, {:map, :map}
+    field :signatures, {:map, {:map, :string}}
+    field :valid_until_ts, :integer
+  end
+
+  def changeset(params) do
+    # TODO: There must be a better way to validate embedded maps?
+    %__MODULE__{}
+    |> cast(params, [:server_name, :verify_keys, :old_verify_keys, :signatures, :valid_until_ts])
+    |> validate_required([:server_name, :verify_keys, :valid_until_ts])
+    |> MatrixServer.validate_change_simple(:verify_keys, fn map ->
+      Enum.all?(map, fn {_, map} ->
+        is_map_key(map, "key")
+      end)
+    end)
+    |> MatrixServer.validate_change_simple(:old_verify_keys, fn map ->
+      Enum.all?(map, fn
+        {_, %{"key" => key, "expired_ts" => expired_ts}}
+        when is_binary(key) and is_integer(expired_ts) ->
+          true
+
+        _ ->
+          false
+      end)
+    end)
+  end
+end

lib/matrix_server_web/federation_client.ex

@@ -2,6 +2,7 @@ defmodule MatrixServerWeb.FederationClient do
   use Tesla
 
   alias MatrixServerWeb.Endpoint
+  alias MatrixServerWeb.Federation.Request.GetSigningKeys
   alias MatrixServerWeb.Router.Helpers, as: RouteHelpers
 
   # TODO: Maybe create database-backed homeserver struct to pass to client function.
@@ -17,7 +18,15 @@ defmodule MatrixServerWeb.FederationClient do
   end
 
   def get_signing_keys(client) do
-    Tesla.get(client, RouteHelpers.key_path(Endpoint, :get_signing_keys))
+    # TODO: Extract into seperate function.
+    # TODO: Check signatures for each verify key.
+    with {:ok, %Tesla.Env{body: body}} <-
+           Tesla.get(client, RouteHelpers.key_path(Endpoint, :get_signing_keys)),
+         %Ecto.Changeset{valid?: true} = cs <- GetSigningKeys.changeset(body) do
+      {:ok, Ecto.Changeset.apply_changes(cs)}
+    else
+      _ -> :error
+    end
   end
 
   # TODO: Create tesla middleware to add signature and headers.
@@ -33,7 +42,7 @@ defmodule MatrixServerWeb.FederationClient do
       destination: server_name
     }
 
-    {:ok, signature, key_id} = MatrixServer.SigningServer.sign_object(object_to_sign)
+    {:ok, signature, key_id} = MatrixServer.KeyServer.sign_object(object_to_sign)
     signatures = %{origin => %{key_id => signature}}
     auth_headers = create_signature_authorization_headers(signatures, origin)
 

priv/repo/migrations/20210812105932_create_signing_key_and_server_key_info_table.exs

@@ -0,0 +1,19 @@
+defmodule MatrixServer.Repo.Migrations.CreateSigningKeyAndServerKeyInfoTable do
+  use Ecto.Migration
+
+  def change do
+    create table(:server_key_info, primary_key: false) do
+      add :server_name, :string, primary_key: true, null: false
+      add :valid_until, :bigint, default: 0, null: false
+    end
+
+    create table(:signing_keys, primary_key: false) do
+      add :server_name,
+          references(:server_key_info, column: :server_name, type: :string, on_delete: :delete_all),
+          null: false
+
+      add :signing_key_id, :string, primary_key: true, null: false
+      add :signing_key, :binary, null: false
+    end
+  end
+end

psql.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+sudo -u postgres psql -d matrix_server_dev