Unwrap KEK outside the main loopprototype/fdb-encryption

2 files changed, 355 insertions, 35 deletions

diff --git a/src/fabric/src/fabric2_encryption.erl b/src/fabric/src/fabric2_encryption.erl
index 86b06a37e..726bc230d 100644
--- a/src/fabric/src/fabric2_encryption.erl
+++ b/src/fabric/src/fabric2_encryption.erl
@@ -35,8 +35,9 @@
 
 -export([
     req_wrapped_kek/1,
-    do_encrypt/5,
-    do_decrypt/5
+    req_unwrap_kek/2,
+    do_encrypt/4,
+    do_decrypt/4
 ]).
 
 
@@ -143,9 +144,9 @@ handle_call({encrypt, WrappedKEK, DbName, DocId, DocRev, Value}, From, St) ->
         waiters := Waiters
     } = St,
 
-    {ok, KEK, AAD} = unwrap_kek(Cache, DbName, WrappedKEK),
+    {ok, Future} = unwrap_kek(Cache, DbName, WrappedKEK),
     {_Pid, Ref} = erlang:spawn_monitor(?MODULE,
-        do_encrypt, [KEK, AAD, DocId, DocRev, Value]),
+        do_encrypt, [Future, DocId, DocRev, Value]),
 
     NewSt = St#{
         waiters := dict:store(Ref, From, Waiters)
@@ -158,9 +159,9 @@ handle_call({decrypt, WrappedKEK, DbName, DocId, DocRev, Value}, From, St) ->
         waiters := Waiters
     } = St,
 
-    {ok, KEK, AAD} = unwrap_kek(Cache, DbName, WrappedKEK),
+    {ok, Future} = unwrap_kek(Cache, DbName, WrappedKEK),
     {_Pid, Ref} = erlang:spawn_monitor(?MODULE,
-        do_decrypt, [KEK, AAD, DocId, DocRev, Value]),
+        do_decrypt, [Future, DocId, DocRev, Value]),
 
     NewSt = St#{
         waiters := dict:store(Ref, From, Waiters)
@@ -172,26 +173,47 @@ handle_cast(Msg, St) ->
     {stop, {bad_cast, Msg}, St}.
 
 
-handle_info({'DOWN', Ref, process, _Pid, Resp}, St) ->
+handle_info({'DOWN', Ref, process, Pid, Resp}, St) ->
     #{
         cache := Cache,
         waiters := Waiters
     } = St,
 
+    Reply = case Resp of
+        {kek, {ok, KEK, WrappedKEK, AAD}} ->
+            Entry = #entry{id = WrappedKEK, kek = KEK, aad = AAD},
+            true = ets:insert(Cache, Entry),
+            {ok, WrappedKEK};
+        {kek, {error, Error}} ->
+            ets:match_delete(Cache, #entry{kek = Pid, _ = '_'}),
+            {error, Error};
+        {retry, Fun, WrappedKEK, DocId, DocRev, Value} ->
+            case ets:lookup(Cache, WrappedKEK) of
+                [#entry{} = Entry] ->
+                    Future = get_future(Entry),
+                    {retry, Fun, Future, DocId, DocRev, Value};
+                [] ->
+                    {error, unable_to_unwrap_kek}
+            end;
+        _ ->
+            Resp
+    end,
+
     case dict:take(Ref, Waiters) of
         {From, Waiters1} ->
-            case Resp of
-                {kek, {ok, KEK, WrappedKEK, AAD}} ->
-                    Entry = #entry{id = WrappedKEK, kek = KEK, aad = AAD},
-                    true = ets:insert(Cache, Entry),
-                    gen_server:reply(From, {ok, WrappedKEK});
+            NewSt = case Reply of
+                {retry, Fun1, Future1, DocId1, DocRev1, Value1} ->
+                    {_, Ref1} = erlang:spawn_monitor(?MODULE,
+                        Fun1, [Future1, DocId1, DocRev1, Value1]),
+                    St#{
+                        waiters := dict:store(Ref1, From, Waiters1)
+                    };
                 _ ->
-                    gen_server:reply(From, Resp)
+                    gen_server:reply(From, Reply),
+                    St#{
+                        waiters := Waiters1
+                    }
             end,
-
-            NewSt = St#{
-                waiters := Waiters1
-            },
             {noreply, NewSt, ?TIMEOUT};
         error ->
             {noreply, St, ?TIMEOUT}
@@ -216,18 +238,47 @@ req_wrapped_kek(DbName) ->
         Resp ->
             exit({kek, Resp})
     catch
+        _:{badmatch, {error, Error}} ->
+            exit({kek, {error, Error}});
         _:Error ->
-            exit({error, Error})
+            exit({kek, {error, Error}})
+    end.
+
+
+req_unwrap_kek(DbName, WrappedKEK) ->
+    process_flag(sensitive, true),
+    try
+        {ok, AAD} = ?PLUGIN:get_aad(DbName),
+        {ok, KEK, WrappedKEK} = ?PLUGIN:unwrap_kek(DbName, WrappedKEK),
+        {ok, KEK, WrappedKEK, AAD}
+    of
+        Resp ->
+            exit({kek, Resp})
+    catch
+        _:{badmatch, {error, Error}} ->
+            exit({kek, {error, Error}});
+        _:Error ->
+            exit({kek, {error, Error}})
     end.
 
 
-do_encrypt(KEK, AAD, DocId, DocRev, Value) ->
+do_encrypt(Future, DocId, DocRev, Value) ->
     process_flag(sensitive, true),
     try
-        {ok, DEK} = get_dek(KEK, DocId, DocRev),
-        {CipherText, CipherTag} = ?aes_gcm_encrypt(DEK, <<0:96>>, AAD, Value),
-        <<CipherTag/binary, CipherText/binary>>
+        case Future() of
+            {ok, KEK, AAD} ->
+                {ok, DEK} = get_dek(KEK, DocId, DocRev),
+                {CipherText, CipherTag} = ?aes_gcm_encrypt(DEK, <<0:96>>,
+                    AAD, Value),
+                <<CipherTag/binary, CipherText/binary>>;
+            Else ->
+                Else
+        end
     of
+        {retry, WrappedKEK} ->
+            exit({retry, do_encrypt, WrappedKEK, DocId, DocRev, Value});
+        {error, Error} ->
+            exit({error, Error});
         Resp ->
             exit({ok, Resp})
     catch
@@ -236,13 +287,22 @@ do_encrypt(KEK, AAD, DocId, DocRev, Value) ->
     end.
 
 
-do_decrypt(KEK, AAD, DocId, DocRev, Value) ->
+do_decrypt(Future, DocId, DocRev, Value) ->
     process_flag(sensitive, true),
     try
-        <<CipherTag:16/binary, CipherText/binary>> = Value,
-        {ok, DEK} = get_dek(KEK, DocId, DocRev),
-        ?aes_gcm_decrypt(DEK, <<0:96>>, AAD, CipherText, CipherTag)
+        case Future() of
+            {ok, KEK, AAD} ->
+                <<CipherTag:16/binary, CipherText/binary>> = Value,
+                {ok, DEK} = get_dek(KEK, DocId, DocRev),
+                ?aes_gcm_decrypt(DEK, <<0:96>>, AAD, CipherText, CipherTag);
+            Else ->
+                Else
+        end
     of
+        {retry, WrappedKEK} ->
+            exit({retry, do_decrypt, WrappedKEK, DocId, DocRev, Value});
+        {error, Error} ->
+            exit({error, Error});
         Resp ->
             exit({ok, Resp})
     catch
@@ -261,14 +321,40 @@ get_dek(KEK, DocId, DocRev) when bit_size(KEK) == 256 ->
 
 unwrap_kek(Cache, DbName, WrappedKEK) ->
     case ets:lookup(Cache, WrappedKEK) of
-        [#entry{id = WrappedKEK, kek = KEK, aad = AAD}] ->
-            {ok, KEK, AAD};
+        [#entry{id = WrappedKEK} = Entry] ->
+            {ok, get_future(Entry)};
         [] ->
-            {ok, AAD} = ?PLUGIN:get_aad(DbName),
-            {ok, KEK, WrappedKEK} = ?PLUGIN:unwrap_kek(DbName, WrappedKEK),
-            Entry = #entry{id = WrappedKEK, kek = KEK, aad = AAD},
+            {Pid, _Ref} = erlang:spawn_monitor(?MODULE,
+                req_unwrap_kek, [DbName, WrappedKEK]),
+            Entry = #entry{id = WrappedKEK, kek = Pid},
             true = ets:insert(Cache, Entry),
-            {ok, KEK, AAD}
+            {ok, get_future(Entry)}
+    end.
+
+
+get_future(#entry{id = WrappedKEK, kek = Pid}) when is_pid(Pid) ->
+    fun() ->
+        Ref = erlang:monitor(process, Pid),
+        receive
+            {'DOWN', Ref, process, Pid, {kek, {ok, KEK, _, AAD}}} ->
+                {ok, KEK, AAD};
+            {'DOWN', Ref, process, Pid, {kek, {error, Error}}} ->
+                {error, Error};
+            {'DOWN', Ref, process, Pid,  noproc} ->
+                {retry, WrappedKEK};
+            {'DOWN', Ref, process, Pid,  {error, Error}} ->
+                {error, Error};
+            {'DOWN', Ref, process, Pid,  Else} ->
+                {error, Else}
+        after
+            ?TIMEOUT ->
+                {error, timeout}
+        end
+    end;
+
+get_future(#entry{kek = KEK, aad = AAD}) ->
+    fun() ->
+        {ok, KEK, AAD}
     end.
 
 
@@ -293,7 +379,8 @@ encrypt_decrypt_test_() ->
         end,
         fun(ok) ->
             [
-                fun test_do_encrypt_decrypt/0
+                fun test_do_encrypt_decrypt/0,
+                fun test_retry/0
             ]
         end
     }.
@@ -305,18 +392,43 @@ test_do_encrypt_decrypt() ->
     DocRev = <<"1-abcdefgh">>,
     Value = term_to_binary({{{[{<<"text">>, <<"test">>}]}, [], false}}),
 
+    Future = get_future(#entry{kek = KEK, aad = AAD}),
+
     {ok, EncResult} = try
-        do_encrypt(KEK, AAD, DocId, DocRev, Value)
+        do_encrypt(Future, DocId, DocRev, Value)
     catch
         exit:ER -> ER
     end,
     ?assertNotEqual(Value, EncResult),
 
     {ok, DecResult} = try
-        do_decrypt(KEK, AAD, DocId, DocRev, EncResult)
+        do_decrypt(Future, DocId, DocRev, EncResult)
     catch
         exit:DR -> DR
     end,
     ?assertEqual(Value, DecResult).
 
+test_retry() ->
+    WrappedKEK = crypto:strong_rand_bytes(32),
+    DocId = <<"0001">>,
+    DocRev = <<"1-abcdefgh">>,
+    Value = <<0:320>>,
+
+    DeadPid = erlang:spawn(fun() -> ok end),
+    Future = get_future(#entry{id = WrappedKEK, kek = DeadPid}),
+
+    EncResult = try
+        do_encrypt(Future, DocId, DocRev, Value)
+    catch
+        exit:ER -> ER
+    end,
+    ?assertMatch({retry, do_encrypt, _, _, _, _}, EncResult),
+
+    DecResult = try
+        do_decrypt(Future, DocId, DocRev, Value)
+    catch
+        exit:DR -> DR
+    end,
+    ?assertMatch({retry, do_decrypt, _, _, _, _}, DecResult).
+
 -endif.
diff --git a/src/fabric/test/fabric2_encryption_tests.erl b/src/fabric/test/fabric2_encryption_tests.erl
new file mode 100644
index 000000000..4b65eedee
--- /dev/null
+++ b/src/fabric/test/fabric2_encryption_tests.erl
@@ -0,0 +1,208 @@
+% 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(fabric2_encryption_tests).
+
+
+-include_lib("couch/include/couch_eunit.hrl").
+-include_lib("eunit/include/eunit.hrl").
+-include("fabric2_test.hrl").
+
+
+-define(PLUGIN, fabric2_encryption_provider).
+-define(WK, <<0:256>>).
+-define(FIXTURE, #{
+    db => <<"db">>,
+    doc_id => <<"0001">>,
+    doc_rev => <<"1-abcdef">>,
+    val => term_to_binary({[], [], false}),
+    enc => <<16#32AA68545ACE6A71466E25089101C5BC378019F304582A9FF9BD4A6F4F:232>>
+}).
+
+
+encryption_basic_test_() ->
+    {
+        setup,
+        fun basic_setup/0,
+        fun teardown/1,
+        [
+            fun test_get_wrapped_kek/0,
+            fun test_basic_encrypt/0,
+            fun test_basic_decrypt/0
+        ]
+    }.
+
+
+basic_setup() ->
+    Ctx = test_util:start_couch([fabric]),
+    meck:new([?PLUGIN], [passthrough]),
+    ok = meck:expect(?PLUGIN, get_kek, 1, {ok, ?WK, ?WK}),
+    ok = meck:expect(?PLUGIN, unwrap_kek, 2, {ok, ?WK, ?WK}),
+    Ctx.
+
+
+teardown(Ctx) ->
+    meck:unload(),
+    test_util:stop_couch(Ctx).
+
+
+test_get_wrapped_kek() ->
+    #{
+        db := DbName
+    } = ?FIXTURE,
+
+    Result = fabric2_encryption:get_wrapped_kek(DbName),
+    ?assertMatch({ok, _}, Result),
+
+    {ok, WrappedKEK} = Result,
+    ?assertEqual(?WK, WrappedKEK).
+
+
+test_basic_encrypt() ->
+    #{
+        db := DbName,
+        doc_id := DocId,
+        doc_rev := DocRev,
+        val := Value,
+        enc := Expected
+    } = ?FIXTURE,
+
+    Result = fabric2_encryption:encrypt(?WK, DbName, DocId, DocRev, Value),
+    ?assertMatch({ok, _}, Result),
+
+    {ok, Encrypted} = Result,
+    ?assertNotEqual(Value, Encrypted),
+    ?assertEqual(Expected, Encrypted).
+
+
+test_basic_decrypt() ->
+    #{
+        db := DbName,
+        doc_id := DocId,
+        doc_rev := DocRev,
+        val := Value,
+        enc := Expected
+    } = ?FIXTURE,
+
+    Result = fabric2_encryption:decrypt(?WK, DbName, DocId, DocRev, Expected),
+    ?assertMatch({ok, _}, Result),
+
+    {ok, Decrypted} = Result,
+    ?assertNotEqual(Expected, Decrypted),
+    ?assertEqual(Value, Decrypted).
+
+
+
+encryption_cache_test_() ->
+    {
+        foreach,
+        fun cache_setup/0,
+        fun teardown/1,
+        [
+            ?TDEF_FE(test_cache_encrypt),
+            ?TDEF_FE(test_cache_decrypt)
+        ]
+    }.
+
+
+cache_setup() ->
+    Ctx = test_util:start_couch([fabric]),
+    meck:new([?PLUGIN], [passthrough]),
+    ok = meck:expect(?PLUGIN, unwrap_kek, 2, {ok, ?WK, ?WK}),
+    Ctx.
+
+
+test_cache_encrypt(_) ->
+    #{
+        db := DbName,
+        doc_id := DocId,
+        doc_rev := DocRev,
+        val := Value,
+        enc := Expected
+    } = ?FIXTURE,
+
+    {ok, V1} = fabric2_encryption:encrypt(?WK, DbName, DocId, DocRev, Value),
+    {ok, V2} = fabric2_encryption:encrypt(?WK, DbName, DocId, DocRev, Value),
+
+    ?assertEqual(Expected, V1),
+    ?assertEqual(Expected, V2),
+    ?assertEqual(1, meck:num_calls(?PLUGIN, unwrap_kek, 2)).
+
+
+test_cache_decrypt(_) ->
+    #{
+        db := DbName,
+        doc_id := DocId,
+        doc_rev := DocRev,
+        val := Value,
+        enc := Expected
+    } = ?FIXTURE,
+
+    {ok, V1} = fabric2_encryption:decrypt(?WK, DbName, DocId, DocRev, Expected),
+    {ok, V2} = fabric2_encryption:decrypt(?WK, DbName, DocId, DocRev, Expected),
+
+    ?assertEqual(Value, V1),
+    ?assertEqual(Value, V2),
+    ?assertEqual(1, meck:num_calls(?PLUGIN, unwrap_kek, 2)).
+
+
+
+encryption_failure_test_() ->
+    {
+        foreach,
+        fun faiure_setup/0,
+        fun teardown/1,
+        [
+            ?TDEF_FE(test_failure_encrypt),
+            ?TDEF_FE(test_failure_decrypt)
+        ]
+    }.
+
+
+faiure_setup() ->
+    Ctx = test_util:start_couch([fabric]),
+    meck:new([?PLUGIN], [passthrough]),
+    ok = meck:expect(?PLUGIN, unwrap_kek, 2, {error, unable_to_unwrap_kek}),
+    Ctx.
+
+
+test_failure_encrypt(_) ->
+    #{
+        db := DbName,
+        doc_id := DocId,
+        doc_rev := DocRev,
+        val := Value
+    } = ?FIXTURE,
+
+    Result1 = fabric2_encryption:encrypt(?WK, DbName, DocId, DocRev, Value),
+    Result2 = fabric2_encryption:encrypt(?WK, DbName, DocId, DocRev, Value),
+
+    %% make sure we don't overwrap error messages and don't cache the errors
+    ?assertEqual({error, unable_to_unwrap_kek}, Result1),
+    ?assertEqual({error, unable_to_unwrap_kek}, Result2),
+    ?assertEqual(2, meck:num_calls(?PLUGIN, unwrap_kek, 2)).
+
+
+test_failure_decrypt(_) ->
+    #{
+        db := DbName,
+        doc_id := DocId,
+        doc_rev := DocRev,
+        enc := Expected
+    } = ?FIXTURE,
+
+    Result1 = fabric2_encryption:decrypt(?WK, DbName, DocId, DocRev, Expected),
+    Result2 = fabric2_encryption:decrypt(?WK, DbName, DocId, DocRev, Expected),
+
+    ?assertEqual({error, unable_to_unwrap_kek}, Result1),
+    ?assertEqual({error, unable_to_unwrap_kek}, Result2),
+    ?assertEqual(2, meck:num_calls(?PLUGIN, unwrap_kek, 2)).