Address review comments

- Rename aegis_key_cache to aegis_server
- Move crypto into workers
- Add aegis_server public API
- Define aegis_key_manager behaviour
- Fix error messages
- Remove part with incomplete functionality
- Remove defensive programming

8 files changed, 196 insertions, 182 deletions

diff --git a/src/aegis/rebar.config.script b/src/aegis/rebar.config.script
index 27752fe81..a2337a3f3 100644
--- a/src/aegis/rebar.config.script
+++ b/src/aegis/rebar.config.script
@@ -27,7 +27,7 @@ case lists:keyfind(aegis_key_manager, 1, CouchConfig) of
             {erl_opts, Opts} -> Opts;
             false -> []
         end,
-        AegisOpts = {d, 'AEGIS_KEY_MANAGER', list_to_existing_atom(Module)},
+        AegisOpts = {d, 'AEGIS_KEY_MANAGER', list_to_atom(Module)},
         NewOpts = [AegisOpts | CurrentOpts],
         lists:keystore(erl_opts, 1, CONFIG, {erl_opts, NewOpts});
     _ ->
diff --git a/src/aegis/src/aegis.app.src b/src/aegis/src/aegis.app.src
index e51f42244..deb152674 100644
--- a/src/aegis/src/aegis.app.src
+++ b/src/aegis/src/aegis.app.src
@@ -16,7 +16,7 @@
   {vsn, git},
   {mod, {aegis_app, []}},
   {registered, [
-    aegis_key_cache
+    aegis_server
   ]},
   {applications,
    [kernel,
diff --git a/src/aegis/src/aegis.erl b/src/aegis/src/aegis.erl
index bdf980e04..0315c7b88 100644
--- a/src/aegis/src/aegis.erl
+++ b/src/aegis/src/aegis.erl
@@ -11,14 +11,11 @@
 % the License.
 
 -module(aegis).
--include("aegis.hrl").
 -include_lib("fabric/include/fabric2.hrl").
 
 
 -define(WRAPPED_KEY, {?DB_AEGIS, 1}).
 
--define(CACHE, aegis_key_cache).
-
 
 -export([
     create/2,
@@ -26,9 +23,7 @@
 
     decrypt/2,
     decrypt/3,
-    decrypt/4,
     encrypt/3,
-    encrypt/4,
     wrap_fold_fun/2
 ]).
 
@@ -38,10 +33,8 @@ create(#{} = Db, _Options) ->
         db_prefix := DbPrefix
     } = Db,
 
-    % Fetch unwrapped key
-    WrappedKey = gen_server:call(?CACHE, {get_wrapped_key, Db}),
+    {ok, WrappedKey} = aegis_server:generate_key(Db),
 
-    % And store it
     FDBKey = erlfdb_tuple:pack(?WRAPPED_KEY, DbPrefix),
     ok = erlfdb:set(Tx, FDBKey, WrappedKey),
 
@@ -50,7 +43,7 @@ create(#{} = Db, _Options) ->
     }.
 
 
-open(#{} = Db, _Options) ->
+open(#{} = Db, Options) ->
     #{
         tx := Tx,
         db_prefix := DbPrefix
@@ -60,35 +53,16 @@ open(#{} = Db, _Options) ->
     FDBKey = erlfdb_tuple:pack(?WRAPPED_KEY, DbPrefix),
     WrappedKey = erlfdb:wait(erlfdb:get(Tx, FDBKey)),
 
-    Db1 = Db#{aegis => WrappedKey},
-
-    case gen_server:call(?CACHE, {maybe_rewrap_key, Db1}) of
-        WrappedKey ->
-            Db1;
-        NewWrappedKey ->
-            FDBKey = erlfdb_tuple:pack(?WRAPPED_KEY, DbPrefix),
-            ok = erlfdb:set(Tx, FDBKey, NewWrappedKey),
-            Db1#{aegis => NewWrappedKey}
-    end.
+    Db#{
+        aegis => WrappedKey
+    }.
 
 
 encrypt(#{} = _Db, _Key, <<>>) ->
     <<>>;
 
 encrypt(#{} = Db, Key, Value) when is_binary(Key), is_binary(Value) ->
-    gen_server:call(?CACHE, {encrypt, Db, Key, Value}).
-
-encrypt(DbKey, UUID, Key, Value) ->
-    EncryptionKey = crypto:strong_rand_bytes(32),
-    <<WrappedKey:320>> = aegis_keywrap:key_wrap(DbKey, EncryptionKey),
-
-    {CipherText, <<CipherTag:128>>} =
-        ?aes_gcm_encrypt(
-           EncryptionKey,
-           <<0:96>>,
-           <<UUID/binary, 0:8, Key/binary>>,
-           Value),
-    <<1:8, WrappedKey:320, CipherTag:128, CipherText/binary>>.
+    aegis_server:encrypt(Db, Key, Value).
 
 
 decrypt(#{} = Db, Rows) when is_list(Rows) ->
@@ -100,29 +74,7 @@ decrypt(#{} = _Db, _Key, <<>>) ->
     <<>>;
 
 decrypt(#{} = Db, Key, Value) when is_binary(Key), is_binary(Value) ->
-    gen_server:call(?CACHE, {decrypt, Db, Key, Value}).
-
-decrypt(DbKey, UUID, Key, Value) ->
-    case Value of
-        <<1:8, WrappedKey:320, CipherTag:128, CipherText/binary>> ->
-            case aegis_keywrap:key_unwrap(DbKey, <<WrappedKey:320>>) of
-                fail ->
-                    erlang:error(decryption_failed);
-                DecryptionKey ->
-                    Decrypted =
-                    ?aes_gcm_decrypt(
-                        DecryptionKey,
-                        <<0:96>>,
-                        <<UUID/binary, 0:8, Key/binary>>,
-                        CipherText,
-                        <<CipherTag:128>>),
-                    if Decrypted /= error -> Decrypted; true ->
-                        erlang:error(decryption_failed)
-                    end
-            end;
-        _ ->
-            erlang:error(not_ciphertext)
-    end.
+    aegis_server:decrypt(Db, Key, Value).
 
 
 wrap_fold_fun(Db, Fun) when is_function(Fun, 2) ->
diff --git a/src/aegis/src/aegis_file_key_manager.erl b/src/aegis/src/aegis_file_key_manager.erl
new file mode 100644
index 000000000..f520bd497
--- /dev/null
+++ b/src/aegis/src/aegis_file_key_manager.erl
@@ -0,0 +1,40 @@
+% 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_file_key_manager).
+
+
+-behaviour(aegis_key_manager).
+
+
+-export([
+    generate_key/1,
+    unwrap_key/2
+]).
+
+
+-define(ROOT_KEY, <<1:256>>).
+
+
+generate_key(#{} = _Db) ->
+    DbKey = crypto:strong_rand_bytes(32),
+    WrappedKey = aegis_keywrap:key_wrap(?ROOT_KEY, DbKey),
+    {ok, DbKey, WrappedKey}.
+
+
+unwrap_key(#{} = _Db, WrappedKey) ->
+    case aegis_keywrap:key_unwrap(?ROOT_KEY, WrappedKey) of
+        fail ->
+            error(unwrap_failed);
+        DbKey ->
+            {ok, DbKey, WrappedKey}
+    end.
diff --git a/src/aegis/src/aegis_key_manager.erl b/src/aegis/src/aegis_key_manager.erl
index e41cfafff..a16c51690 100644
--- a/src/aegis/src/aegis_key_manager.erl
+++ b/src/aegis/src/aegis_key_manager.erl
@@ -13,25 +13,25 @@
 -module(aegis_key_manager).
 
 
--export([
-    key_wrap/1,
-    key_unwrap/1
-]).
+-type key() :: binary().
+-type wrapped_key() :: binary().
+
+-callback generate_key(Db :: #{}) ->
+    {ok, key(), wrapped_key()}.
 
+-callback unwrap_key(Db :: #{}, WrappedKey :: wrapped_key()) ->
+    {ok, key(), wrapped_key()}.
 
--define(ROOT_KEY, <<1:256>>).
+
+-export([
+    generate_key/1,
+    unwrap_key/2
+]).
 
 
-key_wrap(#{} = _Db) ->
-    DbKey = crypto:strong_rand_bytes(32),
-    WrappedKey = aegis_keywrap:key_wrap(?ROOT_KEY, DbKey),
-    {ok, DbKey, WrappedKey}.
+generate_key(#{} = Db) ->
+    ?AEGIS_KEY_MANAGER:generate_key(Db).
 
 
-key_unwrap(#{aegis := WrappedKey} = _Db) ->
-    case aegis_keywrap:key_unwrap(?ROOT_KEY, WrappedKey) of
-        fail ->
-            error(decryption_failed);
-        DbKey ->
-            {ok, DbKey, WrappedKey}
-    end.
+unwrap_key(#{} = Db, WrappedKey) ->
+    ?AEGIS_KEY_MANAGER:unwrap_key(Db, WrappedKey).
diff --git a/src/aegis/src/aegis_key_cache.erl b/src/aegis/src/aegis_server.erl
index 67ac996a8..e5345e744 100644
--- a/src/aegis/src/aegis_key_cache.erl
+++ b/src/aegis/src/aegis_server.erl
@@ -10,15 +10,21 @@
 % License for the specific language governing permissions and limitations under
 % the License.
 
--module(aegis_key_cache).
+-module(aegis_server).
 
 -behaviour(gen_server).
 
 -vsn(1).
 
 
+-include("aegis.hrl").
+
+
 -export([
-    start_link/0
+    start_link/0,
+    generate_key/1,
+    encrypt/3,
+    decrypt/3
 ]).
 
 -export([
@@ -31,8 +37,8 @@
 ]).
 
 -export([
-    get_wrapped_key/1,
-    unwrap_key/1,
+    do_generate_key/1,
+    do_unwrap_key/1,
     do_encrypt/5,
     do_decrypt/5
 ]).
@@ -49,6 +55,21 @@ 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 encrypt(Db :: #{}, Key :: binary(), Value :: binary()) -> binary().
+encrypt(#{} = Db, Key, Value)  when is_binary(Key), is_binary(Value) ->
+    gen_server:call(?MODULE, {encrypt, Db, Key, Value}).
+
+
+-spec decrypt(Db :: #{}, Key :: binary(), Value :: binary()) -> binary().
+decrypt(#{} = Db, Key, Value) when is_binary(Key), is_binary(Value) ->
+    gen_server:call(?MODULE, {decrypt, Db, Key, Value}).
+
+
 %% gen_server functions
 
 init([]) ->
@@ -72,34 +93,22 @@ terminate(_Reason, St) ->
 
     dict:fold(fun(_WrappedKey, WaitList, _) ->
         lists:foreach(fun(#{from := From}) ->
-            gen_server:reply(From, {error, decryption_failed})
+            gen_server:reply(From, {error, decryption_aborted})
         end, WaitList)
     end, ok, Waiters),
 
     dict:fold(fun(Ref, From, _) ->
         erlang:demonitor(Ref),
-        gen_server:reply(From, {error, decryption_failed})
+        gen_server:reply(From, {error, decryption_aborted})
     end, ok, Openers),
     ok.
 
 
-handle_call({get_wrapped_key, Db}, From, #{openers := Openers} = St) ->
-    {_Pid, Ref} = erlang:spawn_monitor(?MODULE, get_wrapped_key, [Db]),
+handle_call({generate_key, Db}, From, #{openers := Openers} = St) ->
+    {_Pid, Ref} = erlang:spawn_monitor(?MODULE, do_generate_key, [Db]),
     Openers1 = dict:store(Ref, From, Openers),
     {noreply, St#{openers := Openers1}, ?TIMEOUT};
 
-handle_call({maybe_rewrap_key, #{aegis := WrappedKey} = Db}, From, St) ->
-    #{
-        openers := Openers,
-        unwrappers := Unwrappers
-    } = St,
-
-    {_Pid, Ref} = erlang:spawn_monitor(?MODULE, unwrap_key, [Db]),
-
-    Openers1 = dict:store(Ref, From, Openers),
-    Unwrappers1 = dict:store(WrappedKey, Ref, Unwrappers),
-    {noreply, St#{openers := Openers1, 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};
@@ -116,7 +125,7 @@ handle_cast(_Msg, St) ->
     {noreply, St}.
 
 
-handle_info({'DOWN', Ref, _, _Pid, {key, {ok, DbKey, WrappedKey}}}, St) ->
+handle_info({'DOWN', Ref, _, _Pid, {ok, DbKey, WrappedKey}}, St) ->
     #{
         cache := Cache,
         openers := Openers,
@@ -124,24 +133,15 @@ handle_info({'DOWN', Ref, _, _Pid, {key, {ok, DbKey, WrappedKey}}}, St) ->
         unwrappers := Unwrappers
     } = St,
 
-    IsOpener = dict:is_key(Ref, Openers),
+    ok = insert(Cache, WrappedKey, DbKey),
 
-    NewSt1 = case dict:take(WrappedKey, Unwrappers) of
-        {Ref, Unwrappers1} ->
-            ok = insert(Cache, WrappedKey, DbKey),
-            St#{unwrappers := Unwrappers1};
-        error when IsOpener ->
-            ok = insert(Cache, WrappedKey, DbKey),
-            St;
+    case dict:take(Ref, Openers) of
+        {From, Openers1} ->
+            gen_server:reply(From, {ok, WrappedKey}),
+            {noreply, St#{openers := Openers1}, ?TIMEOUT};
         error ->
-            %% FIXME! it might be new wrapped key != old wrapped key
-            %% fold over Unwrappers here to find waiters of old key
-            %% by Ref. also need way to store new wrapped key in fdb
-            St
-    end,
-
-    NewSt2 = case dict:take(WrappedKey, Waiters) of
-        {WaitList, Waiters1} ->
+            Unwrappers1 = dict:erase(WrappedKey, Unwrappers),
+            {WaitList, Waiters1} = dict:take(WrappedKey, Waiters),
             lists:foreach(fun(Waiter) ->
                 #{
                     from := From,
@@ -150,17 +150,34 @@ handle_info({'DOWN', Ref, _, _Pid, {key, {ok, DbKey, WrappedKey}}}, St) ->
                 } = Waiter,
                 erlang:spawn(?MODULE, Action, [From, DbKey | Args])
             end, WaitList),
-            NewSt1#{waiters := Waiters1};
-        error ->
-            NewSt1
-    end,
+            NewSt = St#{waiters := Waiters1, unwrappers := Unwrappers1},
+            {noreply, NewSt, ?TIMEOUT}
+    end;
 
-    NewSt3 = maybe_reply(NewSt2, Ref, WrappedKey),
-    {noreply, NewSt3, ?TIMEOUT};
+handle_info({'DOWN', Ref, process, _Pid, {error, Error}}, St) ->
+    #{
+        openers := Openers,
+        waiters := Waiters,
+        unwrappers := Unwrappers
+    } = St,
 
-handle_info({'DOWN', Ref, process, _Pid, Resp}, St) ->
-    NewSt = maybe_reply(St, Ref, Resp),
-    {noreply, NewSt, ?TIMEOUT};
+    case dict:take(Ref, Openers) of
+        {From, Openers1} ->
+            gen_server:reply(From, {error, Error}),
+            {noreply, St#{openers := Openers1}, ?TIMEOUT};
+        error ->
+            {ok, WrappedKey} = 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),
+            lists:foreach(fun(#{from := From}) ->
+                gen_server:reply(From, {error, Error})
+            end, WaitList),
+            NewSt = St#{waiters := Waiters1, unwrappers := Unwrappers1},
+            {noreply, NewSt, ?TIMEOUT}
+    end;
 
 handle_info(_Msg, St) ->
     {noreply, St}.
@@ -172,36 +189,45 @@ code_change(_OldVsn, St, _Extra) ->
 
 %% workers functions
 
-get_wrapped_key(#{} = Db) ->
+do_generate_key(#{} = Db) ->
     process_flag(sensitive, true),
     try
-        ?AEGIS_KEY_MANAGER:key_wrap(Db)
+        aegis_key_manager:generate_key(Db)
     of
         Resp ->
-            exit({key, Resp})
+            exit(Resp)
     catch
         _:Error ->
             exit({error, Error})
     end.
 
 
-unwrap_key(#{aegis := WrappedKey} = Db) ->
+do_unwrap_key(#{aegis := WrappedKey} = Db) ->
     process_flag(sensitive, true),
     try
-        ?AEGIS_KEY_MANAGER:key_unwrap(Db)
+        aegis_key_manager:unwrap_key(Db, WrappedKey)
     of
         Resp ->
-            exit({key, Resp})
+            exit(Resp)
     catch
         _:Error ->
-            exit({key, {error, WrappedKey, Error}})
+            exit({error, Error})
     end.
 
 
 do_encrypt(From, DbKey, #{uuid := UUID}, Key, Value) ->
     process_flag(sensitive, true),
     try
-        aegis:encrypt(DbKey, UUID, Key, Value)
+        EncryptionKey = crypto:strong_rand_bytes(32),
+        <<WrappedKey:320>> = aegis_keywrap:key_wrap(DbKey, EncryptionKey),
+
+        {CipherText, <<CipherTag:128>>} =
+            ?aes_gcm_encrypt(
+               EncryptionKey,
+               <<0:96>>,
+               <<UUID/binary, 0:8, Key/binary>>,
+               Value),
+        <<1:8, WrappedKey:320, CipherTag:128, CipherText/binary>>
     of
         Resp ->
             gen_server:reply(From, Resp)
@@ -214,7 +240,26 @@ do_encrypt(From, DbKey, #{uuid := UUID}, Key, Value) ->
 do_decrypt(From, DbKey, #{uuid := UUID}, Key, Value) ->
     process_flag(sensitive, true),
     try
-        aegis:decrypt(DbKey, UUID, Key, Value)
+        case Value of
+            <<1:8, WrappedKey:320, CipherTag:128, CipherText/binary>> ->
+                case aegis_keywrap:key_unwrap(DbKey, <<WrappedKey:320>>) of
+                    fail ->
+                        erlang:error(decryption_failed);
+                    DecryptionKey ->
+                        Decrypted =
+                        ?aes_gcm_decrypt(
+                            DecryptionKey,
+                            <<0:96>>,
+                            <<UUID/binary, 0:8, Key/binary>>,
+                            CipherText,
+                            <<CipherTag:128>>),
+                        if Decrypted /= error -> Decrypted; true ->
+                            erlang:error(decryption_failed)
+                        end
+                end;
+            _ ->
+                erlang:error(not_ciphertext)
+        end
     of
         Resp ->
             gen_server:reply(From, Resp)
@@ -257,38 +302,12 @@ maybe_spawn_unwrapper(St, #{aegis := WrappedKey} = Db) ->
         true ->
             St;
         false ->
-            {_Pid, Ref} = erlang:spawn_monitor(?MODULE, unwrap_key, [Db]),
+            {_Pid, Ref} = erlang:spawn_monitor(?MODULE, do_unwrap_key, [Db]),
             Unwrappers1 = dict:store(WrappedKey, Ref, Unwrappers),
             St#{unwrappers := Unwrappers1}
     end.
 
 
-maybe_reply(St, Ref, {key, {error, WrappedKey, Error}}) ->
-    #{
-        waiters := Waiters
-    } = St,
-
-    Reply = {error, Error},
-
-    NewSt = case dict:take(WrappedKey, Waiters) of
-        {WaitList, Waiters1} ->
-            [ gen_server:reply(From, Reply) || #{from := From} <- WaitList ],
-            St#{waiters := Waiters1};
-        error ->
-            St
-    end,
-    maybe_reply(NewSt, Ref, Reply);
-
-maybe_reply(#{openers := Openers} = St, Ref, Resp) ->
-    case dict:take(Ref, Openers) of
-        {From, Openers1} ->
-            gen_server:reply(From, Resp),
-            St#{openers := Openers1};
-        error ->
-            St
-    end.
-
-
 %% cache functions
 
 insert(Cache, WrappedKey, DbKey) ->
diff --git a/src/aegis/src/aegis_sup.erl b/src/aegis/src/aegis_sup.erl
index 65f844c4b..6d3ee83d8 100644
--- a/src/aegis/src/aegis_sup.erl
+++ b/src/aegis/src/aegis_sup.erl
@@ -38,8 +38,8 @@ init([]) ->
     },
     Children = [
         #{
-            id => aegis_key_cache,
-            start => {aegis_key_cache, start_link, []},
+            id => aegis_server,
+            start => {aegis_server, start_link, []},
             shutdown => 5000
         }
     ],
diff --git a/src/aegis/test/aegis_key_cache_test.erl b/src/aegis/test/aegis_server_test.erl
index f9b189412..058ca79b2 100644
--- a/src/aegis/test/aegis_key_cache_test.erl
+++ b/src/aegis/test/aegis_server_test.erl
@@ -10,12 +10,12 @@
 % License for the specific language governing permissions and limitations under
 % the License.
 
--module(aegis_key_cache_test).
+-module(aegis_server_test).
 
 -include_lib("eunit/include/eunit.hrl").
 -include_lib("couch/include/couch_eunit.hrl").
 
--define(SERVER, aegis_key_cache).
+-define(SERVER, aegis_server).
 -define(DB, #{aegis => <<0:320>>, uuid => <<0:64>>}).
 -define(VALUE, <<0:8192>>).
 -define(ENCRYPTED, <<1:8, 0:320, 0:4096>>).
@@ -29,8 +29,8 @@ basic_test_() ->
         fun setup/0,
         fun teardown/1,
         [
-            {"cache unwrapped key on get_wrapped_key",
-            {timeout, ?TIMEOUT, fun test_get_wrapped_key/0}},
+            {"cache unwrapped key on generate_key",
+            {timeout, ?TIMEOUT, fun test_generate_key/0}},
             {"cache unwrapped key on encrypt",
             {timeout, ?TIMEOUT, fun test_encrypt/0}},
             {"cache unwrapped key on decrypt",
@@ -44,15 +44,19 @@ basic_test_() ->
 setup() ->
     Ctx = test_util:start_couch([fabric]),
     %% isolate aegis_key_cache from actual crypto
-    meck:new([aegis, aegis_keywrap], [passthrough]),
+    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(_, _) ->
         %% build a line of the waiters
         timer:sleep(20),
         <<0:256>>
     end),
-    ok = meck:expect(aegis, encrypt, 4, ?ENCRYPTED),
-    ok = meck:expect(aegis, decrypt, 4, ?VALUE),
+    ok = meck:expect(aegis_server, do_encrypt, fun(From, _, _, _, _) ->
+        gen_server:reply(From, ?ENCRYPTED)
+    end),
+    ok = meck:expect(aegis_server, do_decrypt, fun(From, _, _, _, _) ->
+        gen_server:reply(From, ?VALUE)
+    end),
     Ctx.
 
 
@@ -61,57 +65,56 @@ teardown(Ctx) ->
     test_util:stop_couch(Ctx).
 
 
-test_get_wrapped_key() ->
-    WrappedKey1 = gen_server:call(?SERVER, {get_wrapped_key, ?DB}),
+test_generate_key() ->
+    {ok, WrappedKey1} = aegis_server:generate_key(?DB),
     ?assertEqual(<<0:320>>, WrappedKey1),
     ?assertEqual(1, meck:num_calls(aegis_keywrap, key_wrap, 2)).
 
 
 test_encrypt() ->
     ?assertEqual(0, meck:num_calls(aegis_keywrap, key_unwrap, 2)),
-    ?assertEqual(0, meck:num_calls(aegis, encrypt, 4)),
+    ?assertEqual(0, meck:num_calls(aegis_server, do_encrypt, 5)),
 
     lists:foreach(fun(I) ->
-        Encrypted = gen_server:call(?SERVER, {encrypt, ?DB, <<I:64>>, ?VALUE}),
+        Encrypted = aegis_server:encrypt(?DB, <<I:64>>, ?VALUE),
         ?assertEqual(?ENCRYPTED, Encrypted)
     end, lists:seq(1, 12)),
 
     ?assertEqual(1, meck:num_calls(aegis_keywrap, key_unwrap, 2)),
-    ?assertEqual(12, meck:num_calls(aegis, encrypt, 4)).
+    ?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, encrypt, 4)),
+    ?assertEqual(0, meck:num_calls(aegis_server, do_encrypt, 5)),
 
     lists:foreach(fun(I) ->
-        Decrypted = gen_server:call(
-            ?SERVER, {decrypt, ?DB, <<I:64>>, ?ENCRYPTED}),
+        Decrypted = aegis_server:decrypt(?DB, <<I:64>>, ?ENCRYPTED),
         ?assertEqual(?VALUE, Decrypted)
     end, lists:seq(1, 12)),
 
     ?assertEqual(1, meck:num_calls(aegis_keywrap, key_unwrap, 2)),
-    ?assertEqual(12, meck:num_calls(aegis, decrypt, 4)).
+    ?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, encrypt, 4)),
+    ?assertEqual(0, meck:num_calls(aegis_server, do_encrypt, 5)),
 
     lists:foreach(fun(I) ->
         Db = ?DB#{aegis => <<I:320>>},
         lists:foreach(fun(J) ->
             Key = <<J:64>>,
-            Out = gen_server:call(?SERVER, {encrypt, Db, Key, ?VALUE}),
+            Out = aegis_server:encrypt(Db, Key, ?VALUE),
             ?assertEqual(?ENCRYPTED, Out),
-            In = gen_server:call(?SERVER, {decrypt, Db, Key, Out}),
+            In = aegis_server:decrypt(Db, Key, Out),
             ?assertEqual(?VALUE, In)
         end, lists:seq(1, 10))
     end, lists:seq(1, 12)),
 
     ?assertEqual(12, meck:num_calls(aegis_keywrap, key_unwrap, 2)),
-    ?assertEqual(120, meck:num_calls(aegis, encrypt, 4)),
-    ?assertEqual(120, meck:num_calls(aegis, decrypt, 4)).
+    ?assertEqual(120, meck:num_calls(aegis_server, do_encrypt, 5)),
+    ?assertEqual(120, meck:num_calls(aegis_server, do_decrypt, 5)).
 
 
 
@@ -135,10 +138,10 @@ error_test_() ->
 
 
 test_encrypt_error() ->
-    Reply = gen_server:call(?SERVER, {encrypt, ?DB, <<1:64>>, ?VALUE}),
-    ?assertEqual({error, decryption_failed}, Reply).
+    Reply = aegis_server:encrypt(?DB, <<1:64>>, ?VALUE),
+    ?assertEqual({error, unwrap_failed}, Reply).
 
 
 test_decrypt_error() ->
-    Reply = gen_server:call(?SERVER, {decrypt, ?DB, <<1:64>>, ?VALUE}),
-    ?assertEqual({error, decryption_failed}, Reply).
+    Reply = aegis_server:decrypt(?DB, <<1:64>>, ?ENCRYPTED),
+    ?assertEqual({error, unwrap_failed}, Reply).