From f238e8f83dc7b77f0424d4f5b1e632a2826f7c46 Mon Sep 17 00:00:00 2001
From: Ilya Averyanov <av@rubybox.ru>
Date: Mon, 27 Jun 2022 23:00:19 +0300
Subject: [PATCH] fix: update JWKS renewal timer correctly

---
 CHANGES-4.3.md                                |  1 +
 apps/emqx_auth_jwt/src/emqx_auth_jwt_svr.erl  |  8 ++++--
 .../test/emqx_auth_jwt_SUITE.erl              | 27 +++++++++++++++++++
 3 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/CHANGES-4.3.md b/CHANGES-4.3.md
index 7799eedda..cbbfc8d8e 100644
--- a/CHANGES-4.3.md
+++ b/CHANGES-4.3.md
@@ -32,6 +32,7 @@ File format:
   subscribed to. [#8288]
 - Ensuring that exhook dispatches the client events are sequential. [#8311]
 - Ensure start dashboard ok event if default_username is missing.
+- Fix key update from JWKS server by JWT auth.
 
 ## v4.3.15
 
diff --git a/apps/emqx_auth_jwt/src/emqx_auth_jwt_svr.erl b/apps/emqx_auth_jwt/src/emqx_auth_jwt_svr.erl
index c018bf1bd..b82ca0b48 100644
--- a/apps/emqx_auth_jwt/src/emqx_auth_jwt_svr.erl
+++ b/apps/emqx_auth_jwt/src/emqx_auth_jwt_svr.erl
@@ -20,6 +20,7 @@
 
 -include_lib("emqx/include/logger.hrl").
 -include_lib("jose/include/jose_jwk.hrl").
+-include_lib("snabbkaffe/include/snabbkaffe.hrl").
 
 -logger_header("[JWT-SVR]").
 
@@ -114,7 +115,8 @@ handle_cast(_Msg, State) ->
 
 handle_info({timeout, _TRef, refresh}, State = #state{addr = Addr}) ->
     NState = try
-                 true = ets:insert(?TAB, {remote, request_jwks(Addr)})
+                 true = ets:insert(?TAB, {remote, request_jwks(Addr)}),
+                 State
              catch _:_ ->
                  State
              end,
@@ -147,7 +149,9 @@ request_jwks(Addr) ->
         {ok, {_Code, _Headers, Body}} ->
             try
                 JwkSet = jose_jwk:from(emqx_json:decode(Body, [return_maps])),
-                {_, Jwks} = JwkSet#jose_jwk.keys, Jwks
+                {_, Jwks} = JwkSet#jose_jwk.keys,
+                ?tp(debug, emqx_auth_jwt_svr_jwks_updated, #{jwks => Jwks, pid => self()}),
+                Jwks
             catch _:_ ->
                 ?LOG(error, "Invalid jwks server response: ~p~n", [Body]),
                 error(badarg)
diff --git a/apps/emqx_auth_jwt/test/emqx_auth_jwt_SUITE.erl b/apps/emqx_auth_jwt/test/emqx_auth_jwt_SUITE.erl
index 2136d5384..1287f43e7 100644
--- a/apps/emqx_auth_jwt/test/emqx_auth_jwt_SUITE.erl
+++ b/apps/emqx_auth_jwt/test/emqx_auth_jwt_SUITE.erl
@@ -22,6 +22,7 @@
 -include_lib("emqx/include/emqx_mqtt.hrl").
 -include_lib("eunit/include/eunit.hrl").
 -include_lib("common_test/include/ct.hrl").
+-include_lib("snabbkaffe/include/snabbkaffe.hrl").
 
 all() -> emqx_ct:all(?MODULE).
 
@@ -201,6 +202,32 @@ t_check_claims_kid_in_header(_Config) ->
     ct:pal("Auth result: ~p~n", [Result0]),
     ?assertMatch({ok, #{auth_result := success, jwt_claims := _}}, Result0).
 
+t_keys_update(init, _Config) ->
+    ok = meck:new(httpc, [passthrough, no_history]),
+    ok = meck:expect(
+           httpc,
+           request,
+           fun(get, _, _, _) ->
+                   {ok,
+                    {200,
+                     [],
+                     jiffy:encode(#{<<"keys">> => []})}}
+           end),
+
+    application:set_env(emqx_auth_jwt, verify_claims, []),
+    application:set_env(emqx_auth_jwt, refresh_interval, 100),
+    application:set_env(emqx_auth_jwt, jwks, "http://localhost:4001/keys.json").
+t_keys_update(_Config) ->
+    ?check_trace(
+       snabbkaffe:block_until(
+         ?match_n_events(2, #{?snk_kind := emqx_auth_jwt_svr_jwks_updated}),
+         _Timeout    = infinity,
+         _BackInTIme = 0),
+       fun(_, Trace) ->
+               ?assertMatch([#{pid := Pid}, #{pid := Pid} |  _],
+                            ?of_kind(emqx_auth_jwt_svr_jwks_updated, Trace))
+       end).
+
 t_check_jwt_acl(init, _Config) ->
     application:set_env(emqx_auth_jwt, verify_claims, [{sub, <<"value">>}]).
 t_check_jwt_acl(_Config) ->