Pass db options to generate_key and rename WrappedKey to AegisConfig

5 files changed, 87 insertions, 70 deletions

diff --git a/src/aegis/src/aegis.erl b/src/aegis/src/aegis.erl
index 0315c7b88..b387324a1 100644
--- a/src/aegis/src/aegis.erl
+++ b/src/aegis/src/aegis.erl
@@ -27,34 +27,35 @@
     wrap_fold_fun/2
 ]).
 
-create(#{} = Db, _Options) ->
+create(#{} = Db, Options) ->
     #{
         tx := Tx,
         db_prefix := DbPrefix
     } = Db,
 
-    {ok, WrappedKey} = aegis_server:generate_key(Db),
+    {ok, AegisConfig} = aegis_server:generate_key(Db, Options),
 
     FDBKey = erlfdb_tuple:pack(?WRAPPED_KEY, DbPrefix),
-    ok = erlfdb:set(Tx, FDBKey, WrappedKey),
+    Packed = erlfdb_tuple:pack({AegisConfig}),
+    ok = erlfdb:set(Tx, FDBKey, Packed),
 
     Db#{
-        aegis => WrappedKey
+        aegis => AegisConfig
     }.
 
 
-open(#{} = Db, Options) ->
+open(#{} = Db, _Options) ->
     #{
         tx := Tx,
         db_prefix := DbPrefix
     } = Db,
 
-    % Fetch wrapped key
     FDBKey = erlfdb_tuple:pack(?WRAPPED_KEY, DbPrefix),
-    WrappedKey = erlfdb:wait(erlfdb:get(Tx, FDBKey)),
+    Packed = erlfdb:wait(erlfdb:get(Tx, FDBKey)),
+    {AegisConfig} = erlfdb_tuple:unpack(Packed),
 
     Db#{
-        aegis => WrappedKey
+        aegis => AegisConfig
     }.
 
 
diff --git a/src/aegis/src/aegis_file_key_manager.erl b/src/aegis/src/aegis_file_key_manager.erl
index f520bd497..6522c76c5 100644
--- a/src/aegis/src/aegis_file_key_manager.erl
+++ b/src/aegis/src/aegis_file_key_manager.erl
@@ -17,7 +17,7 @@
 
 
 -export([
-    generate_key/1,
+    generate_key/2,
     unwrap_key/2
 ]).
 
@@ -25,16 +25,18 @@
 -define(ROOT_KEY, <<1:256>>).
 
 
-generate_key(#{} = _Db) ->
+generate_key(#{} = _Db, _Options) ->
     DbKey = crypto:strong_rand_bytes(32),
     WrappedKey = aegis_keywrap:key_wrap(?ROOT_KEY, DbKey),
-    {ok, DbKey, WrappedKey}.
+    %% just an example of how to represent the arbitrary options
+    AegisConfig = {<<"wrapped_key">>, WrappedKey},
+    {ok, DbKey, AegisConfig}.
 
 
-unwrap_key(#{} = _Db, WrappedKey) ->
+unwrap_key(#{} = _Db, {<<"wrapped_key">>, WrappedKey} = AegisConfig) ->
     case aegis_keywrap:key_unwrap(?ROOT_KEY, WrappedKey) of
         fail ->
             error(unwrap_failed);
         DbKey ->
-            {ok, DbKey, WrappedKey}
+            {ok, DbKey, AegisConfig}
     end.
diff --git a/src/aegis/src/aegis_key_manager.erl b/src/aegis/src/aegis_key_manager.erl
index a16c51690..1ae81ab1d 100644
--- a/src/aegis/src/aegis_key_manager.erl
+++ b/src/aegis/src/aegis_key_manager.erl
@@ -14,24 +14,24 @@
 
 
 -type key() :: binary().
--type wrapped_key() :: binary().
+-type aegis_config() :: term().
 
--callback generate_key(Db :: #{}) ->
-    {ok, key(), wrapped_key()}.
+-callback generate_key(Db :: #{}, DbOptions :: list()) ->
+    {ok, key(), aegis_config()}.
 
--callback unwrap_key(Db :: #{}, WrappedKey :: wrapped_key()) ->
-    {ok, key(), wrapped_key()}.
+-callback unwrap_key(Db :: #{}, AegisConfig :: aegis_config()) ->
+    {ok, key(), aegis_config()}.
 
 
 -export([
-    generate_key/1,
+    generate_key/2,
     unwrap_key/2
 ]).
 
 
-generate_key(#{} = Db) ->
-    ?AEGIS_KEY_MANAGER:generate_key(Db).
+generate_key(#{} = Db, Options) ->
+    ?AEGIS_KEY_MANAGER:generate_key(Db, Options).
 
 
-unwrap_key(#{} = Db, WrappedKey) ->
-    ?AEGIS_KEY_MANAGER:unwrap_key(Db, WrappedKey).
+unwrap_key(#{} = Db, AegisConfig) ->
+    ?AEGIS_KEY_MANAGER:unwrap_key(Db, AegisConfig).
diff --git a/src/aegis/src/aegis_server.erl b/src/aegis/src/aegis_server.erl
index e5345e744..f0b9b80d3 100644
--- a/src/aegis/src/aegis_server.erl
+++ b/src/aegis/src/aegis_server.erl
@@ -20,13 +20,15 @@
 -include("aegis.hrl").
 
 
+%% aegis_server API
 -export([
     start_link/0,
-    generate_key/1,
+    generate_key/2,
     encrypt/3,
     decrypt/3
 ]).
 
+%% gen_server callbacks
 -export([
     init/1,
     terminate/2,
@@ -36,8 +38,9 @@
     code_change/3
 ]).
 
+%% workers callbacks
 -export([
-    do_generate_key/1,
+    do_generate_key/2,
     do_unwrap_key/1,
     do_encrypt/5,
     do_decrypt/5
@@ -55,9 +58,10 @@ start_link() ->
     gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).
 
 
--spec generate_key(Db :: #{}) -> {ok, binary()} | {error, atom()}.
-generate_key(#{} = Db) ->
-    gen_server:call(?MODULE, {generate_key, Db}).
+-spec generate_key(Db :: #{}, Options :: list()) ->
+        {ok, term()} | {error, term()}.
+generate_key(#{} = Db, Options) ->
+    gen_server:call(?MODULE, {generate_key, Db, Options}).
 
 
 -spec encrypt(Db :: #{}, Key :: binary(), Value :: binary()) -> binary().
@@ -91,7 +95,7 @@ terminate(_Reason, St) ->
         waiters := Waiters
     } = St,
 
-    dict:fold(fun(_WrappedKey, WaitList, _) ->
+    dict:fold(fun(_AegisConfig, WaitList, _) ->
         lists:foreach(fun(#{from := From}) ->
             gen_server:reply(From, {error, decryption_aborted})
         end, WaitList)
@@ -104,8 +108,8 @@ terminate(_Reason, St) ->
     ok.
 
 
-handle_call({generate_key, Db}, From, #{openers := Openers} = St) ->
-    {_Pid, Ref} = erlang:spawn_monitor(?MODULE, do_generate_key, [Db]),
+handle_call({generate_key, Db, Options}, From, #{openers := Openers} = St) ->
+    {_, Ref} = erlang:spawn_monitor(?MODULE, do_generate_key, [Db, Options]),
     Openers1 = dict:store(Ref, From, Openers),
     {noreply, St#{openers := Openers1}, ?TIMEOUT};
 
@@ -125,7 +129,7 @@ handle_cast(_Msg, St) ->
     {noreply, St}.
 
 
-handle_info({'DOWN', Ref, _, _Pid, {ok, DbKey, WrappedKey}}, St) ->
+handle_info({'DOWN', Ref, _, _Pid, {ok, DbKey, AegisConfig}}, St) ->
     #{
         cache := Cache,
         openers := Openers,
@@ -133,15 +137,15 @@ handle_info({'DOWN', Ref, _, _Pid, {ok, DbKey, WrappedKey}}, St) ->
         unwrappers := Unwrappers
     } = St,
 
-    ok = insert(Cache, WrappedKey, DbKey),
+    ok = insert(Cache, AegisConfig, DbKey),
 
     case dict:take(Ref, Openers) of
         {From, Openers1} ->
-            gen_server:reply(From, {ok, WrappedKey}),
+            gen_server:reply(From, {ok, AegisConfig}),
             {noreply, St#{openers := Openers1}, ?TIMEOUT};
         error ->
-            Unwrappers1 = dict:erase(WrappedKey, Unwrappers),
-            {WaitList, Waiters1} = dict:take(WrappedKey, Waiters),
+            Unwrappers1 = dict:erase(AegisConfig, Unwrappers),
+            {WaitList, Waiters1} = dict:take(AegisConfig, Waiters),
             lists:foreach(fun(Waiter) ->
                 #{
                     from := From,
@@ -166,12 +170,12 @@ handle_info({'DOWN', Ref, process, _Pid, {error, Error}}, St) ->
             gen_server:reply(From, {error, Error}),
             {noreply, St#{openers := Openers1}, ?TIMEOUT};
         error ->
-            {ok, WrappedKey} = dict:fold(fun
+            {ok, AegisConfig} = dict:fold(fun
                 (K, V, _) when V == Ref -> {ok, K};
                 (_, _, Acc) -> Acc
             end, not_found, Unwrappers),
-            Unwrappers1 = dict:erase(WrappedKey, Unwrappers),
-            {WaitList, Waiters1} = dict:take(WrappedKey, Waiters),
+            Unwrappers1 = dict:erase(AegisConfig, Unwrappers),
+            {WaitList, Waiters1} = dict:take(AegisConfig, Waiters),
             lists:foreach(fun(#{from := From}) ->
                 gen_server:reply(From, {error, Error})
             end, WaitList),
@@ -189,10 +193,10 @@ code_change(_OldVsn, St, _Extra) ->
 
 %% workers functions
 
-do_generate_key(#{} = Db) ->
+do_generate_key(#{} = Db, Options) ->
     process_flag(sensitive, true),
     try
-        aegis_key_manager:generate_key(Db)
+        aegis_key_manager:generate_key(Db, Options)
     of
         Resp ->
             exit(Resp)
@@ -202,10 +206,10 @@ do_generate_key(#{} = Db) ->
     end.
 
 
-do_unwrap_key(#{aegis := WrappedKey} = Db) ->
+do_unwrap_key(#{aegis := AegisConfig} = Db) ->
     process_flag(sensitive, true),
     try
-        aegis_key_manager:unwrap_key(Db, WrappedKey)
+        aegis_key_manager:unwrap_key(Db, AegisConfig)
     of
         Resp ->
             exit(Resp)
@@ -271,13 +275,15 @@ do_decrypt(From, DbKey, #{uuid := UUID}, Key, Value) ->
 
 %% private functions
 
-maybe_spawn_worker(St, From, Action, #{aegis := WrappedKey} = Db, Key, Value) ->
+maybe_spawn_worker(St, From, Action, Db, Key, Value) ->
     #{
         cache := Cache,
         waiters := Waiters
     } = St,
 
-    case lookup(Cache, WrappedKey) of
+    #{aegis := AegisConfig} = Db,
+
+    case lookup(Cache, AegisConfig) of
         {ok, DbKey} ->
             erlang:spawn(?MODULE, Action, [From, DbKey, Db, Key, Value]),
             St;
@@ -288,37 +294,37 @@ maybe_spawn_worker(St, From, Action, #{aegis := WrappedKey} = Db, Key, Value) ->
                 action => Action,
                 args => [Db, Key, Value]
             },
-            Waiters1 = dict:append(WrappedKey, Waiter, Waiters),
+            Waiters1 = dict:append(AegisConfig, Waiter, Waiters),
             NewSt#{waiters := Waiters1}
      end.
 
 
-maybe_spawn_unwrapper(St, #{aegis := WrappedKey} = Db) ->
+maybe_spawn_unwrapper(St, #{aegis := AegisConfig} = Db) ->
     #{
         unwrappers := Unwrappers
     } = St,
 
-    case dict:is_key(WrappedKey, Unwrappers) of
+    case dict:is_key(AegisConfig, Unwrappers) of
         true ->
             St;
         false ->
             {_Pid, Ref} = erlang:spawn_monitor(?MODULE, do_unwrap_key, [Db]),
-            Unwrappers1 = dict:store(WrappedKey, Ref, Unwrappers),
+            Unwrappers1 = dict:store(AegisConfig, Ref, Unwrappers),
             St#{unwrappers := Unwrappers1}
     end.
 
 
 %% cache functions
 
-insert(Cache, WrappedKey, DbKey) ->
-    Entry = #entry{id = WrappedKey, key = DbKey},
+insert(Cache, AegisConfig, DbKey) ->
+    Entry = #entry{id = AegisConfig, key = DbKey},
     true = ets:insert(Cache, Entry),
     ok.
 
 
-lookup(Cache, WrappedKey) ->
-    case ets:lookup(Cache, WrappedKey) of
-        [#entry{id = WrappedKey, key = DbKey}] ->
+lookup(Cache, AegisConfig) ->
+    case ets:lookup(Cache, AegisConfig) of
+        [#entry{id = AegisConfig, key = DbKey}] ->
             {ok, DbKey};
         [] ->
             {error, not_found}
diff --git a/src/aegis/test/aegis_server_test.erl b/src/aegis/test/aegis_server_test.erl
index 058ca79b2..f0782856a 100644
--- a/src/aegis/test/aegis_server_test.erl
+++ b/src/aegis/test/aegis_server_test.erl
@@ -43,13 +43,18 @@ basic_test_() ->
 
 setup() ->
     Ctx = test_util:start_couch([fabric]),
-    %% isolate aegis_key_cache from actual crypto
-    meck:new([aegis_server, aegis_keywrap], [passthrough]),
-    ok = meck:expect(aegis_keywrap, key_wrap, 2, <<0:320>>),
-    ok = meck:expect(aegis_keywrap, key_unwrap, fun(_, _) ->
+    meck:new([aegis_server, aegis_key_manager], [passthrough]),
+    ok = meck:expect(aegis_key_manager, generate_key, fun(Db, _) ->
+        DbKey = <<0:256>>,
+        #{aegis := AegisConfig} = Db,
+        {ok, DbKey, AegisConfig}
+    end),
+    ok = meck:expect(aegis_key_manager, unwrap_key, fun(Db, _) ->
         %% build a line of the waiters
         timer:sleep(20),
-        <<0:256>>
+        DbKey = <<0:256>>,
+        #{aegis := AegisConfig} = Db,
+        {ok, DbKey, AegisConfig}
     end),
     ok = meck:expect(aegis_server, do_encrypt, fun(From, _, _, _, _) ->
         gen_server:reply(From, ?ENCRYPTED)
@@ -66,13 +71,13 @@ teardown(Ctx) ->
 
 
 test_generate_key() ->
-    {ok, WrappedKey1} = aegis_server:generate_key(?DB),
+    {ok, WrappedKey1} = aegis_server:generate_key(?DB, []),
     ?assertEqual(<<0:320>>, WrappedKey1),
-    ?assertEqual(1, meck:num_calls(aegis_keywrap, key_wrap, 2)).
+    ?assertEqual(1, meck:num_calls(aegis_key_manager, generate_key, 2)).
 
 
 test_encrypt() ->
-    ?assertEqual(0, meck:num_calls(aegis_keywrap, key_unwrap, 2)),
+    ?assertEqual(0, meck:num_calls(aegis_key_manager, unwrap_key, 2)),
     ?assertEqual(0, meck:num_calls(aegis_server, do_encrypt, 5)),
 
     lists:foreach(fun(I) ->
@@ -80,12 +85,12 @@ test_encrypt() ->
         ?assertEqual(?ENCRYPTED, Encrypted)
     end, lists:seq(1, 12)),
 
-    ?assertEqual(1, meck:num_calls(aegis_keywrap, key_unwrap, 2)),
+    ?assertEqual(1, meck:num_calls(aegis_key_manager, unwrap_key, 2)),
     ?assertEqual(12, meck:num_calls(aegis_server, do_encrypt, 5)).
 
 
 test_decrypt() ->
-    ?assertEqual(0, meck:num_calls(aegis_keywrap, key_unwrap, 2)),
+    ?assertEqual(0, meck:num_calls(aegis_key_manager, unwrap_key, 2)),
     ?assertEqual(0, meck:num_calls(aegis_server, do_encrypt, 5)),
 
     lists:foreach(fun(I) ->
@@ -93,16 +98,17 @@ test_decrypt() ->
         ?assertEqual(?VALUE, Decrypted)
     end, lists:seq(1, 12)),
 
-    ?assertEqual(1, meck:num_calls(aegis_keywrap, key_unwrap, 2)),
+    ?assertEqual(1, meck:num_calls(aegis_key_manager, unwrap_key, 2)),
     ?assertEqual(12, meck:num_calls(aegis_server, do_decrypt, 5)).
 
 
 test_multibase() ->
-    ?assertEqual(0, meck:num_calls(aegis_keywrap, key_unwrap, 2)),
+    ?assertEqual(0, meck:num_calls(aegis_key_manager, unwrap_key, 2)),
     ?assertEqual(0, meck:num_calls(aegis_server, do_encrypt, 5)),
+    ?assertEqual(0, meck:num_calls(aegis_server, do_decrypt, 5)),
 
     lists:foreach(fun(I) ->
-        Db = ?DB#{aegis => <<I:320>>},
+        Db = ?DB#{aegis => {<<"wrapped_key">>, <<I:320>>}},
         lists:foreach(fun(J) ->
             Key = <<J:64>>,
             Out = aegis_server:encrypt(Db, Key, ?VALUE),
@@ -112,7 +118,7 @@ test_multibase() ->
         end, lists:seq(1, 10))
     end, lists:seq(1, 12)),
 
-    ?assertEqual(12, meck:num_calls(aegis_keywrap, key_unwrap, 2)),
+    ?assertEqual(12, meck:num_calls(aegis_key_manager, unwrap_key, 2)),
     ?assertEqual(120, meck:num_calls(aegis_server, do_encrypt, 5)),
     ?assertEqual(120, meck:num_calls(aegis_server, do_decrypt, 5)).
 
@@ -123,8 +129,10 @@ error_test_() ->
         foreach,
         fun() ->
             Ctx = setup(),
-            ok = meck:delete(aegis_keywrap, key_unwrap, 2),
-            ok = meck:expect(aegis_keywrap, key_unwrap, 2, fail),
+            ok = meck:delete(aegis_key_manager, unwrap_key, 2),
+            ok = meck:expect(aegis_key_manager, unwrap_key, fun(_, _) ->
+                error(unwrap_failed)
+            end),
             Ctx
         end,
         fun teardown/1,