Fix Desktop.Auth login key race under concurrent first access

lib/desktop/auth.ex

@@ -7,24 +7,60 @@ defmodule Desktop.Auth do
 
   import Plug.Conn
   alias Desktop.OS
+  require Logger
   @behaviour Plug
 
+  @table __MODULE__
+  @key {__MODULE__, :key}
+
   defp key() do
     # key should stay the same during application run,
     # but be different on each instance
-    case :persistent_term.get({__MODULE__, :key}, nil) do
-      nil ->
-        key = :crypto.strong_rand_bytes(32)
-        :persistent_term.put({__MODULE__, :key}, key)
-        key
-
-      key ->
-        key
+    case :persistent_term.get(@key, nil) do
+      nil -> init_key()
+      key -> key
+    end
+  end
+
+  # `persistent_term` get/put is not atomic: concurrent first access from
+  # Window.prepare_url/1 and this plug can mint different keys and leave the
+  # webview on a blank "Unauthorized" page. ETS insert_new elects one winner;
+  # persistent_term remains the fast read path.
+  defp init_key() do
+    table = table!()
+    key = :crypto.strong_rand_bytes(32)
+
+    if :ets.insert_new(table, {:key, key}) do
+      store_key(key)
+    else
+      [{:key, key}] = :ets.lookup(table, :key)
+      store_key(key)
+    end
+  end
+
+  defp store_key(key) do
+    :persistent_term.put(@key, key)
+    key
+  end
+
+  defp table!() do
+    case :ets.whereis(@table) do
+      :undefined ->
+        try do
+          :ets.new(@table, [:named_table, :public, :set, read_concurrency: true])
+        rescue
+          ArgumentError -> table!()
+        end
+
+      tid ->
+        tid
     end
   end
 
   def set_key(key) do
-    :persistent_term.put({__MODULE__, :key}, Base.decode32!(key))
+    decoded = Base.decode32!(key)
+    :ets.insert(table!(), {:key, decoded})
+    :persistent_term.put(@key, decoded)
   end
 
   def login_key() do
@@ -44,10 +80,17 @@ defmodule Desktop.Auth do
 
   defp require_auth(conn) do
     conn = fetch_query_params(conn)
+    k = conn.query_params["k"] || ""
 
-    if OS.mobile?() or Plug.Crypto.secure_compare(login_key(), conn.query_params["k"] || "") do
+    if OS.mobile?() or Plug.Crypto.secure_compare(login_key(), k) do
       put_session(conn, :user, true)
     else
+      has_k? = k != ""
+
+      Logger.warning(
+        "Desktop.Auth Unauthorized path=#{conn.request_path} has_k=#{has_k?} peer=#{inspect(conn.remote_ip)}"
+      )
+
       conn
       |> resp(401, "Unauthorized")
       |> halt()

test/auth_test.exs

@@ -0,0 +1,108 @@
+defmodule Desktop.AuthTest do
+  use ExUnit.Case, async: false
+
+  setup do
+    reset_auth_key!()
+    :ok
+  end
+
+  test "login_key/0 is idempotent" do
+    key1 = Desktop.Auth.login_key()
+    key2 = Desktop.Auth.login_key()
+
+    assert key1 == key2
+    assert byte_size(key1) > 0
+  end
+
+  test "login_key/0 returns one key under concurrent first access" do
+    parent = self()
+
+    for _ <- 1..40 do
+      spawn(fn ->
+        send(parent, {:key, Desktop.Auth.login_key()})
+      end)
+    end
+
+    keys =
+      for _ <- 1..40 do
+        receive do
+          {:key, key} -> key
+        after
+          5_000 -> flunk("timeout waiting for auth key")
+        end
+      end
+
+    assert length(Enum.uniq(keys)) == 1
+    assert hd(keys) == Desktop.Auth.login_key()
+  end
+
+  test "set_key/1 is visible to login_key/0 and Auth plug" do
+    raw = :crypto.strong_rand_bytes(32)
+    encoded = Base.encode32(raw)
+    login = Base.encode32(raw, padding: false)
+
+    Desktop.Auth.set_key(encoded)
+    assert Desktop.Auth.login_key() == login
+
+    opts = Desktop.Auth.init([])
+
+    conn =
+      Plug.Test.conn(:get, "/?k=#{login}")
+      |> Plug.Test.init_test_session(%{})
+      |> Desktop.Auth.call(opts)
+
+    refute conn.halted
+  end
+
+  test "Auth plug rejects missing key and logs a warning" do
+    import ExUnit.CaptureLog
+
+    _ = Desktop.Auth.login_key()
+    opts = Desktop.Auth.init([])
+
+    log =
+      capture_log(fn ->
+        conn =
+          Plug.Test.conn(:get, "/")
+          |> Plug.Test.init_test_session(%{})
+          |> Desktop.Auth.call(opts)
+
+        assert conn.halted
+        assert conn.status == 401
+        assert conn.resp_body == "Unauthorized"
+      end)
+
+    assert log =~ "Desktop.Auth Unauthorized"
+    assert log =~ "has_k=false"
+  end
+
+  test "Auth plug rejects wrong key and logs has_k=true" do
+    import ExUnit.CaptureLog
+
+    _ = Desktop.Auth.login_key()
+    opts = Desktop.Auth.init([])
+
+    log =
+      capture_log(fn ->
+        conn =
+          Plug.Test.conn(:get, "/dashboard?k=WRONG")
+          |> Plug.Test.init_test_session(%{})
+          |> Desktop.Auth.call(opts)
+
+        assert conn.halted
+        assert conn.status == 401
+      end)
+
+    assert log =~ "Desktop.Auth Unauthorized path=/dashboard"
+    assert log =~ "has_k=true"
+  end
+
+  defp reset_auth_key!() do
+    :persistent_term.erase({Desktop.Auth, :key})
+
+    case :ets.whereis(Desktop.Auth) do
+      :undefined -> :ok
+      _tid -> :ets.delete_all_objects(Desktop.Auth)
+    end
+  end
+end