fix(phase-91): wire remote jwks operator diagnostics

Author: szTheory

lib/lockspire/diagnostics/remote_jwks.ex

@@ -182,6 +182,24 @@ defmodule Lockspire.Diagnostics.RemoteJwks do
     |> Map.new()
   end
 
+  @spec snapshot(t()) :: map()
+  def snapshot(%__MODULE__{} = incident) do
+    %{
+      class: incident.class,
+      consumer: incident.consumer,
+      stage: incident.stage,
+      subreason: incident.subreason,
+      fetch_status: incident.fetch_status,
+      target_safety_reason: incident.target_safety_reason,
+      cached_entry_present?: incident.cached_entry_present?,
+      forced_refresh_attempted?: incident.forced_refresh_attempted?,
+      requested_kid_present_in_cached_set?: incident.requested_kid_present_in_cached_set?,
+      remediation: incident.remediation
+    }
+    |> Enum.reject(fn {_key, value} -> is_nil(value) end)
+    |> Map.new()
+  end
+
   @spec summarize_client(Client.t()) :: summary()
   def summarize_client(%Client{} = client) do
     if remote_jwks_client?(client) do
@@ -272,12 +290,11 @@ defmodule Lockspire.Diagnostics.RemoteJwks do
     "Confirm overlap-based key rollover and keep both old and new keys published until Lockspire can refresh and verify a fresh JWT."
   end
 
-  defp remote_jwks_client?(%Client{
-         jwks_uri: jwks_uri,
-         token_endpoint_auth_method: :private_key_jwt
-       })
-       when is_binary(jwks_uri) and jwks_uri != "",
-       do: true
+  defp remote_jwks_client?(%Client{jwks_uri: jwks_uri} = client)
+       when is_binary(jwks_uri) and jwks_uri != "" do
+    client.token_endpoint_auth_method == :private_key_jwt or
+      not is_nil(client.authorization_encrypted_response_alg)
+  end
 
   defp remote_jwks_client?(_client), do: false
 

lib/lockspire/protocol/client_auth/private_key_jwt.ex

@@ -31,6 +31,7 @@ defmodule Lockspire.Protocol.ClientAuth.PrivateKeyJwt do
                ),
              :ok <- validate_claims(verified_assertion, verified_client, opts),
              :ok <- record_replay(verified_assertion, verified_client, opts) do
+          clear_remote_jwks_diagnostic(client, opts)
           :ok
         else
           {:error, reason, remote_jwks_incident} ->
@@ -368,6 +369,7 @@ defmodule Lockspire.Protocol.ClientAuth.PrivateKeyJwt do
 
     Observability.emit(:client_auth, action, %{}, metadata)
     append_audit_event(reason, client, metadata, opts)
+    persist_remote_jwks_diagnostic(client, remote_jwks_incident, opts)
   end
 
   defp append_audit_event(reason, %Client{} = client, metadata, opts) do
@@ -427,6 +429,41 @@ defmodule Lockspire.Protocol.ClientAuth.PrivateKeyJwt do
 
   defp replay_store(opts), do: Keyword.get(opts, :jti_store, Keyword.fetch!(opts, :client_store))
 
+  defp persist_remote_jwks_diagnostic(%Client{} = client, %RemoteJwks{} = incident, opts) do
+    with store when not is_nil(store) <- Keyword.get(opts, :client_store),
+         true <- function_exported?(store, :update_client, 2) do
+      metadata =
+        client.metadata
+        |> ensure_metadata()
+        |> Map.put("remote_jwks_diagnostic", RemoteJwks.snapshot(incident))
+
+      _ = store.update_client(client, %{metadata: metadata})
+      :ok
+    else
+      _other -> :ok
+    end
+  end
+
+  defp persist_remote_jwks_diagnostic(_client, _incident, _opts), do: :ok
+
+  defp clear_remote_jwks_diagnostic(%Client{jwks_uri: jwks_uri} = client, opts)
+       when is_binary(jwks_uri) do
+    with store when not is_nil(store) <- Keyword.get(opts, :client_store),
+         true <- function_exported?(store, :update_client, 2),
+         metadata when is_map(metadata) <- client.metadata,
+         true <- Map.has_key?(metadata, "remote_jwks_diagnostic") do
+      _ = store.update_client(client, %{metadata: Map.delete(metadata, "remote_jwks_diagnostic")})
+      :ok
+    else
+      _other -> :ok
+    end
+  end
+
+  defp clear_remote_jwks_diagnostic(_client, _opts), do: :ok
+
+  defp ensure_metadata(metadata) when is_map(metadata), do: metadata
+  defp ensure_metadata(_metadata), do: %{}
+
   defp server_policy_store(opts),
     do:
       Keyword.get_lazy(opts, :server_policy_store, fn ->

lib/lockspire/protocol/jarm/client_key_resolver.ex

@@ -37,13 +37,15 @@ defmodule Lockspire.Protocol.Jarm.ClientKeyResolver do
     end
   end
 
-  defp do_resolve(%Client{jwks_uri: jwks_uri} = _client, params, opts) when is_binary(jwks_uri) do
+  defp do_resolve(%Client{jwks_uri: jwks_uri} = client, params, opts) when is_binary(jwks_uri) do
     fetcher = Keyword.get(opts, :jwks_fetcher, Config.jwks_fetcher())
+    opts = Keyword.put_new(opts, :client, client)
 
     with {:ok, jwk_set} <- fetcher.get_keys(jwks_uri, jwks_fetcher_opts(opts)),
          {_modules, jwks} <- JOSE.JWK.to_map(jwk_set) do
       case select_key(jwks, params) do
         {:ok, jwk} ->
+          clear_remote_jwks_diagnostic(client_from_opts_or_nil(opts), opts)
           {:ok, jwk, :jwks_uri}
 
         {:error, :jarm_encryption_key_unavailable} ->
@@ -53,7 +55,9 @@ defmodule Lockspire.Protocol.Jarm.ClientKeyResolver do
       {:error, {:jwks_fetch_failed, _reason} = fetch_error} ->
         emit_remote_failure(
           :jarm_encryption_key_fetch_failed,
-          RemoteJwks.classify_fetch_error(:jarm, fetch_error)
+          RemoteJwks.classify_fetch_error(:jarm, fetch_error),
+          client,
+          opts
         )
 
         {:error, :jarm_encryption_key_fetch_failed}
@@ -74,6 +78,7 @@ defmodule Lockspire.Protocol.Jarm.ClientKeyResolver do
 
         case select_key(jwks, params) do
           {:ok, jwk} ->
+            clear_remote_jwks_diagnostic(client_from_opts_or_nil(opts), opts)
             {:ok, jwk, :jwks_uri}
 
           {:error, :jarm_encryption_key_unavailable} ->
@@ -86,7 +91,9 @@ defmodule Lockspire.Protocol.Jarm.ClientKeyResolver do
                   :requested_kid_present_in_cached_set?,
                   requested_kid_present?(jwks, params)
                 )
-              )
+              ),
+              client_from_opts_or_nil(opts),
+              opts
             )
 
             {:error, :jarm_encryption_key_unavailable}
@@ -95,7 +102,9 @@ defmodule Lockspire.Protocol.Jarm.ClientKeyResolver do
       {:error, {:jwks_fetch_failed, _reason} = fetch_error} ->
         emit_remote_failure(
           :jarm_encryption_key_fetch_failed,
-          RemoteJwks.classify_fetch_error(:jarm, fetch_error, remote_opts)
+          RemoteJwks.classify_fetch_error(:jarm, fetch_error, remote_opts),
+          client_from_opts_or_nil(opts),
+          opts
         )
 
         {:error, :jarm_encryption_key_fetch_failed}
@@ -180,6 +189,47 @@ defmodule Lockspire.Protocol.Jarm.ClientKeyResolver do
     Observability.emit(:jarm, :failed, %{}, metadata)
   end
 
+  defp emit_remote_failure(reason_code, remote_jwks_incident, %Client{} = client, opts) do
+    emit_remote_failure(reason_code, remote_jwks_incident)
+    persist_remote_jwks_diagnostic(client, remote_jwks_incident, opts)
+  end
+
+  defp client_from_opts_or_nil(opts), do: Keyword.get(opts, :client)
+
+  defp persist_remote_jwks_diagnostic(%Client{} = client, %RemoteJwks{} = incident, opts) do
+    with store when not is_nil(store) <- Keyword.get(opts, :client_store, Config.repo!()),
+         true <- function_exported?(store, :update_client, 2) do
+      metadata =
+        client.metadata
+        |> ensure_metadata()
+        |> Map.put("remote_jwks_diagnostic", RemoteJwks.snapshot(incident))
+
+      _ = store.update_client(client, %{metadata: metadata})
+      :ok
+    else
+      _other -> :ok
+    end
+  end
+
+  defp persist_remote_jwks_diagnostic(_client, _incident, _opts), do: :ok
+
+  defp clear_remote_jwks_diagnostic(%Client{} = client, opts) do
+    with store when not is_nil(store) <- Keyword.get(opts, :client_store, Config.repo!()),
+         true <- function_exported?(store, :update_client, 2),
+         metadata when is_map(metadata) <- client.metadata,
+         true <- Map.has_key?(metadata, "remote_jwks_diagnostic") do
+      _ = store.update_client(client, %{metadata: Map.delete(metadata, "remote_jwks_diagnostic")})
+      :ok
+    else
+      _other -> :ok
+    end
+  end
+
+  defp clear_remote_jwks_diagnostic(_client, _opts), do: :ok
+
+  defp ensure_metadata(metadata) when is_map(metadata), do: metadata
+  defp ensure_metadata(_metadata), do: %{}
+
   defp jwks_fetcher_opts(opts) do
     Config.jwks_fetcher_opts()
     |> Keyword.merge(Keyword.get(opts, :jwks_fetcher_opts, []))

lib/mix/tasks/lockspire.doctor.ex

@@ -0,0 +1,23 @@
+defmodule Mix.Tasks.Lockspire.Doctor do
+  @moduledoc """
+  Dispatcher for Lockspire runtime diagnostic subcommands.
+  """
+
+  use Mix.Task
+
+  @shortdoc "Runs Lockspire runtime diagnostic commands"
+
+  @impl Mix.Task
+  def run(["remote-jwks" | rest]) do
+    Mix.Task.run("lockspire.doctor.remote_jwks", rest)
+  end
+
+  def run(_args) do
+    Mix.raise("""
+    Unknown doctor command.
+
+    Supported commands:
+      mix lockspire.doctor remote-jwks --client CLIENT_ID
+    """)
+  end
+end

test/lockspire/admin/clients_test.exs

@@ -253,6 +253,35 @@ defmodule Lockspire.Admin.ClientsTest do
              "mix lockspire.doctor remote-jwks --client admin-remote-jwks-summary"
   end
 
+  test "remote_jwks_summary/1 also applies to JARM-only jwks_uri clients" do
+    {:ok, client} =
+      Repository.register_client(%Client{
+        client_id: "admin-remote-jwks-jarm",
+        client_secret_hash: "sha256:remote:jarm",
+        client_type: :confidential,
+        name: "Admin Remote JWKS JARM",
+        redirect_uris: ["https://remote.example.com/callback"],
+        allowed_scopes: ["openid"],
+        allowed_grant_types: ["authorization_code"],
+        allowed_response_types: ["code"],
+        token_endpoint_auth_method: :client_secret_basic,
+        authorization_encrypted_response_alg: :RSA_OAEP_256,
+        authorization_encrypted_response_enc: :A256GCM,
+        pkce_required: true,
+        subject_type: :public,
+        created_at: DateTime.utc_now(),
+        jwks_uri: "https://remote.example.com/.well-known/jwks.json",
+        metadata: %{}
+      })
+
+    summary = Clients.remote_jwks_summary(client)
+
+    assert summary.applicable? == true
+    assert summary.status == :supported
+    assert summary.command_hint =~
+             "mix lockspire.doctor remote-jwks --client admin-remote-jwks-jarm"
+  end
+
   test "remote_jwks_summary/1 reuses the shared incident taxonomy from client metadata" do
     {:ok, client} =
       Repository.register_client(%Client{

test/lockspire/protocol/client_auth_test.exs

@@ -20,6 +20,7 @@ defmodule Lockspire.Protocol.ClientAuthTest do
     Process.delete(:recorded_audit_events)
     Process.delete(:force_replay_store_error)
     Process.delete(:client_secret_jwt_secret)
+    Process.delete(:updated_remote_client)
     :ok
   end
 
@@ -91,6 +92,7 @@ defmodule Lockspire.Protocol.ClientAuthTest do
                )
 
       assert Process.get(:fetched_jwks_uri) == "https://keys.example.test/client.jwks.json"
+      assert Process.get(:updated_remote_client) == nil
     end
 
     test "refreshes remote jwks once on key mismatch and retries verification" do
@@ -123,6 +125,7 @@ defmodule Lockspire.Protocol.ClientAuthTest do
 
       assert Process.get(:fetched_jwks_uri) == "https://keys.example.test/client.jwks.json"
       assert Process.get(:refreshed_jwks_uri) == "https://keys.example.test/client.jwks.json"
+      assert Process.get(:updated_remote_client) == nil
     end
 
     test "rejects algorithms outside the effective issuer allowlist" do
@@ -345,6 +348,10 @@ defmodule Lockspire.Protocol.ClientAuthTest do
                         remote_jwks_forced_refresh_attempted?: true,
                         remote_jwks_requested_kid_present_in_cached_set?: false
                       }}
+
+      assert %{"remote_jwks_diagnostic" => diagnostic} = Process.get(:updated_remote_client)
+      assert diagnostic[:class] == :remote_jwks_key_unavailable
+      assert diagnostic[:consumer] == :private_key_jwt
     end
 
     test "emits shared remote jwks incident metadata for guarded fetch failures" do
@@ -388,6 +395,10 @@ defmodule Lockspire.Protocol.ClientAuthTest do
                         remote_jwks_fetch_status: 503,
                         remote_jwks_forced_refresh_attempted?: false
                       }}
+
+      assert %{"remote_jwks_diagnostic" => diagnostic} = Process.get(:updated_remote_client)
+      assert diagnostic[:class] == :remote_jwks_fetch_failed
+      assert diagnostic[:fetch_status] == 503
     end
 
     test "emits shared remote jwks signature-invalid metadata when the refreshed kid exists" do
@@ -439,6 +450,10 @@ defmodule Lockspire.Protocol.ClientAuthTest do
                         remote_jwks_forced_refresh_attempted?: true,
                         remote_jwks_requested_kid_present_in_cached_set?: true
                       }}
+
+      assert %{"remote_jwks_diagnostic" => diagnostic} = Process.get(:updated_remote_client)
+      assert diagnostic[:class] == :remote_jwks_signature_invalid
+      assert diagnostic[:consumer] == :private_key_jwt
     end
   end
 
@@ -750,6 +765,11 @@ defmodule Lockspire.Protocol.ClientAuthTest do
     end
 
     def fetch_client_by_id(_), do: {:ok, nil}
+
+    def update_client(_client, attrs) do
+      Process.put(:updated_remote_client, attrs.metadata)
+      {:ok, struct(Client, attrs)}
+    end
   end
 
   defmodule RemoteJwksFetcher do