fix jwks refresh for client stores

Author: woylie

lib/oidcc/plug/authorization_callback.ex

@@ -141,6 +141,16 @@ defmodule Oidcc.Plug.AuthorizationCallback do
           | :state_not_verified
           | {:missing_request_param, param :: String.t()}
 
+  # what Oidcc.Plug.Authorize writes into the session. peer_ip and useragent are
+  # read back from the client, so the checks tolerate them being absent.
+  @typep session() :: %{
+           nonce: String.t(),
+           peer_ip: :inet.ip_address() | nil,
+           useragent: String.t() | nil,
+           pkce_verifier: String.t(),
+           state_verifier: integer()
+         }
+
   @impl Plug
   def init(opts),
     do:
@@ -185,21 +195,15 @@ defmodule Oidcc.Plug.AuthorizationCallback do
            {:ok, code} <- fetch_request_param(params, "code"),
            scope = Map.get(params, "scope", "openid"),
            token_opts =
-             prepare_retrieve_opts(
-               opts,
-               scope,
-               session.nonce,
-               redirect_uri,
-               session.pkce_verifier
-             ),
+             prepare_retrieve_opts(client_context, session, scope, redirect_uri, opts),
            {:ok, token} <-
              retrieve_token(
                code,
                client_context,
                retrieve_userinfo?,
                Map.merge(profile_opts, token_opts)
              ),
-           userinfo_opts = prepare_userinfo_opts(opts),
+           userinfo_opts = prepare_userinfo_opts(client_context, opts),
            {:ok, userinfo} <- retrieve_userinfo(token, client_context, userinfo_opts, retrieve_userinfo?) do
         {:ok, {token, userinfo}}
       end
@@ -217,44 +221,37 @@ defmodule Oidcc.Plug.AuthorizationCallback do
   end
 
   @spec prepare_retrieve_opts(
-          opts :: opts(),
+          client_context :: ClientContext.t(),
+          session :: session(),
           scope :: String.t(),
-          nonce :: String.t() | :any,
           redirect_uri :: String.t(),
-          pkce_verifier :: String.t() | :none
+          opts :: opts()
         ) :: :oidcc_token.retrieve_opts()
-  defp prepare_retrieve_opts(opts, scope, nonce, redirect_uri, pkce_verifier) do
+  defp prepare_retrieve_opts(client_context, session, scope, redirect_uri, opts) do
     scopes = :oidcc_scope.parse(scope)
 
-    refresh_jwks = Utils.get_refresh_jwks_fun(opts)
-
     opts
     |> Keyword.take([:request_opts, :preferred_auth_methods])
     |> Map.new()
     |> Map.merge(%{
-      nonce: nonce,
+      nonce: session.nonce,
       scope: scopes,
       redirect_uri: redirect_uri,
-      pkce_verifier: pkce_verifier,
-      refresh_jwks: refresh_jwks
+      pkce_verifier: session.pkce_verifier
     })
-    |> case do
-      %{pkce_verifier: :none} = opts -> Map.delete(opts, :pkce_verifier)
-      opts -> opts
-    end
+    |> Utils.put_refresh_jwks(client_context, opts)
   end
 
-  @spec prepare_userinfo_opts(opts :: opts()) :: :oidcc_userinfo.retrieve_opts()
-  defp prepare_userinfo_opts(opts) do
-    refresh_jwks = Utils.get_refresh_jwks_fun(opts)
-
-    %{refresh_jwks: refresh_jwks}
+  @spec prepare_userinfo_opts(client_context :: ClientContext.t(), opts :: opts()) ::
+          :oidcc_userinfo.retrieve_opts()
+  defp prepare_userinfo_opts(client_context, opts) do
+    Utils.put_refresh_jwks(%{}, client_context, opts)
   end
 
   # The session is written by Oidcc.Plug.Authorize. Without it, there is nothing
   # to validate. Reject the request to prevent CSRF in that case.
   @spec fetch_authorize_session(conn :: Plug.Conn.t()) ::
-          {:ok, map()} | {:error, error()}
+          {:ok, session()} | {:error, error()}
   defp fetch_authorize_session(conn) do
     case get_session(conn, Authorize.get_session_name()) do
       %{

lib/oidcc/plug/utils.ex

@@ -29,17 +29,52 @@ defmodule Oidcc.Plug.Utils do
   end
 
   @doc """
-  Returns a function to refresh the JWKS for a provider.
+  Adds the JWKS refresh function to `map`, unless there is none.
+
+  The key is omitted rather than set to `nil` because `oidcc` only skips the
+  refresh when the key is absent.
   """
-  @spec get_refresh_jwks_fun(Keyword.t()) ::
+  @spec put_refresh_jwks(map(), ClientContext.t(), Keyword.t()) :: map()
+  def put_refresh_jwks(map, client_context, opts) do
+    case get_refresh_jwks_fun(client_context, opts) do
+      nil -> map
+      refresh_jwks -> Map.put(map, :refresh_jwks, refresh_jwks)
+    end
+  end
+
+  # oidcc calls the refresh function with (jwks, kid), while
+  # c:Oidcc.Plug.ClientStore.refresh_jwks/1 takes the client context => the store
+  # callback is wrapped rather than captured directly.
+  @spec get_refresh_jwks_fun(ClientContext.t(), Keyword.t()) ::
           :oidcc_jwt_util.refresh_jwks_for_unknown_kid_fun() | nil
-  def get_refresh_jwks_fun(opts) do
-    if client_store = Keyword.get(opts, :client_store) do
-      if function_exported?(client_store, :refresh_jwks, 1),
-        do: &client_store.refresh_jwks/1
-    else
-      provider = Keyword.fetch!(opts, :provider)
-      :oidcc_jwt_util.refresh_jwks_fun(provider)
+  defp get_refresh_jwks_fun(client_context, opts) do
+    case Keyword.get(opts, :client_store) do
+      nil ->
+        provider = Keyword.fetch!(opts, :provider)
+        :oidcc_jwt_util.refresh_jwks_fun(provider)
+
+      client_store ->
+        client_store_refresh_jwks_fun(client_store, client_context)
+    end
+  end
+
+  @spec client_store_refresh_jwks_fun(module(), ClientContext.t()) ::
+          :oidcc_jwt_util.refresh_jwks_for_unknown_kid_fun() | nil
+  defp client_store_refresh_jwks_fun(client_store, client_context) do
+    Code.ensure_loaded!(client_store)
+
+    if function_exported?(client_store, :refresh_jwks, 1),
+      do: fn _jwks, _kid -> refresh_jwks_from_store(client_store, client_context) end
+  end
+
+  # The callback returns a JOSE.JWK struct, while oidcc puts the result straight
+  # into the client context record => convert it.
+  @spec refresh_jwks_from_store(module(), ClientContext.t()) ::
+          {:ok, :jose_jwk.key()} | {:error, term()}
+  defp refresh_jwks_from_store(client_store, client_context) do
+    case client_store.refresh_jwks(client_context) do
+      {:ok, %JOSE.JWK{} = jwks} -> {:ok, JOSE.JWK.to_record(jwks)}
+      other -> other
     end
   end
 

lib/oidcc/plug/validate_jwt_token.ex

@@ -49,9 +49,10 @@ defmodule Oidcc.Plug.ValidateJwtToken do
   """
   @typedoc since: "0.1.0"
   @type opts :: [
-          provider: GenServer.name(),
-          client_id: String.t() | (-> String.t()),
-          client_secret: String.t() | (-> String.t()),
+          provider: GenServer.name() | nil,
+          client_store: module() | nil,
+          client_id: String.t() | (-> String.t()) | nil,
+          client_secret: String.t() | (-> String.t()) | nil,
           send_inactive_token_response: (conn :: Plug.Conn.t() -> Plug.Conn.t()),
           validate_opts: Oidcc.Token.retrieve_opts()
         ]
@@ -76,6 +77,7 @@ defmodule Oidcc.Plug.ValidateJwtToken do
       opts
       |> Keyword.validate!([
         :provider,
+        :client_store,
         :client_id,
         :client_secret,
         send_inactive_token_response: &__MODULE__.send_inactive_token_response/1,
@@ -89,16 +91,9 @@ defmodule Oidcc.Plug.ValidateJwtToken do
   def call(%Plug.Conn{private: %{ExtractAuthorization => access_token}} = conn, opts) do
     send_inactive_token_response = Keyword.fetch!(opts, :send_inactive_token_response)
 
-    refresh_jwks = Utils.get_refresh_jwks_fun(opts)
-
-    validate_opts =
-      Map.merge(
-        opts[:validate_opts],
-        %{nonce: :any, refresh_jwks: refresh_jwks}
-      )
-
     with {:ok, client_context} <-
            Utils.get_client_context(conn, opts),
+         validate_opts = prepare_validate_opts(client_context, opts),
          {:ok, claims} <-
            Oidcc.Token.validate_id_token(access_token, client_context, validate_opts) do
       put_private(conn, __MODULE__, claims)
@@ -119,6 +114,15 @@ defmodule Oidcc.Plug.ValidateJwtToken do
     """
   end
 
+  @spec prepare_validate_opts(client_context :: Oidcc.ClientContext.t(), opts :: opts()) ::
+          Oidcc.Token.retrieve_opts()
+  defp prepare_validate_opts(client_context, opts) do
+    opts
+    |> Keyword.fetch!(:validate_opts)
+    |> Map.put(:nonce, :any)
+    |> Utils.put_refresh_jwks(client_context, opts)
+  end
+
   @doc false
   @spec send_inactive_token_response(conn :: Plug.Conn.t()) :: Plug.Conn.t()
   def send_inactive_token_response(conn) do

test/oidcc/plug/utils_test.exs

@@ -128,7 +128,30 @@ defmodule Oidcc.Plug.UtilsTest do
     end
   end
 
-  describe "get_refresh_jwks_fun/1" do
+  defmodule ClientStoreWithoutRefresh do
+    @moduledoc false
+    @behaviour Oidcc.Plug.ClientStore
+
+    alias Oidcc.Plug.ClientStore
+
+    @impl ClientStore
+    def get_client_context(_conn), do: {:ok, %{}}
+  end
+
+  defmodule ClientStoreWithRefresh do
+    @moduledoc false
+    @behaviour Oidcc.Plug.ClientStore
+
+    alias Oidcc.Plug.ClientStore
+
+    @impl ClientStore
+    def get_client_context(_conn), do: {:ok, %{}}
+
+    @impl ClientStore
+    def refresh_jwks(context), do: {:refreshed, context}
+  end
+
+  describe "put_refresh_jwks/3" do
     test "uses oidcc_jwt_util for provider configuration" do
       refresh_fun = :test_refresh_fun
 
@@ -138,45 +161,40 @@ defmodule Oidcc.Plug.UtilsTest do
           refresh_fun
         end do
         opts = [provider: :test_provider]
-        assert Utils.get_refresh_jwks_fun(opts) == refresh_fun
+
+        assert Utils.put_refresh_jwks(%{}, :client_context, opts) == %{refresh_jwks: refresh_fun}
       end
     end
 
-    test "returns nil when client_store doesn't implement refresh_jwks" do
-      defmodule ClientStoreWithoutRefresh do
-        @moduledoc false
-        @behaviour Oidcc.Plug.ClientStore
-
-        alias Oidcc.Plug.ClientStore
-
-        @impl ClientStore
-        def get_client_context(_conn), do: {:ok, %{}}
-      end
+    test "adds a fun with the arity oidcc calls it with" do
+      opts = [client_store: ClientStoreWithRefresh]
 
-      opts = [client_store: ClientStoreWithoutRefresh]
+      assert %{refresh_jwks: refresh_jwks} = Utils.put_refresh_jwks(%{}, :client_context, opts)
 
-      assert Utils.get_refresh_jwks_fun(opts) == nil
+      # oidcc invokes the refresh fun as fun(jwks, kid)
+      assert is_function(refresh_jwks, 2)
     end
 
-    test "returns client_store.refresh_jwks function when implemented" do
-      defmodule ClientStoreWithRefresh do
-        @moduledoc false
-        @behaviour Oidcc.Plug.ClientStore
+    test "passes the client context to client_store.refresh_jwks/1" do
+      opts = [client_store: ClientStoreWithRefresh]
 
-        alias Oidcc.Plug.ClientStore
+      %{refresh_jwks: refresh_jwks} = Utils.put_refresh_jwks(%{}, :client_context, opts)
 
-        @impl ClientStore
-        def get_client_context(_conn), do: {:ok, %{}}
+      assert refresh_jwks.(:stale_jwks, "unknown_kid") == {:refreshed, :client_context}
+    end
 
-        @impl ClientStore
-        def refresh_jwks(arg), do: {:refreshed, arg}
-      end
+    test "omits the key for a client_store without the callback" do
+      opts = [client_store: ClientStoreWithoutRefresh]
 
-      opts = [client_store: ClientStoreWithRefresh]
+      # oidcc only skips the refresh when the key is absent, so it must not be
+      # set to nil
+      assert Utils.put_refresh_jwks(%{}, :client_context, opts) == %{}
+    end
+
+    test "keeps existing keys" do
+      opts = [client_store: ClientStoreWithoutRefresh]
 
-      refresh_fun = Utils.get_refresh_jwks_fun(opts)
-      assert is_function(refresh_fun, 1)
-      assert refresh_fun.(:test_arg) == {:refreshed, :test_arg}
+      assert Utils.put_refresh_jwks(%{nonce: :any}, :client_context, opts) == %{nonce: :any}
     end
   end
 end