wiparchive/aegis_key_cache_rnewsonaegis_key_cache_rnewson

1 files changed, 51 insertions, 160 deletions

diff --git a/src/aegis/src/aegis_server.erl b/src/aegis/src/aegis_server.erl
index 04cbf0e23..29b816248 100644
--- a/src/aegis/src/aegis_server.erl
+++ b/src/aegis/src/aegis_server.erl
@@ -41,8 +41,6 @@
 
 %% workers callbacks
 -export([
-    do_init_db/2,
-    do_open_db/2,
     do_encrypt/5,
     do_decrypt/5
 ]).
@@ -52,7 +50,7 @@
 -define(TIMEOUT, 10000).
 
 
--record(entry, {id, key}).
+-record(entry, {uuid, encryption_key}).
 
 
 start_link() ->
@@ -61,12 +59,35 @@ start_link() ->
 
 -spec init_db(Db :: #{}, Options :: list()) -> boolean().
 init_db(#{} = Db, Options) ->
-    gen_server:call(?MODULE, {init_db, Db, Options}).
+    #{
+        uuid := UUID
+    } = Db,
+    case ?AEGIS_KEY_MANAGER:init_db(Db, Options) of
+        {ok, Key} ->
+            ok = gen_server:call(?MODULE, {insert_key, UUID, Key}),
+            true;
+        false ->
+            false
+    end.
 
 
 -spec open_db(Db :: #{}, Options :: list()) -> boolean().
 open_db(#{} = Db, Options) ->
-    gen_server:call(?MODULE, {open_db, Db, Options}).
+    #{
+        uuid := UUID
+    } = Db,
+    case gen_server:call(?MODULE, {has_key, UUID}) of
+        true ->
+            true;
+        false ->
+            case ?AEGIS_KEY_MANAGER:open_db(Db, Options) of
+                {ok, Key} ->
+                    ok = gen_server:call(?MODULE, {insert_key, UUID, Key}),
+                    true;
+                false ->
+                    false
+            end
+    end.
 
 
 -spec encrypt(Db :: #{}, Key :: binary(), Value :: binary()) -> binary().
@@ -83,61 +104,41 @@ decrypt(#{} = Db, Key, Value) when is_binary(Key), is_binary(Value) ->
 
 init([]) ->
     process_flag(sensitive, true),
-    Cache = ets:new(?MODULE, [set, private, {keypos, #entry.id}]),
+    Cache = ets:new(?MODULE, [set, private, {keypos, #entry.uuid}]),
 
     St = #{
-        cache => Cache,
-        openers => dict:new(),
-        waiters => dict:new(),
-        unwrappers => dict:new()
+        cache => Cache
     },
     {ok, St, ?INIT_TIMEOUT}.
 
 
-terminate(_Reason, St) ->
-    #{
-        openers := Openers,
-        waiters := Waiters
-    } = St,
-
-    dict:fold(fun(_AegisConfig, WaitList, _) ->
-        lists:foreach(fun(#{from := From}) ->
-            gen_server:reply(From, {error, operation_aborted})
-        end, WaitList)
-    end, ok, Waiters),
-
-    dict:fold(fun(Ref, From, _) ->
-        erlang:demonitor(Ref),
-        gen_server:reply(From, {error, operation_aborted})
-    end, ok, Openers),
+terminate(_Reason, _St) ->
     ok.
 
-
-handle_call({init_db, Db, Options}, From, #{openers := Openers} = St) ->
+handle_call({insert_key, UUID, EncryptionKey}, _From, St) ->
     #{
-        uuid := UUID
-    } = Db,
-
-    {_, Ref} = erlang:spawn_monitor(?MODULE, do_init_db, [Db, Options]),
-    Openers1 = dict:store(Ref, {UUID, From}, Openers),
-    {noreply, St#{openers := Openers1}, ?TIMEOUT};
+        cache := Cache
+    } = St,
+    true = insert(Cache, UUID, EncryptionKey),
+    {reply, ok, St, ?TIMEOUT};
 
-handle_call({open_db, Db, Options}, From, #{openers := Openers} = St) ->
+handle_call({has_key, UUID}, _From, St) ->
     #{
-        uuid := UUID
-    } = Db,
+        cache := Cache
+    } = St,
+    case lookup(Cache, UUID) of
+        {ok, _Key} ->
+            {reply, true, St, ?TIMEOUT};
+        {error, not_found} ->
+            {reply, false, St, ?TIMEOUT}
+    end.
 
-    {_, Ref} = erlang:spawn_monitor(?MODULE, do_open_db, [Db, Options]),
-    Openers1 = dict:store(Ref, {UUID, From}, Openers),
-    {noreply, St#{openers := Openers1}, ?TIMEOUT};
 
 handle_call({encrypt, Db, Key, Value}, From, St) ->
-    NewSt = maybe_spawn_worker(St, From, do_encrypt, Db, Key, Value),
-    {noreply, NewSt, ?TIMEOUT};
+    maybe_spawn_worker(St, From, do_encrypt, Db, Key, Value);
 
 handle_call({decrypt, Db, Key, Value}, From, St) ->
-    NewSt = maybe_spawn_worker(St, From, do_decrypt, Db, Key, Value),
-    {noreply, NewSt, ?TIMEOUT};
+    maybe_spawn_worker(St, From, do_decrypt, Db, Key, Value);
 
 handle_call(_Msg, _From, St) ->
     {noreply, St}.
@@ -147,65 +148,6 @@ handle_cast(_Msg, St) ->
     {noreply, St}.
 
 
-handle_info({'DOWN', Ref, _, _Pid, false}, #{openers := Openers} = St) ->
-    {{_UUID, From}, Openers1} = dict:take(Ref, Openers),
-    gen_server:reply(From, false),
-    {noreply, St#{openers := Openers1}, ?TIMEOUT};
-
-handle_info({'DOWN', Ref, _, _Pid, {ok, DbKey}}, St) ->
-    #{
-        cache := Cache,
-        openers := Openers,
-        waiters := Waiters,
-        unwrappers := Unwrappers
-    } = St,
-
-    case dict:take(Ref, Openers) of
-        {{UUID, From}, Openers1} ->
-            ok = insert(Cache, UUID, DbKey),
-            gen_server:reply(From, true),
-            {noreply, St#{openers := Openers1}, ?TIMEOUT};
-        error ->
-            {UUID, Unwrappers1} = dict:take(Ref, Unwrappers),
-            ok = insert(Cache, UUID, DbKey),
-            Unwrappers2 = dict:erase(UUID, Unwrappers1),
-
-            {WaitList, Waiters1} = dict:take(UUID, Waiters),
-            lists:foreach(fun(Waiter) ->
-                #{
-                    from := From,
-                    action := Action,
-                    args := Args
-                } = Waiter,
-                erlang:spawn(?MODULE, Action, [From, DbKey | Args])
-            end, WaitList),
-            NewSt = St#{waiters := Waiters1, unwrappers := Unwrappers2},
-            {noreply, NewSt, ?TIMEOUT}
-    end;
-
-handle_info({'DOWN', Ref, process, _Pid, {error, Error}}, St) ->
-    #{
-        openers := Openers,
-        waiters := Waiters,
-        unwrappers := Unwrappers
-    } = St,
-
-    case dict:take(Ref, Openers) of
-        {From, Openers1} ->
-            gen_server:reply(From, {error, Error}),
-            {noreply, St#{openers := Openers1}, ?TIMEOUT};
-        error ->
-            {UUID, Unwrappers1} = dict:take(Ref, Unwrappers),
-            Unwrappers2 = dict:erase(UUID, Unwrappers1),
-
-            {WaitList, Waiters1} = dict:take(UUID, Waiters),
-            lists:foreach(fun(#{from := From}) ->
-                gen_server:reply(From, {error, Error})
-            end, WaitList),
-            NewSt = St#{waiters := Waiters1, unwrappers := Unwrappers2},
-            {noreply, NewSt, ?TIMEOUT}
-    end;
-
 handle_info(_Msg, St) ->
     {noreply, St}.
 
@@ -216,32 +158,6 @@ code_change(_OldVsn, St, _Extra) ->
 
 %% workers functions
 
-do_init_db(#{} = Db, Options) ->
-    process_flag(sensitive, true),
-    try
-        ?AEGIS_KEY_MANAGER:init_db(Db, Options)
-    of
-        Resp ->
-            exit(Resp)
-    catch
-        _:Error ->
-            exit({error, Error})
-    end.
-
-
-do_open_db(#{} = Db, Options) ->
-    process_flag(sensitive, true),
-    try
-        ?AEGIS_KEY_MANAGER:open_db(Db, Options)
-    of
-        Resp ->
-            exit(Resp)
-    catch
-        _:Error ->
-            exit({error, Error})
-    end.
-
-
 do_encrypt(From, DbKey, #{uuid := UUID}, Key, Value) ->
     process_flag(sensitive, true),
     try
@@ -300,53 +216,28 @@ do_decrypt(From, DbKey, #{uuid := UUID}, Key, Value) ->
 
 maybe_spawn_worker(St, From, Action, #{uuid := UUID} = Db, Key, Value) ->
     #{
-        cache := Cache,
-        waiters := Waiters
+        cache := Cache
     } = St,
 
     case lookup(Cache, UUID) of
         {ok, DbKey} ->
             erlang:spawn(?MODULE, Action, [From, DbKey, Db, Key, Value]),
-            St;
+            {noreply, St, ?TIMEOUT};
         {error, not_found} ->
-            NewSt = maybe_spawn_unwrapper(St, Db),
-            Waiter = #{
-                from => From,
-                action => Action,
-                args => [Db, Key, Value]
-            },
-            Waiters1 = dict:append(UUID, Waiter, Waiters),
-            NewSt#{waiters := Waiters1}
+            {reply, {error, encryption_key_not_found}, St, ?TIMEOUT}
      end.
 
 
-maybe_spawn_unwrapper(St, #{uuid := UUID} = Db) ->
-    #{
-        unwrappers := Unwrappers
-    } = St,
-
-    case dict:is_key(UUID, Unwrappers) of
-        true ->
-            St;
-        false ->
-            {_Pid, Ref} = erlang:spawn_monitor(?MODULE, do_unwrap_key, [Db]),
-            Unwrappers1 = dict:store(UUID, Ref, Unwrappers),
-            Unwrappers2 = dict:store(Ref, UUID, Unwrappers1),
-            St#{unwrappers := Unwrappers2}
-    end.
-
-
 %% cache functions
 
 insert(Cache, UUID, DbKey) ->
-    Entry = #entry{id = UUID, key = DbKey},
-    true = ets:insert(Cache, Entry),
-    ok.
+    Entry = #entry{uuid = UUID, encryption_key = DbKey},
+    ets:insert(Cache, Entry).
 
 
 lookup(Cache, UUID) ->
     case ets:lookup(Cache, UUID) of
-        [#entry{id = UUID, key = DbKey}] ->
+        [#entry{uuid = UUID, encryption_key = DbKey}] ->
             {ok, DbKey};
         [] ->
             {error, not_found}