provide caching of JWKS keys

4 files changed, 119 insertions, 0 deletions

diff --git a/src/jwks.erl b/src/jwks.erl
index d694d2e7b..d6c44deb4 100644
--- a/src/jwks.erl
+++ b/src/jwks.erl
@@ -16,11 +16,42 @@
 -module(jwks).
 
 -export([
+    get_key/3,
     get_keyset/1
 ]).
 
 -include_lib("public_key/include/public_key.hrl").
 
+get_key(Url, Kty, Kid) ->
+    case lookup(Url, Kty, Kid) of
+        {ok, Key} ->
+            {ok, Key};
+        {error, not_found} ->
+            update_cache(Url),
+            lookup(Url, Kty, Kid)
+    end.
+
+
+lookup(Url, Kty, Kid) ->
+    case ets_lru:lookup_d(jwks_cache_lru, {Url, Kty, Kid}) of
+        {ok, Key} ->
+            {ok, Key};
+        not_found ->
+            {error, not_found}
+    end.
+
+
+update_cache(Url) ->
+    case get_keyset(Url) of
+        {ok, KeySet} ->
+            [ets_lru:insert(jwks_cache_lru, {Url, Kty, Kid}, Key)
+                || {{Kty, Kid}, Key} <- KeySet],
+            ok;
+        {error, Reason} ->
+            {error, Reason}
+    end.
+
+
 get_keyset(Url) ->
     ReqHeaders = [],
     case ibrowse:send_req(Url, ReqHeaders, get) of
diff --git a/src/jwtf.app.src b/src/jwtf.app.src
index 304bb9e0a..5fd9c2562 100644
--- a/src/jwtf.app.src
+++ b/src/jwtf.app.src
@@ -14,11 +14,13 @@
     {description, "JSON Web Token Functions"},
     {vsn, git},
     {registered, []},
+    {mod, { jwtf_app, []}},
     {applications, [
         kernel,
         stdlib,
         b64url,
         crypto,
+        ets_lru,
         jiffy,
         public_key
     ]},
diff --git a/src/jwtf_app.erl b/src/jwtf_app.erl
new file mode 100644
index 000000000..92a26d558
--- /dev/null
+++ b/src/jwtf_app.erl
@@ -0,0 +1,26 @@
+%%%-------------------------------------------------------------------
+%% @doc jwtf public API
+%% @end
+%%%-------------------------------------------------------------------
+
+-module(jwtf_app).
+
+-behaviour(application).
+
+%% Application callbacks
+-export([start/2, stop/1]).
+
+%%====================================================================
+%% API
+%%====================================================================
+
+start(_StartType, _StartArgs) ->
+    jwtf_sup:start_link().
+
+%%--------------------------------------------------------------------
+stop(_State) ->
+    ok.
+
+%%====================================================================
+%% Internal functions
+%%====================================================================
diff --git a/src/jwtf_sup.erl b/src/jwtf_sup.erl
new file mode 100644
index 000000000..2256ac53a
--- /dev/null
+++ b/src/jwtf_sup.erl
@@ -0,0 +1,60 @@
+%%%-------------------------------------------------------------------
+%% @doc epep top level supervisor.
+%% @end
+%%%-------------------------------------------------------------------
+
+-module(jwtf_sup).
+
+-behaviour(supervisor).
+
+%% API
+-export([start_link/0]).
+
+%% Supervisor callbacks
+-export([init/1]).
+
+-define(SERVER, ?MODULE).
+
+%%====================================================================
+%% API functions
+%%====================================================================
+
+start_link() ->
+    supervisor:start_link({local, ?SERVER}, ?MODULE, []).
+
+%%====================================================================
+%% Supervisor callbacks
+%%====================================================================
+
+%% Child :: {Id,StartFunc,Restart,Shutdown,Type,Modules}
+init([]) ->
+    Children = [
+        {jwks_cache_lru,
+            {ets_lru, start_link, [jwks_cache_lru, lru_opts()]},
+            permanent, 5000, worker, [ets_lru]}
+    ],
+    {ok, { {one_for_all, 0, 1}, Children} }.
+
+%%====================================================================
+%% Internal functions
+%%====================================================================
+
+lru_opts() ->
+    case config:get_integer("jwtf_cache", "max_objects", 50) of
+        MxObjs when MxObjs > 0 ->
+            [{max_objects, MxObjs}];
+        _ ->
+            []
+    end ++
+    case config:get_integer("jwtf_cache", "max_size", 0) of
+        MxSize when MxSize > 0 ->
+            [{max_size, MxSize}];
+        _ ->
+            []
+    end ++
+    case config:get_integer("jwtf_cache", "max_lifetime", 0) of
+        MxLT when MxLT > 0 ->
+            [{max_lifetime, MxLT}];
+        _ ->
+            []
+    end.