Make aegis into app and add key cache server

7 files changed, 482 insertions, 22 deletions

diff --git a/rel/reltool.config b/rel/reltool.config
index 1e64a808d..b59c95f55 100644
--- a/rel/reltool.config
+++ b/rel/reltool.config
@@ -27,6 +27,7 @@
         syntax_tools,
         xmerl,
         %% couchdb
+        aegis,
         b64url,
         bear,
         chttpd,
diff --git a/src/aegis/src/aegis.app.src b/src/aegis/src/aegis.app.src
index 51b608df9..e51f42244 100644
--- a/src/aegis/src/aegis.app.src
+++ b/src/aegis/src/aegis.app.src
@@ -14,12 +14,15 @@
  [
   {description, "If it's good enough for Zeus, it's good enough for CouchDB"},
   {vsn, git},
+  {mod, {aegis_app, []}},
+  {registered, [
+    aegis_key_cache
+  ]},
   {applications,
    [kernel,
     stdlib,
     crypto,
     couch_log,
-    base64,
     erlfdb
    ]},
   {env,[]},
diff --git a/src/aegis/src/aegis.erl b/src/aegis/src/aegis.erl
index dc8271f36..ca38a7d5b 100644
--- a/src/aegis/src/aegis.erl
+++ b/src/aegis/src/aegis.erl
@@ -14,8 +14,6 @@
 -include("aegis.hrl").
 -include_lib("fabric/include/fabric2.hrl").
 
-%% TODO - get from key manager
--define(ROOT_KEY, <<1:256>>).
 
 -define(WRAPPED_KEY, {?DB_AEGIS, 1}).
 
@@ -26,32 +24,31 @@
 
     decrypt/2,
     decrypt/3,
+    decrypt/4,
     encrypt/3,
+    encrypt/4,
     wrap_fold_fun/2
 ]).
 
-create(#{} = Db, Options) ->
+create(#{} = Db, _Options) ->
     #{
         tx := Tx,
         db_prefix := DbPrefix
     } = Db,
 
-    % Generate new key
-    DbKey = crypto:strong_rand_bytes(32),
-
-    % protect it with root key
-    WrappedKey = aegis_keywrap:key_wrap(?ROOT_KEY, DbKey),
+    % Fetch unwrapped key
+    WrappedKey = gen_server:call(aegis_key_cache, {get_wrapped_key, Db}),
 
     % And store it
     FDBKey = erlfdb_tuple:pack(?WRAPPED_KEY, DbPrefix),
     ok = erlfdb:set(Tx, FDBKey, WrappedKey),
 
     Db#{
-        aegis => DbKey
+        aegis => WrappedKey
     }.
 
 
-open(#{} = Db, Options) ->
+open(#{} = Db, _Options) ->
     #{
         tx := Tx,
         db_prefix := DbPrefix
@@ -61,11 +58,10 @@ open(#{} = Db, Options) ->
     FDBKey = erlfdb_tuple:pack(?WRAPPED_KEY, DbPrefix),
     WrappedKey = erlfdb:wait(erlfdb:get(Tx, FDBKey)),
 
-    % Unwrap it
-    DbKey = aegis_keywrap:key_unwrap(?ROOT_KEY, WrappedKey),
+    %% maybe ask to rewrap and store if updated?
 
     Db#{
-        aegis => DbKey
+        aegis => WrappedKey
     }.
 
 
@@ -73,11 +69,9 @@ encrypt(#{} = _Db, _Key, <<>>) ->
     <<>>;
 
 encrypt(#{} = Db, Key, Value) when is_binary(Key), is_binary(Value) ->
-    #{
-        uuid := UUID,
-        aegis := DbKey
-    } = Db,
+    gen_server:call(aegis_key_cache, {encrypt, Db, Key, Value}).
 
+encrypt(DbKey, UUID, Key, Value) ->
     EncryptionKey = crypto:strong_rand_bytes(32),
     <<WrappedKey:320>> = aegis_keywrap:key_wrap(DbKey, EncryptionKey),
 
@@ -99,11 +93,9 @@ decrypt(#{} = _Db, _Key, <<>>) ->
     <<>>;
 
 decrypt(#{} = Db, Key, Value) when is_binary(Key), is_binary(Value) ->
-    #{
-        uuid := UUID,
-        aegis := DbKey
-    } = Db,
+    gen_server:call(aegis_key_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
diff --git a/src/aegis/src/aegis_app.erl b/src/aegis/src/aegis_app.erl
new file mode 100644
index 000000000..4a5a11f0c
--- /dev/null
+++ b/src/aegis/src/aegis_app.erl
@@ -0,0 +1,26 @@
+% 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_app).
+
+-behaviour(application).
+
+
+-export([start/2, stop/1]).
+
+
+start(_StartType, _StartArgs) ->
+    aegis_sup:start_link().
+
+
+stop(_State) ->
+    ok.
diff --git a/src/aegis/src/aegis_key_cache.erl b/src/aegis/src/aegis_key_cache.erl
new file mode 100644
index 000000000..9e4ba2fd4
--- /dev/null
+++ b/src/aegis/src/aegis_key_cache.erl
@@ -0,0 +1,280 @@
+% 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_cache).
+
+-behaviour(gen_server).
+
+-vsn(1).
+
+
+-export([
+    start_link/0
+]).
+
+-export([
+    init/1,
+    terminate/2,
+    handle_call/3,
+    handle_cast/2,
+    handle_info/2,
+    code_change/3
+]).
+
+-export([
+    get_wrapped_key/1,
+    unwrap_key/1,
+    do_encrypt/4,
+    do_decrypt/4
+]).
+
+
+-define(ROOT_KEY, <<1:256>>).
+
+-define(INIT_TIMEOUT, 60000).
+
+-define(TIMEOUT, 10000).
+
+
+-record(entry, {id, key}).
+
+
+start_link() ->
+    gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).
+
+
+%% gen_server functions
+
+init([]) ->
+    process_flag(sensitive, true),
+    Cache = ets:new(?MODULE, [set, private, {keypos, #entry.id}]),
+
+    St = #{
+        cache => Cache,
+        clients => dict:new(),
+        waiters => dict:new(),
+        unwrappers => dict:new()
+    },
+    {ok, St, ?INIT_TIMEOUT}.
+
+
+terminate(_Reason, St) ->
+    #{
+        clients := Clients,
+        waiters := Waiters
+    } = St,
+
+    dict:fold(fun(_WrappedKey, WaitList, _) ->
+        lists:foreach(fun(#{from := From}) ->
+            gen_server:reply(From, {error, decryption_failed})
+        end, WaitList)
+    end, ok, Waiters),
+
+    dict:fold(fun(Ref, From, _) ->
+        erlang:demonitor(Ref),
+        gen_server:reply(From, {error, decryption_failed})
+    end, ok, Clients),
+    ok.
+
+
+handle_call({get_wrapped_key, Db}, From, #{clients := Clients} = St) ->
+    {_Pid, Ref} = erlang:spawn_monitor(?MODULE, get_wrapped_key, [Db]),
+    Clients1 = dict:store(Ref, From, Clients),
+    {noreply, St#{clients := Clients1}, ?TIMEOUT};
+
+handle_call({encrypt, Db, Key, Value}, From, St) ->
+    NewSt = maybe_spawn_worker(St, From, do_encrypt, Db, Key, Value),
+    {noreply, NewSt, ?TIMEOUT};
+
+handle_call({decrypt, Db, Key, Value}, From, St) ->
+    NewSt = maybe_spawn_worker(St, From, do_decrypt, Db, Key, Value),
+    {noreply, NewSt, ?TIMEOUT};
+
+handle_call(_Msg, _From, St) ->
+    {noreply, St}.
+
+
+handle_cast(_Msg, St) ->
+    {noreply, St}.
+
+
+handle_info({'DOWN', Ref, _, _Pid, {key, {ok, DbKey, WrappedKey}}}, St) ->
+    #{
+        cache := Cache,
+        clients := Clients,
+        waiters := Waiters,
+        unwrappers := Unwrappers
+    } = St,
+
+    NewSt1 = case dict:take(WrappedKey, Unwrappers) of
+        {Ref, Unwrappers1} ->
+            ok = insert(Cache, WrappedKey, DbKey),
+            St#{unwrappers := Unwrappers1};
+        error ->
+            %% FIXME! it might be new wrapped key != old wrapped key
+            %% fold here to search for it based on ref
+            St
+    end,
+
+    NewSt2 = case dict:take(WrappedKey, Waiters) of
+        {WaitList, Waiters1} ->
+            Clients1 = lists:foldl(fun(Waiter, Acc) ->
+                #{
+                    from := From,
+                    action := Action,
+                    args := Args
+                } = Waiter,
+
+                {_Pid1, Ref1} = erlang:spawn_monitor(
+                    ?MODULE, Action, [DbKey | Args]),
+
+                dict:store(Ref1, From, Acc)
+            end, Clients, WaitList),
+
+            NewSt1#{clients := Clients1, waiters := Waiters1};
+        error ->
+            NewSt1
+    end,
+
+    NewSt3 = maybe_reply(NewSt2, Ref, WrappedKey),
+    {noreply, NewSt3, ?TIMEOUT};
+
+handle_info({'DOWN', Ref, process, _Pid, Resp}, St) ->
+    NewSt = maybe_reply(St, Ref, Resp),
+    {noreply, NewSt, ?TIMEOUT};
+
+handle_info(_Msg, St) ->
+    {noreply, St}.
+
+
+code_change(_OldVsn, St, _Extra) ->
+    {ok, St}.
+
+
+%% workers functions
+
+maybe_spawn_worker(St, From, Action, #{aegis := WrappedKey} = Db, Key, Value) ->
+    #{
+        cache := Cache,
+        clients := Clients,
+        waiters := Waiters,
+        unwrappers := Unwrappers
+    } = St,
+
+    case lookup(Cache, WrappedKey) of
+        {ok, DbKey} ->
+            {_Pid, Ref} = erlang:spawn_monitor(
+                ?MODULE, Action, [DbKey, Db, Key, Value]),
+            Clients1 = dict:store(Ref, From, Clients),
+            St#{clients := Clients1};
+        {error, not_found} ->
+            NewSt = case dict:is_key(WrappedKey, Unwrappers) of
+                true ->
+                    St;
+                false ->
+                    {_Pid, Ref} = erlang:spawn_monitor(
+                        ?MODULE, unwrap_key, [Db]),
+                    Unwrappers1 = dict:store(WrappedKey, Ref, Unwrappers),
+                    St#{unwrappers := Unwrappers1}
+            end,
+            Waiter = #{
+                from => From,
+                action => Action,
+                args => [Db, Key, Value]
+            },
+            Waiters1 = dict:append(WrappedKey, Waiter, Waiters),
+            NewSt#{waiters := Waiters1}
+     end.
+
+
+maybe_reply(#{clients := Clients} = St, Ref, Resp) ->
+    case dict:take(Ref, Clients) of
+        {From, Clients1} ->
+            gen_server:reply(From, Resp),
+            St#{clients := Clients1};
+        error ->
+            St
+    end.
+
+
+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}
+    of
+        Resp ->
+            exit({key, Resp})
+    catch
+        _:Error ->
+            exit({error, Error})
+    end.
+
+
+unwrap_key(#{aegis := WrappedKey} = _Db) ->
+    process_flag(sensitive, true),
+    try
+        %% this could be atom fail, throw error is so !!
+        DbKey = aegis_keywrap:key_unwrap(?ROOT_KEY, WrappedKey),
+        {ok, DbKey, WrappedKey}
+    of
+        Resp ->
+            exit({key, Resp})
+    catch
+        _:Error ->
+            %% FIXME! add tag key and WrappedKey so we can respond to Waiters
+            exit({error, Error})
+    end.
+
+
+do_encrypt(DbKey, #{uuid := UUID}, Key, Value) ->
+    process_flag(sensitive, true),
+    try
+        aegis:encrypt(DbKey, UUID, Key, Value)
+    of
+        Resp ->
+            exit(Resp)
+    catch
+        _:Error ->
+            exit({error, Error})
+    end.
+
+
+do_decrypt(DbKey, #{uuid := UUID}, Key, Value) ->
+    process_flag(sensitive, true),
+    try
+        aegis:decrypt(DbKey, UUID, Key, Value)
+    of
+        Resp ->
+            exit(Resp)
+    catch
+        _:Error ->
+            exit({error, Error})
+    end.
+
+
+%% cache functions
+
+insert(Cache, WrappedKey, DbKey) ->
+    Entry = #entry{id = WrappedKey, key = DbKey},
+    true = ets:insert(Cache, Entry),
+    ok.
+
+
+lookup(Cache, WrappedKey) ->
+    case ets:lookup(Cache, WrappedKey) of
+        [#entry{id = WrappedKey, key = DbKey}] ->
+            {ok, DbKey};
+        [] ->
+            {error, not_found}
+    end.
diff --git a/src/aegis/src/aegis_sup.erl b/src/aegis/src/aegis_sup.erl
new file mode 100644
index 000000000..65f844c4b
--- /dev/null
+++ b/src/aegis/src/aegis_sup.erl
@@ -0,0 +1,46 @@
+% 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_sup).
+
+-behaviour(supervisor).
+
+-vsn(1).
+
+
+-export([
+    start_link/0
+]).
+
+-export([
+    init/1
+]).
+
+
+start_link() ->
+    supervisor:start_link({local, ?MODULE}, ?MODULE, []).
+
+
+init([]) ->
+    Flags = #{
+        strategy => one_for_one,
+        intensity => 5,
+        period => 10
+    },
+    Children = [
+        #{
+            id => aegis_key_cache,
+            start => {aegis_key_cache, start_link, []},
+            shutdown => 5000
+        }
+    ],
+    {ok, {Flags, Children}}.
diff --git a/src/aegis/test/aegis_key_cache_test.erl b/src/aegis/test/aegis_key_cache_test.erl
new file mode 100644
index 000000000..2e6680e6f
--- /dev/null
+++ b/src/aegis/test/aegis_key_cache_test.erl
@@ -0,0 +1,112 @@
+% 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_cache_test).
+
+-include_lib("eunit/include/eunit.hrl").
+-include_lib("couch/include/couch_eunit.hrl").
+
+-define(SERVER, aegis_key_cache).
+-define(DB, #{aegis => <<0:320>>, uuid => <<0:64>>}).
+-define(VALUE, <<0:8192>>).
+-define(ENCRYPTED, <<1:8, 0:320, 0:4096>>).
+
+
+
+basic_test_() ->
+    {
+        foreach,
+        fun setup/0,
+        fun teardown/1,
+        [
+            {"cache unwrapped key on get_wrapped_key",
+            fun test_get_wrapped_key/0},
+            {"cache unwrapped key on encrypt",
+            fun test_encrypt/0},
+            {"cache unwrapped key on decrypt",
+            fun test_decrypt/0},
+            {"cache unwrapped key per database",
+            fun test_multibase/0}
+        ]
+    }.
+
+
+setup() ->
+    Ctx = test_util:start_couch([fabric]),
+    %% isolate aegis_key_cache from actual crypto
+    meck:new([aegis, 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(50),
+        <<0:256>>
+    end),
+    ok = meck:expect(aegis, encrypt, 4, ?ENCRYPTED),
+    ok = meck:expect(aegis, decrypt, 4, ?VALUE),
+    Ctx.
+
+
+teardown(Ctx) ->
+    meck:unload(),
+    test_util:stop_couch(Ctx).
+
+
+test_get_wrapped_key() ->
+    WrappedKey1 = gen_server:call(?SERVER, {get_wrapped_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)),
+
+    lists:foreach(fun(I) ->
+        Encrypted = gen_server:call(?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)).
+
+
+test_decrypt() ->
+    ?assertEqual(0, meck:num_calls(aegis_keywrap, key_unwrap, 2)),
+    ?assertEqual(0, meck:num_calls(aegis, encrypt, 4)),
+
+    lists:foreach(fun(I) ->
+        Decrypted = gen_server:call(
+            ?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)).
+
+test_multibase() ->
+    ?assertEqual(0, meck:num_calls(aegis_keywrap, key_unwrap, 2)),
+    ?assertEqual(0, meck:num_calls(aegis, encrypt, 4)),
+
+    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}),
+            ?assertEqual(?ENCRYPTED, Out),
+            In = gen_server:call(?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)).