Extract aegis_keywrap parts in shim of key manager

3 files changed, 74 insertions, 19 deletions

diff --git a/src/aegis/src/aegis.erl b/src/aegis/src/aegis.erl
index ca38a7d5b..2817564e5 100644
--- a/src/aegis/src/aegis.erl
+++ b/src/aegis/src/aegis.erl
@@ -17,6 +17,8 @@
 
 -define(WRAPPED_KEY, {?DB_AEGIS, 1}).
 
+-define(CACHE, aegis_key_cache).
+
 
 -export([
     create/2,
@@ -37,7 +39,7 @@ create(#{} = Db, _Options) ->
     } = Db,
 
     % Fetch unwrapped key
-    WrappedKey = gen_server:call(aegis_key_cache, {get_wrapped_key, Db}),
+    WrappedKey = gen_server:call(?CACHE, {get_wrapped_key, Db}),
 
     % And store it
     FDBKey = erlfdb_tuple:pack(?WRAPPED_KEY, DbPrefix),
@@ -58,18 +60,23 @@ open(#{} = Db, _Options) ->
     FDBKey = erlfdb_tuple:pack(?WRAPPED_KEY, DbPrefix),
     WrappedKey = erlfdb:wait(erlfdb:get(Tx, FDBKey)),
 
-    %% maybe ask to rewrap and store if updated?
+    Db1 = Db#{aegis => WrappedKey},
 
-    Db#{
-        aegis => WrappedKey
-    }.
+    case gen_server:call(?CACHE, {unwrap_key, Db1}) of
+        WrappedKey ->
+            Db1;
+        NewWrappedKey ->
+            FDBKey = erlfdb_tuple:pack(?WRAPPED_KEY, DbPrefix),
+            ok = erlfdb:set(Tx, FDBKey, NewWrappedKey),
+            Db1#{aegis => NewWrappedKey}
+    end.
 
 
 encrypt(#{} = _Db, _Key, <<>>) ->
     <<>>;
 
 encrypt(#{} = Db, Key, Value) when is_binary(Key), is_binary(Value) ->
-    gen_server:call(aegis_key_cache, {encrypt, Db, Key, Value}).
+    gen_server:call(?CACHE, {encrypt, Db, Key, Value}).
 
 encrypt(DbKey, UUID, Key, Value) ->
     EncryptionKey = crypto:strong_rand_bytes(32),
@@ -93,7 +100,7 @@ decrypt(#{} = _Db, _Key, <<>>) ->
     <<>>;
 
 decrypt(#{} = Db, Key, Value) when is_binary(Key), is_binary(Value) ->
-    gen_server:call(aegis_key_cache, {decrypt, Db, Key, Value}).
+    gen_server:call(?CACHE, {decrypt, Db, Key, Value}).
 
 decrypt(DbKey, UUID, Key, Value) ->
     case Value of
diff --git a/src/aegis/src/aegis_key_cache.erl b/src/aegis/src/aegis_key_cache.erl
index aac946c0e..ed2de58f3 100644
--- a/src/aegis/src/aegis_key_cache.erl
+++ b/src/aegis/src/aegis_key_cache.erl
@@ -91,6 +91,18 @@ handle_call({get_wrapped_key, Db}, From, #{clients := Clients} = St) ->
     Clients1 = dict:store(Ref, From, Clients),
     {noreply, St#{clients := Clients1}, ?TIMEOUT};
 
+handle_call({unwrap_key, #{aegis := WrappedKey} = Db}, From, St) ->
+    #{
+        clients := Clients,
+        unwrappers := Unwrappers
+    } = St,
+
+    {_Pid, Ref} = erlang:spawn_monitor(?MODULE, unwrap_key, [Db]),
+
+    Clients1 = dict:store(Ref, From, Clients),
+    Unwrappers1 = dict:store(WrappedKey, Ref, Unwrappers),
+    {noreply, St#{clients := Clients1, unwrappers := Unwrappers1}, ?TIMEOUT};
+
 handle_call({encrypt, Db, Key, Value}, From, St) ->
     NewSt = maybe_spawn_worker(St, From, do_encrypt, Db, Key, Value),
     {noreply, NewSt, ?TIMEOUT};
@@ -115,13 +127,19 @@ handle_info({'DOWN', Ref, _, _Pid, {key, {ok, DbKey, WrappedKey}}}, St) ->
         unwrappers := Unwrappers
     } = St,
 
+    IsGetWrappedKeyClient = dict:is_key(Ref, Clients),
+
     NewSt1 = case dict:take(WrappedKey, Unwrappers) of
         {Ref, Unwrappers1} ->
             ok = insert(Cache, WrappedKey, DbKey),
             St#{unwrappers := Unwrappers1};
+        error when IsGetWrappedKeyClient ->
+            ok = insert(Cache, WrappedKey, DbKey),
+            St;
         error ->
             %% FIXME! it might be new wrapped key != old wrapped key
-            %% fold here to search for it based on ref
+            %% fold over Unwrappers here to find waiters of old key
+            %% by Ref. also need way to store new wrapped key in fdb
             St
     end,
 
@@ -222,12 +240,10 @@ maybe_reply(#{clients := Clients} = St, Ref, Resp) ->
     end.
 
 
-get_wrapped_key(#{} = _Db) ->
+get_wrapped_key(#{} = Db) ->
     process_flag(sensitive, true),
     try
-        DbKey = crypto:strong_rand_bytes(32),
-        WrappedKey = aegis_keywrap:key_wrap(?ROOT_KEY, DbKey),
-        {ok, DbKey, WrappedKey}
+        ?AEGIS_KEY_MANAGER:key_wrap(Db)
     of
         Resp ->
             exit({key, Resp})
@@ -237,15 +253,10 @@ get_wrapped_key(#{} = _Db) ->
     end.
 
 
-unwrap_key(#{aegis := WrappedKey} = _Db) ->
+unwrap_key(#{aegis := WrappedKey} = Db) ->
     process_flag(sensitive, true),
     try
-        case aegis_keywrap:key_unwrap(?ROOT_KEY, WrappedKey) of
-            fail ->
-                error(decryption_failed);
-            DbKey ->
-                {ok, DbKey, WrappedKey}
-        end
+        ?AEGIS_KEY_MANAGER:key_unwrap(Db)
     of
         Resp ->
             exit({key, Resp})
diff --git a/src/aegis/src/aegis_key_manager.erl b/src/aegis/src/aegis_key_manager.erl
new file mode 100644
index 000000000..e41cfafff
--- /dev/null
+++ b/src/aegis/src/aegis_key_manager.erl
@@ -0,0 +1,37 @@
+% Licensed under the Apache License, Version 2.0 (the "License"); you may not
+% use this file except in compliance with the License. You may obtain a copy of
+% the License at
+%
+%   http://www.apache.org/licenses/LICENSE-2.0
+%
+% Unless required by applicable law or agreed to in writing, software
+% distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+% WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+% License for the specific language governing permissions and limitations under
+% the License.
+
+-module(aegis_key_manager).
+
+
+-export([
+    key_wrap/1,
+    key_unwrap/1
+]).
+
+
+-define(ROOT_KEY, <<1:256>>).
+
+
+key_wrap(#{} = _Db) ->
+    DbKey = crypto:strong_rand_bytes(32),
+    WrappedKey = aegis_keywrap:key_wrap(?ROOT_KEY, DbKey),
+    {ok, DbKey, WrappedKey}.
+
+
+key_unwrap(#{aegis := WrappedKey} = _Db) ->
+    case aegis_keywrap:key_unwrap(?ROOT_KEY, WrappedKey) of
+        fail ->
+            error(decryption_failed);
+        DbKey ->
+            {ok, DbKey, WrappedKey}
+    end.