Add encryption for database valuesaegis-old

6 files changed, 325 insertions, 27 deletions

diff --git a/rebar.config.script b/rebar.config.script
index 6f9f65c73..118a99e53 100644
--- a/rebar.config.script
+++ b/rebar.config.script
@@ -114,6 +114,7 @@ os:putenv("COUCHDB_APPS_CONFIG_DIR", filename:join([COUCHDB_ROOT, "rel/apps"])).
 SubDirs = [
     %% must be compiled first as it has a custom behavior
     "src/couch_epi",
+    "src/aegis",
     "src/couch_log",
     "src/chttpd",
     "src/couch",
diff --git a/rel/reltool.config b/rel/reltool.config
index 9fbf28544..1e64a808d 100644
--- a/rel/reltool.config
+++ b/rel/reltool.config
@@ -90,6 +90,7 @@
     {app, xmerl, [{incl_cond, include}]},
 
     %% couchdb
+    {app, aegis, [{incl_cond, include}]},
     {app, b64url, [{incl_cond, include}]},
     {app, bear, [{incl_cond, include}]},
     {app, chttpd, [{incl_cond, include}]},
diff --git a/src/aegis/src/aegis.app.src b/src/aegis/src/aegis.app.src
new file mode 100644
index 000000000..eb3018f1b
--- /dev/null
+++ b/src/aegis/src/aegis.app.src
@@ -0,0 +1,29 @@
+% 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.
+
+{application, aegis,
+ [{description, "If it's good enough for Zeus, it's good enough for CouchDB"},
+  {vsn, git},
+  {registered, []},
+  {applications,
+   [kernel,
+    stdlib,
+    crypto,
+    erlfdb
+   ]},
+  {env,[]},
+  {modules, []},
+
+  {maintainers, []},
+  {licenses, []},
+  {links, []}
+ ]}.
diff --git a/src/aegis/src/aegis.erl b/src/aegis/src/aegis.erl
new file mode 100644
index 000000000..5937b4c8e
--- /dev/null
+++ b/src/aegis/src/aegis.erl
@@ -0,0 +1,186 @@
+% 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).
+
+-define(IS_AEGIS_FUTURE, {aegis_future, _}).
+%% encapsulation violation :/
+-define(IS_FUTURE, {erlfdb_future, _, _}).
+
+%% Assume old crypto api
+-define(hmac(Key, PlainText), crypto:hmac(sha256, Key, PlainText)).
+-define(aes_gcm_encrypt(Key, IV, AAD, Data),
+        crypto:block_encrypt(aes_gcm, Key, IV, {AAD, Data, 16})).
+-define(aes_gcm_decrypt(Key, IV, AAD, CipherText, CipherTag),
+        crypto:block_decrypt(aes_gcm, Key, IV, {AAD, CipherText, CipherTag})).
+
+%% Replace macros if new crypto api is available
+-ifdef(OTP_RELEASE).
+-if(?OTP_RELEASE >= 22).
+-undef(hmac).
+-define(hmac(Key, PlainText), crypto:mac(hmac, sha256, Key, PlainText)).
+-undef(aes_gcm_encrypt).
+-define(aes_gcm_encrypt(Key, IV, AAD, Data),
+        crypto:crypto_one_time_aead(aes_256_gcm, Key, IV, Data, AAD, 16, true)).
+-undef(aes_gcm_decrypt).
+-define(aes_gcm_decrypt(Key, IV, AAD, CipherText, CipherTag),
+        crypto:crypto_one_time_aead(aes_256_gcm, Key, IV, CipherText,
+        AAD, CipherTag, false)).
+-endif.
+-endif.
+
+-export([
+    fold_range/6,
+    fold_range/7,
+    fold_range_future/5,
+    fold_range_wait/4,
+    get/3,
+    get_range/4,
+    get_range/5,
+    get_range_startswith/3,
+    get_range_startswith/4,
+    set/4,
+    wait/1
+]).
+
+
+fold_range(EncryptionContext, DbOrTx, StartKey, EndKey, Fun, Acc) ->
+    fold_range(EncryptionContext, DbOrTx, StartKey, EndKey, Fun, Acc, []).
+
+
+fold_range(EncryptionContext, DbOrTx, StartKey, EndKey, Fun, Acc, Options) ->
+    validate_encryption_context(EncryptionContext),
+    erlfdb:fold_range(DbOrTx, StartKey, EndKey, decrypt_fun(EncryptionContext, Fun), Acc, Options).
+
+
+fold_range_future(EncryptionContext, TxOrSs, StartKey, EndKey, Options) ->
+    validate_encryption_context(EncryptionContext),
+    Future = erlfdb:fold_range_future(TxOrSs, StartKey, EndKey, Options),
+    {aegis_fold_future, EncryptionContext, Future}.
+
+
+fold_range_wait(Tx, {aegis_fold_future, EncryptionContext, Future}, Fun, Acc) ->
+    validate_encryption_context(EncryptionContext),
+    erlfdb:fold_range_wait(Tx, Future, decrypt_fun(EncryptionContext, Fun), Acc).
+
+
+get(EncryptionContext, DbOrTx, Key) ->
+    validate_encryption_context(EncryptionContext),
+    Result = erlfdb:get(DbOrTx, Key),
+    decrypt(EncryptionContext, Key, Result).
+
+
+get_range(EncryptionContext, DbOrTx, StartKey, EndKey) ->
+    get_range(EncryptionContext, DbOrTx, StartKey, EndKey, []).
+
+
+get_range(EncryptionContext, DbOrTx, StartKey, EndKey, Options) ->
+    validate_encryption_context(EncryptionContext),
+    Result = erlfdb:get_range(DbOrTx, StartKey, EndKey, Options),
+    decrypt(EncryptionContext, Result).
+
+
+get_range_startswith(EncryptionContext, DbOrTx, Prefix) ->
+    get_range_startswith(EncryptionContext, DbOrTx, Prefix, []).
+
+
+get_range_startswith(EncryptionContext, DbOrTx, Prefix, Options) ->
+    validate_encryption_context(EncryptionContext),
+    Result = erlfdb:get_range_startswith(DbOrTx, Prefix, Options),
+    decrypt(EncryptionContext, Result).
+
+
+set(EncryptionContext, DbOrTx, Key, Value) ->
+    validate_encryption_context(EncryptionContext),
+    erlfdb:set(DbOrTx, Key, encrypt(EncryptionContext, Key, Value)).
+
+
+wait({aegis_future, EncryptionContext, Future}) ->
+    Value = erlfdb:wait(Future),
+    decrypt(EncryptionContext, Value);
+
+wait({aegis_future, EncryptionContext, Key, Future}) ->
+    Value = erlfdb:wait(Future),
+    decrypt(EncryptionContext, Key, Value);
+
+wait(Result) ->
+    Result.
+
+
+%% Private functions
+
+validate_encryption_context(#{uuid := _UUID}) ->
+    ok;
+validate_encryption_context(_) ->
+    error(invalid_encryption_context).
+
+
+-define(DUMMY_KEY, <<1:256>>).
+
+encrypt(#{uuid := UUID}, Key, Value) ->
+    {CipherText, <<CipherTag:128>>} =
+        ?aes_gcm_encrypt(
+           derive(?DUMMY_KEY, Key),
+           <<0:96>>,
+           UUID,
+           Value),
+    <<1:8, CipherTag:128, CipherText/binary>>.
+
+
+decrypt(EncryptionContext, ?IS_FUTURE = Future) ->
+    decrypt_future(EncryptionContext, Future);
+
+decrypt(EncryptionContext, {Key, Value})
+  when is_binary(Key), is_binary(Value) ->
+    decrypt(EncryptionContext, Key, Value);
+
+decrypt(EncryptionContext, Rows) when is_list(Rows) ->
+    [{Key, decrypt(EncryptionContext, Row)} || {Key, _} = Row <- Rows].
+
+
+decrypt(EncryptionContext, Key, ?IS_FUTURE = Future) ->
+    decrypt_future(EncryptionContext, Key, Future);
+
+decrypt(#{uuid := UUID}, Key, Value) when is_binary(Value) ->
+    <<1:8, CipherTag:128, CipherText/binary>> = Value,
+    Decrypted =
+        ?aes_gcm_decrypt(
+           derive(?DUMMY_KEY, Key),
+           <<0:96>>,
+           UUID,
+           CipherText,
+           <<CipherTag:128>>),
+    case Decrypted of
+        error ->
+            erlang:error(decryption_failed);
+        Decrypted ->
+            Decrypted
+    end;
+
+decrypt(_EncryptionContext, _Key, Value) when not is_binary(Value) ->
+    Value.
+
+
+decrypt_future(EncryptionContext, ?IS_FUTURE = Future) ->
+    {aegis_future, EncryptionContext, Future}.
+
+decrypt_future(EncryptionContext, Key, ?IS_FUTURE = Future) ->
+    {aegis_future, EncryptionContext, Key, Future}.
+
+decrypt_fun(EncryptionContext, Fun) ->
+    fun(Rows, Acc) ->
+            Fun(decrypt(EncryptionContext, Rows), Acc)
+    end.
+
+derive(KEK, KeyMaterial) when bit_size(KEK) == 256 ->
+    PlainText = <<1:16, "aegis", 0:8, KeyMaterial/binary, 256:16>>,
+    <<_:256>> = ?hmac(KEK, PlainText).
diff --git a/src/aegis/test/aegis_basic_test.erl b/src/aegis/test/aegis_basic_test.erl
new file mode 100644
index 000000000..061f7242a
--- /dev/null
+++ b/src/aegis/test/aegis_basic_test.erl
@@ -0,0 +1,77 @@
+% 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_basic_test).
+
+-include_lib("eunit/include/eunit.hrl").
+
+-define(DB, #{uuid => <<"foo">>}).
+
+get_set_db_test() ->
+    Db = erlfdb_util:get_test_db([empty]),
+    Key = <<"foo">>,
+    Value = <<"bar">>,
+    ?assertEqual(ok, aegis:set(?DB, Db, Key, Value)),
+    ?assertNotEqual(Value, erlfdb:get(Db, Key)),
+    ?assertEqual(Value, aegis:get(?DB, Db, Key)).
+
+get_set_tx_test() ->
+    Db = erlfdb_util:get_test_db([empty]),
+    Key = <<"foo">>,
+    Value = <<"bar">>,
+    ?assertEqual(ok, aegis:set(?DB, Db, Key, Value)),
+    Tx = erlfdb:create_transaction(Db),
+    Future = aegis:get(?DB, Tx, Key),
+    ?assertEqual(Value, aegis:wait(Future)).
+
+get_range_test() ->
+    Db = erlfdb_util:get_test_db([empty]),
+    Rows = [{<<"foo1">>, <<"bar1">>},
+            {<<"foo2">>, <<"bar2">>},
+            {<<"foo3">>, <<"bar3">>}],
+    [aegis:set(?DB, Db, K, V) || {K, V} <- Rows],
+    ?assertNotEqual(Rows, erlfdb:get_range(Db, <<"foo1">>, <<"foo9">>)),
+    ?assertEqual(Rows, aegis:get_range(?DB, Db, <<"foo1">>, <<"foo9">>)).
+
+get_range_startswith_test() ->
+    Db = erlfdb_util:get_test_db([empty]),
+    Rows = [{<<"foo1">>, <<"bar1">>},
+            {<<"foo2">>, <<"bar2">>},
+            {<<"foo3">>, <<"bar3">>}],
+    [aegis:set(?DB, Db, K, V) || {K, V} <- Rows],
+    ?assertNotEqual(Rows, erlfdb:get_range_startswith(Db, <<"foo">>)),
+    ?assertEqual(Rows, aegis:get_range_startswith(?DB, Db, <<"foo">>)).
+
+fold_range_test() ->
+    Db = erlfdb_util:get_test_db([empty]),
+    Rows = [{<<"foo1">>, <<"bar1">>},
+            {<<"foo2">>, <<"bar2">>},
+            {<<"foo3">>, <<"bar3">>}],
+    {_, Values} = lists:unzip(Rows),
+    [aegis:set(?DB, Db, K, V) || {K, V} <- Rows],
+    Fun = fun(NewRows, Acc) -> [NewRows | Acc] end,
+    ?assertNotEqual(Values, lists:reverse(erlfdb:fold_range(Db, <<"foo1">>, <<"foo9">>, Fun, []))),
+    ?assertEqual(Values, lists:reverse(aegis:fold_range(?DB, Db, <<"foo1">>, <<"foo9">>, Fun, []))).
+
+fold_range_future_test() ->
+    Db = erlfdb_util:get_test_db([empty]),
+    Rows = [{<<"foo1">>, <<"bar1">>},
+            {<<"foo2">>, <<"bar2">>},
+            {<<"foo3">>, <<"bar3">>}],
+    {_, Values} = lists:unzip(Rows),
+    [aegis:set(?DB, Db, K, V) || {K, V} <- Rows],
+    Fun = fun(NewRows, Acc) -> [NewRows | Acc] end,
+    Tx = erlfdb:create_transaction(Db),
+    Future = aegis:fold_range_future(?DB, Tx, <<"foo1">>, <<"foo9">>, []),
+    ?assertEqual(Values, lists:reverse(aegis:fold_range_wait(Tx, Future, Fun, []))).
+
+
diff --git a/src/fabric/src/fabric2_fdb.erl b/src/fabric/src/fabric2_fdb.erl
index 2295a5648..c77651839 100644
--- a/src/fabric/src/fabric2_fdb.erl
+++ b/src/fabric/src/fabric2_fdb.erl
@@ -406,12 +406,16 @@ get_info(#{} = Db) ->
         tx := Tx,
         db_prefix := DbPrefix
     } = ensure_current(Db),
-    get_info_wait(get_info_future(Tx, DbPrefix)).
+    get_info_wait(get_info_future(Db, Tx, DbPrefix)).
 
 
+%% Needs more thought.
 get_info_future(Tx, DbPrefix) ->
+    get_info_future(will_crash, Tx, DbPrefix).
+
+get_info_future(Db, Tx, DbPrefix) ->
     {CStart, CEnd} = erlfdb_tuple:range({?DB_CHANGES}, DbPrefix),
-    ChangesFuture = erlfdb:get_range(Tx, CStart, CEnd, [
+    ChangesFuture = aegis:get_range(Db, Tx, CStart, CEnd, [
         {streaming_mode, exact},
         {limit, 1},
         {reverse, true}
@@ -538,12 +542,12 @@ get_all_revs(#{} = Db, DocId) ->
 
         Prefix = erlfdb_tuple:pack({?DB_REVS, DocId}, DbPrefix),
         Options = [{streaming_mode, want_all}],
-        Future = erlfdb:get_range_startswith(Tx, Prefix, Options),
+        Future = aegis:get_range_startswith(Db, Tx, Prefix, Options),
         lists:map(fun({K, V}) ->
             Key = erlfdb_tuple:unpack(K, DbPrefix),
             Val = erlfdb_tuple:unpack(V),
             fdb_to_revinfo(Key, Val)
-        end, erlfdb:wait(Future))
+        end, aegis:wait(Future))
     end).
 
 
@@ -563,7 +567,7 @@ get_winning_revs_future(#{} = Db, DocId, NumRevs) ->
 
     {StartKey, EndKey} = erlfdb_tuple:range({?DB_REVS, DocId}, DbPrefix),
     Options = [{reverse, true}, {limit, NumRevs}],
-    erlfdb:fold_range_future(Tx, StartKey, EndKey, Options).
+    aegis:fold_range_future(Db, Tx, StartKey, EndKey, Options).
 
 
 get_winning_revs_wait(#{} = Db, RangeFuture) ->
@@ -571,7 +575,7 @@ get_winning_revs_wait(#{} = Db, RangeFuture) ->
         tx := Tx,
         db_prefix := DbPrefix
     } = ensure_current(Db),
-    RevRows = erlfdb:fold_range_wait(Tx, RangeFuture, fun({K, V}, Acc) ->
+    RevRows = aegis:fold_range_wait(Tx, RangeFuture, fun({K, V}, Acc) ->
         Key = erlfdb_tuple:unpack(K, DbPrefix),
         Val = erlfdb_tuple:unpack(V),
         [fdb_to_revinfo(Key, Val) | Acc]
@@ -589,7 +593,7 @@ get_non_deleted_rev(#{} = Db, DocId, RevId) ->
 
     BaseKey = {?DB_REVS, DocId, true, RevPos, Rev},
     Key = erlfdb_tuple:pack(BaseKey, DbPrefix),
-    case erlfdb:wait(erlfdb:get(Tx, Key)) of
+    case aegis:wait(aegis:get(Db, Tx, Key)) of
         not_found ->
             not_found;
         Val ->
@@ -630,7 +634,7 @@ get_doc_body_wait(#{} = Db0, DocId, RevInfo, Future) ->
         rev_path := RevPath
     } = RevInfo,
 
-    RevBodyRows = erlfdb:fold_range_wait(Tx, Future, fun({_K, V}, Acc) ->
+    RevBodyRows = aegis:fold_range_wait(Tx, Future, fun({_K, V}, Acc) ->
         [V | Acc]
     end, []),
     BodyRows = lists:reverse(RevBodyRows),
@@ -645,11 +649,11 @@ get_local_doc(#{} = Db0, <<?LOCAL_DOC_PREFIX, _/binary>> = DocId) ->
     } = Db = ensure_current(Db0),
 
     Key = erlfdb_tuple:pack({?DB_LOCAL_DOCS, DocId}, DbPrefix),
-    Rev = erlfdb:wait(erlfdb:get(Tx, Key)),
+    Rev = aegis:wait(aegis:get(Db, Tx, Key)),
 
     Prefix = erlfdb_tuple:pack({?DB_LOCAL_DOC_BODIES, DocId}, DbPrefix),
-    Future = erlfdb:get_range_startswith(Tx, Prefix),
-    Chunks = lists:map(fun({_K, V}) -> V end, erlfdb:wait(Future)),
+    Future = aegis:get_range_startswith(Db, Tx, Prefix),
+    Chunks = lists:map(fun({_K, V}) -> V end, aegis:wait(Future)),
 
     fdb_to_local_doc(Db, DocId, Rev, Chunks).
 
@@ -742,7 +746,7 @@ write_doc(#{} = Db0, Doc, NewWinner0, OldWinner, ToUpdate, ToRemove) ->
     lists:foreach(fun(RI0) ->
         RI = RI0#{winner := false},
         {K, V, undefined} = revinfo_to_fdb(Tx, DbPrefix, DocId, RI),
-        ok = erlfdb:set(Tx, K, V)
+        ok = aegis:set(Db, Tx, K, V)
     end, ToUpdate),
 
     lists:foreach(fun(RI0) ->
@@ -780,7 +784,7 @@ write_doc(#{} = Db0, Doc, NewWinner0, OldWinner, ToUpdate, ToRemove) ->
         _ ->
             ADKey = erlfdb_tuple:pack({?DB_ALL_DOCS, DocId}, DbPrefix),
             ADVal = erlfdb_tuple:pack(NewRevId),
-            ok = erlfdb:set(Tx, ADKey, ADVal)
+            ok = aegis:set(Db, Tx, ADKey, ADVal)
     end,
 
     % _changes
@@ -855,7 +859,7 @@ write_local_doc(#{} = Db0, Doc) ->
 
     {LDocKey, LDocVal, NewSize, Rows} = local_doc_to_fdb(Db, Doc),
 
-    {WasDeleted, PrevSize} = case erlfdb:wait(erlfdb:get(Tx, LDocKey)) of
+    {WasDeleted, PrevSize} = case aegis:wait(aegis:get(Db, Tx, LDocKey)) of
         <<255, RevBin/binary>> ->
             case erlfdb_tuple:unpack(RevBin) of
                 {?CURR_LDOC_FORMAT, _Rev, Size} ->
@@ -874,11 +878,11 @@ write_local_doc(#{} = Db0, Doc) ->
             erlfdb:clear(Tx, LDocKey),
             erlfdb:clear_range_startswith(Tx, BPrefix);
         false ->
-            erlfdb:set(Tx, LDocKey, LDocVal),
+            aegis:set(Db, Tx, LDocKey, LDocVal),
             % Make sure to clear the whole range, in case there was a larger
             % document body there before.
             erlfdb:clear_range_startswith(Tx, BPrefix),
-            lists:foreach(fun({K, V}) -> erlfdb:set(Tx, K, V) end, Rows)
+            lists:foreach(fun({K, V}) -> aegis:set(Db, Tx, K, V) end, Rows)
     end,
 
     case {WasDeleted, Doc#doc.deleted} of
@@ -902,7 +906,7 @@ read_attachment(#{} = Db, DocId, AttId) ->
     } = ensure_current(Db),
 
     AttKey = erlfdb_tuple:pack({?DB_ATTS, DocId, AttId}, DbPrefix),
-    case erlfdb:wait(erlfdb:get_range_startswith(Tx, AttKey)) of
+    case aegis:wait(aegis:get_range_startswith(Db, Tx, AttKey)) of
         not_found ->
             throw({not_found, missing});
         KVs ->
@@ -921,11 +925,11 @@ write_attachment(#{} = Db, DocId, Data) when is_binary(Data) ->
     Chunks = chunkify_binary(Data),
 
     IdKey = erlfdb_tuple:pack({?DB_ATT_NAMES, DocId, AttId}, DbPrefix),
-    ok = erlfdb:set(Tx, IdKey, <<>>),
+    ok = aegis:set(Db, Tx, IdKey, <<>>),
 
     lists:foldl(fun(Chunk, ChunkId) ->
         AttKey = erlfdb_tuple:pack({?DB_ATTS, DocId, AttId, ChunkId}, DbPrefix),
-        ok = erlfdb:set(Tx, AttKey, Chunk),
+        ok = aegis:set(Db, Tx, AttKey, Chunk),
         ChunkId + 1
     end, 0, Chunks),
     {ok, AttId}.
@@ -939,7 +943,7 @@ get_last_change(#{} = Db) ->
 
     {Start, End} = erlfdb_tuple:range({?DB_CHANGES}, DbPrefix),
     Options = [{limit, 1}, {reverse, true}],
-    case erlfdb:get_range(Tx, Start, End, Options) of
+    case aegis:get_range(Db, Tx, Start, End, Options) of
         [] ->
             vs_to_seq(fabric2_util:seq_zero_vs());
         [{K, _V}] ->
@@ -960,14 +964,14 @@ fold_range(TxOrDb, RangePrefix, UserFun, UserAcc, Options) ->
     case fabric2_util:get_value(limit, Options) of 0 -> UserAcc; _ ->
         FAcc = get_fold_acc(Db, RangePrefix, UserFun, UserAcc, Options),
         try
-            fold_range(Tx, FAcc)
+            fold_range(Db, Tx, FAcc)
         after
             erase(?PDICT_FOLD_ACC_STATE)
         end
     end.
 
 
-fold_range(Tx, FAcc) ->
+fold_range(#{} = Db, Tx, FAcc) ->
     #fold_acc{
         start_key = Start,
         end_key = End,
@@ -983,14 +987,14 @@ fold_range(Tx, FAcc) ->
     try
         #fold_acc{
             user_acc = FinalUserAcc
-        } = erlfdb:fold_range(Tx, Start, End, Callback, FAcc, Opts),
+        } = aegis:fold_range(Db, Tx, Start, End, Callback, FAcc, Opts),
         FinalUserAcc
     catch error:{erlfdb_error, ?TRANSACTION_TOO_OLD} when DoRestart ->
         % Possibly handle cluster_version_changed and future_version as well to
         % continue iteration instead fallback to transactional and retrying
         % from the beginning which is bound to fail when streaming data out to a
         % socket.
-        fold_range(Tx, restart_fold(Tx, FAcc))
+        fold_range(Db, Tx, restart_fold(Tx, FAcc))
     end.
 
 
@@ -1193,7 +1197,7 @@ write_doc_body(#{} = Db0, #doc{} = Doc) ->
 
     Rows = doc_to_fdb(Db, Doc),
     lists:foreach(fun({Key, Value}) ->
-        ok = erlfdb:set(Tx, Key, Value)
+        ok = aegis:set(Db, Tx, Key, Value)
     end, Rows).
 
 
@@ -1246,12 +1250,12 @@ cleanup_attachments(Db, DocId, NewDoc, ToRemove) ->
 
     AttPrefix = erlfdb_tuple:pack({?DB_ATT_NAMES, DocId}, DbPrefix),
     Options = [{streaming_mode, want_all}],
-    Future = erlfdb:get_range_startswith(Tx, AttPrefix, Options),
+    Future = aegis:get_range_startswith(Db, Tx, AttPrefix, Options),
 
     ExistingIdSet = lists:foldl(fun({K, _}, Acc) ->
         {?DB_ATT_NAMES, DocId, AttId} = erlfdb_tuple:unpack(K, DbPrefix),
         sets:add_element(AttId, Acc)
-    end, sets:new(), erlfdb:wait(Future)),
+    end, sets:new(), aegis:wait(Future)),
 
     AttsToRemove = sets:subtract(ExistingIdSet, ActiveIdSet),