Convert aegis key cach to LRU with hard expiration time

4 files changed, 327 insertions, 20 deletions

diff --git a/rel/overlay/etc/default.ini b/rel/overlay/etc/default.ini
index a1e3c5851..66680a4e8 100644
--- a/rel/overlay/etc/default.ini
+++ b/rel/overlay/etc/default.ini
@@ -764,3 +764,16 @@ opts = #{budget => 100, target => 2500, window => 60000, sensitivity => 1000}
 ; performance. This value must be at least 10000 and cannot be set to higher than
 ; 10000000, the default transaction size limit.
 ;size_limit = 10000000
+
+[aegis]
+; The maximum number of entries in the key cache.
+; Once the limit reached the least recently used entries are eviceted.
+;cache_limit = 100000
+
+; The period in seconds for how long each entry kept in cache.
+; This is not affected by access time, i.e. the keys are always removed
+; once expired and re-fetched on a next encrypt/decrypt operation.
+;cache_max_age_sec = 1800
+
+; The interval in seconds of how often the expiration check runs.
+;cache_expiration_check_sec = 10
diff --git a/src/aegis/src/aegis_server.erl b/src/aegis/src/aegis_server.erl
index be8202ced..21932626c 100644
--- a/src/aegis/src/aegis_server.erl
+++ b/src/aegis/src/aegis_server.erl
@@ -44,9 +44,13 @@
 -define(KEY_CHECK, aegis_key_check).
 -define(INIT_TIMEOUT, 60000).
 -define(TIMEOUT, 10000).
+-define(CACHE_LIMIT, 100000).
+-define(CACHE_MAX_AGE_SEC, 1800).
+-define(CACHE_EXPIRATION_CHECK_SEC, 10).
+-define(LAST_ACCESSED_INACTIVITY_SEC, 10).
 
 
--record(entry, {uuid, encryption_key}).
+-record(entry, {uuid, encryption_key, counter, last_accessed, expires_at}).
 
 
 start_link() ->
@@ -84,7 +88,7 @@ encrypt(#{} = Db, Key, Value) when is_binary(Key), is_binary(Value) ->
         uuid := UUID
     } = Db,
 
-    case ets:member(?KEY_CHECK, UUID) of
+    case is_key_fresh(UUID) of
         true ->
             case gen_server:call(?MODULE, {encrypt, Db, Key, Value}) of
                 CipherText when is_binary(CipherText) ->
@@ -109,7 +113,7 @@ decrypt(#{} = Db, Key, Value) when is_binary(Key), is_binary(Value) ->
         uuid := UUID
     } = Db,
 
-    case ets:member(?KEY_CHECK, UUID) of
+    case is_key_fresh(UUID) of
         true ->
             case gen_server:call(?MODULE, {decrypt, Db, Key, Value}) of
                 PlainText when is_binary(PlainText) ->
@@ -133,10 +137,16 @@ decrypt(#{} = Db, Key, Value) when is_binary(Key), is_binary(Value) ->
 init([]) ->
     process_flag(sensitive, true),
     Cache = ets:new(?MODULE, [set, private, {keypos, #entry.uuid}]),
+    ByAccess = ets:new(?MODULE,
+        [ordered_set, private, {keypos, #entry.counter}]),
     ets:new(?KEY_CHECK, [named_table, protected, {read_concurrency, true}]),
 
+    erlang:send_after(0, self(), maybe_remove_expired),
+
     St = #{
-        cache => Cache
+        cache => Cache,
+        by_access => ByAccess,
+        counter => 0
     },
     {ok, St, ?INIT_TIMEOUT}.
 
@@ -146,15 +156,18 @@ terminate(_Reason, _St) ->
 
 
 handle_call({insert_key, UUID, DbKey}, _From, #{cache := Cache} = St) ->
-    ok = insert(Cache, UUID, DbKey),
-    {reply, ok, St, ?TIMEOUT};
+    case ets:lookup(Cache, UUID) of
+        [#entry{uuid = UUID} = Entry] ->
+            delete(St, Entry);
+        [] ->
+            ok
+    end,
+    NewSt = insert(St, UUID, DbKey),
+    {reply, ok, NewSt, ?TIMEOUT};
 
 handle_call({encrypt, #{uuid := UUID} = Db, Key, Value}, From, St) ->
-    #{
-        cache := Cache
-    } = St,
 
-    {ok, DbKey} = lookup(Cache, UUID),
+    {ok, DbKey} = lookup(St, UUID),
 
     erlang:spawn(fun() ->
         process_flag(sensitive, true),
@@ -172,11 +185,8 @@ handle_call({encrypt, #{uuid := UUID} = Db, Key, Value}, From, St) ->
     {noreply, St, ?TIMEOUT};
 
 handle_call({decrypt, #{uuid := UUID} = Db, Key, Value}, From, St) ->
-    #{
-        cache := Cache
-    } = St,
 
-    {ok, DbKey} = lookup(Cache, UUID),
+    {ok, DbKey} = lookup(St, UUID),
 
     erlang:spawn(fun() ->
         process_flag(sensitive, true),
@@ -197,10 +207,22 @@ handle_call(_Msg, _From, St) ->
     {noreply, St}.
 
 
+handle_cast({accessed, UUID}, St) ->
+    NewSt = bump_last_accessed(St, UUID),
+    {noreply, NewSt};
+
+
 handle_cast(_Msg, St) ->
     {noreply, St}.
 
 
+handle_info(maybe_remove_expired, St) ->
+    remove_expired_entries(St),
+    CheckInterval = erlang:convert_time_unit(
+        expiration_check_interval(), second, millisecond),
+    erlang:send_after(CheckInterval, self(), maybe_remove_expired),
+    {noreply, St};
+
 handle_info(_Msg, St) ->
     {noreply, St}.
 
@@ -257,19 +279,134 @@ do_decrypt(DbKey, #{uuid := UUID}, Key, Value) ->
     end.
 
 
+is_key_fresh(UUID) ->
+    Now = fabric2_util:now(sec),
+
+    case ets:lookup(?KEY_CHECK, UUID) of
+        [{UUID, ExpiresAt}] when ExpiresAt >= Now -> true;
+        _ -> false
+    end.
+
+
 %% cache functions
 
-insert(Cache, UUID, DbKey) ->
-    Entry = #entry{uuid = UUID, encryption_key = DbKey},
+insert(St, UUID, DbKey) ->
+    #{
+        cache := Cache,
+        by_access := ByAccess,
+        counter := Counter
+    } = St,
+
+    Now = fabric2_util:now(sec),
+    ExpiresAt = Now + max_age(),
+
+    Entry = #entry{
+        uuid = UUID,
+        encryption_key = DbKey,
+        counter = Counter,
+        last_accessed = Now,
+        expires_at = ExpiresAt
+    },
+
     true = ets:insert(Cache, Entry),
-    true = ets:insert(?KEY_CHECK, {UUID, true}),
-    ok.
+    true = ets:insert_new(ByAccess, Entry),
+    true = ets:insert(?KEY_CHECK, {UUID, ExpiresAt}),
+
+    CacheLimit = cache_limit(),
+    CacheSize = ets:info(Cache, size),
 
+    case CacheSize > CacheLimit of
+        true ->
+            LRUKey = ets:first(ByAccess),
+            [LRUEntry] = ets:lookup(ByAccess, LRUKey),
+            delete(St, LRUEntry);
+        false ->
+            ok
+    end,
 
-lookup(Cache, UUID) ->
+    St#{counter := Counter + 1}.
+
+
+lookup(#{cache := Cache}, UUID) ->
     case ets:lookup(Cache, UUID) of
-        [#entry{uuid = UUID, encryption_key = DbKey}] ->
+        [#entry{uuid = UUID, encryption_key = DbKey} = Entry] ->
+            maybe_bump_last_accessed(Entry),
             {ok, DbKey};
         [] ->
             {error, not_found}
     end.
+
+
+delete(St, #entry{uuid = UUID} = Entry) ->
+    #{
+        cache := Cache,
+        by_access := ByAccess
+    } = St,
+
+    true = ets:delete(?KEY_CHECK, UUID),
+    true = ets:delete_object(Cache, Entry),
+    true = ets:delete_object(ByAccess, Entry).
+
+
+maybe_bump_last_accessed(#entry{last_accessed = LastAccessed} = Entry) ->
+    case fabric2_util:now(sec) > LastAccessed + ?LAST_ACCESSED_INACTIVITY_SEC of
+        true ->
+            gen_server:cast(?MODULE, {accessed, Entry#entry.uuid});
+        false ->
+            ok
+    end.
+
+
+bump_last_accessed(St, UUID) ->
+    #{
+        cache := Cache,
+        by_access := ByAccess,
+        counter := Counter
+    } = St,
+
+
+    [#entry{counter = OldCounter} = Entry0] = ets:lookup(Cache, UUID),
+
+    Entry = Entry0#entry{
+        last_accessed = fabric2_util:now(sec),
+        counter = Counter
+    },
+
+    true = ets:insert(Cache, Entry),
+    true = ets:insert_new(ByAccess, Entry),
+
+    ets:delete(ByAccess, OldCounter),
+
+    St#{counter := Counter + 1}.
+
+
+remove_expired_entries(St) ->
+    #{
+        cache := Cache,
+        by_access := ByAccess
+    } = St,
+
+    MatchConditions = [{'=<', '$1', fabric2_util:now(sec)}],
+
+    KeyCheckMatchHead = {'_', '$1'},
+    KeyCheckExpired = [{KeyCheckMatchHead, MatchConditions, [true]}],
+    Count = ets:select_delete(?KEY_CHECK, KeyCheckExpired),
+
+    CacheMatchHead = #entry{expires_at = '$1', _ = '_'},
+    CacheExpired = [{CacheMatchHead, MatchConditions, [true]}],
+    Count = ets:select_delete(Cache, CacheExpired),
+    Count = ets:select_delete(ByAccess, CacheExpired).
+
+
+
+max_age() ->
+    config:get_integer("aegis", "cache_max_age_sec", ?CACHE_MAX_AGE_SEC).
+
+
+expiration_check_interval() ->
+    config:get_integer(
+        "aegis", "cache_expiration_check_sec", ?CACHE_EXPIRATION_CHECK_SEC).
+
+
+cache_limit() ->
+    config:get_integer("aegis", "cache_limit", ?CACHE_LIMIT).
diff --git a/src/aegis/test/aegis_server_test.erl b/src/aegis/test/aegis_server_test.erl
index 0f23a3fd9..0f96798b7 100644
--- a/src/aegis/test/aegis_server_test.erl
+++ b/src/aegis/test/aegis_server_test.erl
@@ -163,3 +163,152 @@ test_disabled_decrypt() ->
     Db = ?DB#{is_encrypted => aegis_server:open_db(?DB)},
     Decrypted = aegis:decrypt(Db, <<1:64>>, ?ENCRYPTED),
     ?assertEqual(?ENCRYPTED, Decrypted).
+
+
+
+lru_cache_with_expiration_test_() ->
+    {
+        foreach,
+        fun() ->
+            %% this has to be be set before start of aegis server
+            %% for config param "cache_expiration_check_sec" to be picked up
+            meck:new([config, aegis_server, fabric2_util], [passthrough]),
+            ok = meck:expect(config, get_integer, fun
+                ("aegis", "cache_limit", _) -> 5;
+                ("aegis", "cache_max_age_sec", _) -> 130;
+                ("aegis", "cache_expiration_check_sec", _) -> 1;
+                (_, _, Default) -> Default
+            end),
+            Ctx = setup(),
+            ok = meck:expect(fabric2_util, now, fun(sec) ->
+                get(time) == undefined andalso put(time, 10),
+                Now = get(time),
+                put(time, Now + 10),
+                Now
+            end),
+            Ctx
+        end,
+        fun teardown/1,
+        [
+            {"counter moves forward on access bump",
+            {timeout, ?TIMEOUT, fun test_advance_counter/0}},
+            {"oldest entries evicted",
+            {timeout, ?TIMEOUT, fun test_evict_old_entries/0}},
+            {"access bump preserves entries",
+            {timeout, ?TIMEOUT, fun test_bump_accessed/0}},
+            {"expired entries removed",
+            {timeout, ?TIMEOUT, fun test_remove_expired/0}}
+        ]
+    }.
+
+
+test_advance_counter() ->
+    ?assertEqual(0, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)),
+
+    ok = meck:expect(aegis_server, handle_cast, fun({accessed, _} = Msg, St) ->
+        #{counter := Counter} = St,
+        get(counter) == undefined andalso put(counter, 0),
+        OldCounter = get(counter),
+        put(counter, Counter),
+        ?assert(Counter > OldCounter),
+        meck:passthrough([Msg, St])
+    end),
+
+    lists:foreach(fun(I) ->
+        Db = ?DB#{uuid => <<I:64>>},
+        aegis_server:encrypt(Db, <<I:64>>, ?VALUE),
+        aegis_server:encrypt(Db, <<(I+1):64>>, ?VALUE)
+    end, lists:seq(1, 10)),
+
+    ?assertEqual(10, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)).
+
+
+test_evict_old_entries() ->
+    ?assertEqual(0, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)),
+
+    %% overflow cache
+    lists:foreach(fun(I) ->
+        Db = ?DB#{uuid => <<I:64>>},
+        aegis_server:encrypt(Db, <<I:64>>, ?VALUE)
+    end, lists:seq(1, 10)),
+
+    ?assertEqual(10, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)),
+
+    %% confirm that newest keys are still in cache
+    lists:foreach(fun(I) ->
+        Db = ?DB#{uuid => <<I:64>>},
+        aegis_server:encrypt(Db, <<(I+1):64>>, ?VALUE)
+    end, lists:seq(6, 10)),
+
+    ?assertEqual(10, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)),
+
+    %% confirm that oldest keys been eviced and needed re-fetch
+    lists:foreach(fun(I) ->
+        Db = ?DB#{uuid => <<I:64>>},
+        aegis_server:encrypt(Db, <<(I+1):64>>, ?VALUE)
+    end, lists:seq(1, 5)),
+
+    ?assertEqual(15, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)).
+
+
+test_bump_accessed() ->
+    ?assertEqual(0, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)),
+
+    %% fill the cache
+    lists:foreach(fun(I) ->
+        Db = ?DB#{uuid => <<I:64>>},
+        aegis_server:encrypt(Db, <<I:64>>, ?VALUE)
+    end, lists:seq(1, 5)),
+
+    ?assertEqual(5, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)),
+
+    %% bump oldest key and then insert a new key to trigger eviction
+    aegis_server:encrypt(?DB#{uuid => <<1:64>>}, <<1:64>>, ?VALUE),
+    aegis_server:encrypt(?DB#{uuid => <<6:64>>}, <<6:64>>, ?VALUE),
+    ?assertEqual(6, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)),
+
+    %% confirm that former oldest key is still in cache
+    aegis_server:encrypt(?DB#{uuid => <<1:64>>}, <<2:64>>, ?VALUE),
+    ?assertEqual(6, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)),
+
+    %% confirm that the second oldest key been evicted by the new insert
+    aegis_server:encrypt(?DB#{uuid => <<2:64>>}, <<3:64>>, ?VALUE),
+    ?assertEqual(7, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)).
+
+
+test_remove_expired() ->
+    ?assertEqual(0, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)),
+
+    %% to detect when maybe_remove_expired called
+    ok = meck:expect(aegis_server, handle_info,fun
+        (maybe_remove_expired, St) ->
+            meck:passthrough([maybe_remove_expired, St])
+    end),
+
+    %% fill the cache. first key expires a 140, last at 180 of "our" time
+    lists:foreach(fun(I) ->
+        Db = ?DB#{uuid => <<I:64>>},
+        aegis_server:encrypt(Db, <<I:64>>, ?VALUE)
+    end, lists:seq(1, 5)),
+
+    ?assertEqual(5, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)),
+
+    %% confirm enties are still in cache and wind up our "clock" to 160
+    lists:foreach(fun(I) ->
+        Db = ?DB#{uuid => <<I:64>>},
+        aegis_server:encrypt(Db, <<I:64>>, ?VALUE)
+    end, lists:seq(1, 5)),
+
+    ?assertEqual(5, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)),
+
+    %% wait for remove_expired_entries to be triggered
+    meck:reset(aegis_server),
+    meck:wait(aegis_server, handle_info, [maybe_remove_expired, '_'], 2500),
+
+    %% 3 "oldest" entries should be removed, 2 yet to expire still in cache
+    lists:foreach(fun(I) ->
+        Db = ?DB#{uuid => <<I:64>>},
+        aegis_server:encrypt(Db, <<I:64>>, ?VALUE)
+    end, lists:seq(1, 5)),
+
+    ?assertEqual(8, meck:num_calls(?AEGIS_KEY_MANAGER, open_db, 1)).
diff --git a/src/fabric/src/fabric2_util.erl b/src/fabric/src/fabric2_util.erl
index 9b6d18c58..136762b34 100644
--- a/src/fabric/src/fabric2_util.erl
+++ b/src/fabric/src/fabric2_util.erl
@@ -41,6 +41,7 @@
     all_docs_view_opts/1,
 
     iso8601_timestamp/0,
+    now/1,
     do_recovery/0,
 
     pmap/2,
@@ -348,6 +349,13 @@ iso8601_timestamp() ->
     io_lib:format(Format, [Year, Month, Date, Hour, Minute, Second]).
 
 
+now(ms) ->
+    {Mega, Sec, Micro} = os:timestamp(),
+    (Mega * 1000000 + Sec) * 1000 + round(Micro / 1000);
+now(sec) ->
+    now(ms) div 1000.
+
+
 do_recovery() ->
     config:get_boolean("couchdb",
         "enable_database_recovery", false).