path: root/src/ht.erl

%%% Copyright (c) 2014, NORDUnet A/S.
%%% See LICENSE for licensing information.
%%%
%%% Implementation of a history tree as described in Efficient Data
%%% Structures for Tamper-Evident Logging [0]. This implementation
%%% follows RFC 6962 and differs from [0] only in how non-full trees
%%% are handled.
%%%
%%% [0] https://www.usenix.org/event/sec09/tech/full_papers/crosby.pdf
%%%
%%% Hashes of inner nodes and leaves are stored in arrays, one per
%%% layer with layer 0 being where the leaves are. The total number of
%%% arrays is equal to the depth of the tree. The depth of the tree is
%%% ceil(lg2(number of leaves)).
%%% Let {r,i} denote the hash with index i on layer r. The first leaf
%%% is {0,0}, second is {0,1} and n:th is {0,n-1}.
%%% The parent of {r,i} is {r+1,floor(i/2)} (not strictly true because
%%% of "placeholder nodes", see update_parent/4).
%%% The sibling of {r,i} is {r,i+1} when i is even and {r,i-1} when i
%%% is odd.

-module(ht).
-behaviour(gen_server).

-export([size/0, add/1, tree_hash/0, tree_hash/1]).
-export([path/2, consistency/2]).
-export([start_link/0, start_link/1, stop/0]).
-export([init/1, handle_call/3, terminate/2, handle_cast/2, handle_info/2,
         code_change/3]).

-include("$CTROOT/plop/include/plop.hrl").
-include_lib("eunit/include/eunit.hrl").
-import(lists, [foreach/2, foldl/3, reverse/1]).

%% Data types.
-record(tree, {version :: integer(),
               evaluated :: integer(),
               store :: ts:tree_store()}).
-type tree() :: #tree{}.

%%%%%%%%%%%%%%%%%%%%
%% Public interface.
start_link() ->
    gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).
start_link(NEntries) ->
    gen_server:start_link({local, ?MODULE}, ?MODULE, [NEntries], []).
stop() ->
    gen_server:call(?MODULE, stop).
size() ->
    gen_server:call(?MODULE, size).
add(Entry) ->
    gen_server:call(?MODULE, {add, Entry}).
tree_hash() ->
    gen_server:call(?MODULE, tree_hash).
tree_hash(Version) ->
    gen_server:call(?MODULE, {tree_hash, Version}).
path(I, V) ->
    gen_server:call(?MODULE, {path, I, V}).
consistency(V1, V2) ->
    gen_server:call(?MODULE, {consistency, V1, V2}).

%% gen_server callbacks
init([]) ->
    {ok, new()};
init(Args) ->
    {ok, new(Args)}.
handle_cast(_Request, State) ->
    {noreply, State}.
handle_info(_Info, State) ->
    {noreply, State}.
code_change(_OldVersion, State, _Extra) ->
    {ok, State}.
terminate(_Reason, _State) ->
    ok.
handle_call(stop, _From, State) ->
    {stop, normal, stopped, State};
handle_call(size, _From, State) ->
    {reply, State#tree.version + 1, State};
handle_call({add, Entry}, _From, State) ->
    {reply, ok, add(State, Entry)};
handle_call(tree_hash, _From, State) ->
    {NewState, Hash} = head(State, State#tree.version),
    {reply, Hash, NewState};
handle_call({tree_hash, Version}, _From, State) ->
    {NewState, Hash} = head(State, Version),
    {reply, Hash, NewState};
handle_call({path, Index, Version}, _From, State) ->
    {NewState, Path} = path(State, Index, Version),
    {reply, Path, NewState};
handle_call({consistency, _Version1, _Version2}, _From, State) ->
    {reply, nyi, State}.

%%%%%%%%%%%%%%%%%%%%
%% Private.

%% @doc Return a list of hashes showing the path from leaf Index to
%% the tree head in the tree of version Version.
-spec path(tree(), non_neg_integer(), non_neg_integer()) -> {tree(), list()}.
path(Tree, _Index, -1) ->
    {Tree, []};
path(Tree, Index, Version) ->
    {Tree, path(update(Tree, Version), 0, Index, Version, Version, [])}.

-spec path(tree(), non_neg_integer(), non_neg_integer(), non_neg_integer(), non_neg_integer(), list()) -> list().
path(_, _, _, 0, _, Acc) ->
    reverse(Acc);
path(Tree, Layer, I, ILast, Version, Acc) ->
    path(Tree, Layer + 1, parent(I), parent(ILast), Version,
         case sibling(I) of
             Sib when Sib == ILast ->
                 %% We're at the edge of the layer and might need to
                 %% recompute an old tree.
                 [old_version_tree_head(Tree, Version, Layer) | Acc];
             Sib when Sib < ILast ->
                 %% Just use sibling.
                 [get_hash(Tree, {Sib, Layer}) | Acc];
             _ ->
                 %% Sibling is larger than ILast so doesn't exist.
                 Acc
         end).

%% @doc updates the tree and returns new tree plus hash for IR
-spec get_hash(tree(), tuple()) -> binary().
get_hash(Tree, IR) ->
    ts:retrieve_hash(Tree#tree.store, IR).

-spec head(tree(), integer()) -> {tree(), binary()}.
head(Tree, -1) ->
    {Tree, hash(<<"">>)};
head(Tree = #tree{version = V}, Version) when Version == V ->
    NewTree = update(Tree),
    {NewTree, get_hash(NewTree, {0, depth(Tree) - 1})};
head(Tree = #tree{version = V}, Version) when Version > V ->
    {Tree, enotimetravel};
head(Tree, Version) ->
    NewTree = update(Tree, Version),
    {NewTree, old_version_tree_head(NewTree, Version)}.

-spec old_version_tree_head(tree(), non_neg_integer()) -> binary().
old_version_tree_head(Tree, Version) ->
    old_version_tree_head(Tree, Version, -1).

-spec old_version_tree_head(tree(), non_neg_integer(), integer()) -> binary().
old_version_tree_head(Tree, Version, BreakAtLayer) ->
    true = Tree#tree.evaluated >= Version,      % ASSERTION
    %% Go up the tree from the rightmost leaf (index=Version) until a
    %% left node is found. (There is always one -- the head is a left
    %% node.)
    {FirstLeftR, FirstLeftI} = first_left_node(0, Version, BreakAtLayer),

    %% Walk up the tree from this lowest left node up to and including
    %% the last right node, rehashing as we go. Calculate the parent
    %% hash of that node and its sibling. Return that hash.
    last_right_node_rehash(Tree, Version, FirstLeftR, FirstLeftI,
                           get_hash(Tree, {FirstLeftI, FirstLeftR}),
                           BreakAtLayer).

-spec last_right_node_rehash(tree(), non_neg_integer(), non_neg_integer(),
                             non_neg_integer(), binary(), integer()) ->
                                    binary().
last_right_node_rehash(_, _, Layer, _, RightNodeHash, BAL) when Layer == BAL ->
    %% Bailing out at Layer.
    RightNodeHash;
last_right_node_rehash(_, _, _, 0, RightNodeHash, _) ->
    %% Index is 0, we're done.
    RightNodeHash;
last_right_node_rehash(Tree, Version, Layer, Index, RightNodeHash, BAL) ->
    last_right_node_rehash(
      Tree, Version, Layer + 1, parent(Index),
      case right_node_p(Index) of
          true ->
              %% Rehash parent using sibling.
              mkinnerhash(get_hash(Tree, {Index - 1, Layer}), RightNodeHash);
          false ->
              %% Just use the incoming hash.
              RightNodeHash
      end,
      BAL).

-spec first_left_node(non_neg_integer(), non_neg_integer(), non_neg_integer()) ->
                             {non_neg_integer(), non_neg_integer()}.
first_left_node(Layer, Index, BAL) when Layer == BAL ->
    {Layer, Index};
first_left_node(Layer, Index, BAL) ->
    case right_node_p(Index) of
        true -> first_left_node(Layer + 1, parent(Index), BAL);
        false -> {Layer, Index}
    end.

%% @doc Add an entry but don't update the tree.
-spec add(tree(), binary()) -> tree().
add(Tree = #tree{version = V, store = Store}, Entry) ->
    NewVersion = V + 1,
    LeafIndex = NewVersion,
    LeafHash = mkleafhash(Entry),
    Tree#tree{version = NewVersion,
              store = ts:store(Store, {LeafIndex, 0}, LeafHash)}.

-spec new() -> tree().
new() ->
    #tree{version = -1,
          evaluated = -1,
          store = ts:new()}.

-spec new([non_neg_integer()]) -> tree().
new([Version]) when is_integer(Version) ->
    foldl(fun(#mtl{entry = E}, Tree) ->
                  D = (E#timestamped_entry.entry)#plop_entry.data,
                  add(Tree, D) % Return value -> Tree in next invocation.
          end, new(), db:get_by_index_sorted(0, Version));
new([List]) when is_list(List) ->
    foldl(fun(D, Tree) ->
                  add(Tree, D) % Return value -> Tree in next invocation.
          end, new(), List).

update(Tree) ->
    update(Tree, Tree#tree.version).

%% @doc Calculate hashes in Tree up to and including node with index
%% equal to Version. Update Tree.evaluated to reflect the new state.
-spec update(tree(), non_neg_integer()) -> tree().
update(Tree, 0) ->
    %% A version 0 tree needs no updating.
    Tree;
update(Tree = #tree{evaluated = E}, V) when E >= V ->
    %% Evaluated enough already. Nothing to do.
    Tree;
update(Tree = #tree{version = MaxV}, V) when V > MaxV ->
    %% Asking for more than we've got. Do as much as possible.
    update(Tree, MaxV);
update(Tree = #tree{evaluated = Evaluated}, Version) ->
    NewTree = update_layer(Tree, 0, Evaluated + 1, Version),
    NewTree#tree{evaluated = Version}.

%% @doc Update the tree wrt the leaves ICur..ILast.
-spec update_layer(tree(), non_neg_integer(), non_neg_integer(),
                   non_neg_integer()) -> tree().
update_layer(Tree, _Layer, _ICur, 0) ->         % Done
    Tree;
update_layer(Tree, Layer, ICur, ILast) ->
    %% Update parents on next upper layer, starting with a left
    %% child <= ICur and ending with ILast. Recurse with next layer.
    NewStore = update_parent(Tree#tree.store, Layer,
                             strip_bits_bottom(ICur, 1), ILast),
    update_layer(Tree#tree{store = NewStore}, Layer + 1,
                 parent(ICur), parent(ILast)).

%% @doc Update parents of I..ILast, on Layer+1. I has to be a left child.
-spec update_parent(ts:tree_store(), non_neg_integer(), non_neg_integer(),
                    non_neg_integer()) -> ts:tree_store().
update_parent(S, Layer, I, ILast) when I >= ILast ->
    %% We're done updating parents. If ILast is a left child, copy it
    %% to where its parent would've been were it a right child. This
    %% is a "placeholder node" which simplifies creating incomplete
    %% ("non-frozen") trees.
    case right_node_p(ILast) of
        true -> S;
        _ -> ts:append(S, Layer + 1, ts:retrieve_hash(S, {ILast, Layer}))
    end;
update_parent(S, Layer, I, ILast) ->
    false = right_node_p(I),                    % ASSERTION
    %% Make an inner node hash of I and its sibling. Store it as
    %% parent. Recurse with next pair of leaves.
    %% TODO: This is the only place where we store rather than
    %% append. Consider changing this to an append. This would make it
    %% possible for ts to use other data structures for the storage
    %% (like a binary per layer). If so, don't forget to truncate
    %% Layer+1 before appending if there's already a parent for I
    %% there.
    update_parent(ts:store(S, {parent(I), Layer + 1},
                  mkinnerhash(ts:retrieve_hash(S, {I, Layer}),
                              ts:retrieve_hash(S, {I + 1, Layer}))),
                  Layer, I + 2, ILast).

%% @doc Parent of {i, r} is at {i/2, r+1} (unless it's a "placeholder").
parent(I) ->
    I bsr 1.

-spec right_node_p(integer()) -> boolean().
right_node_p(Index) ->
    case Index band 1 of
        1 -> true;
        _ -> false
    end.

-spec sibling(non_neg_integer()) -> non_neg_integer().
sibling(Index) ->
    case right_node_p(Index) of
        true -> Index - 1;
        false -> Index + 1
    end.

strip_bits_bottom(N, Nbits) ->
    (N bsr Nbits) bsl Nbits.

%% @doc Return position of highest bit set, counting from the least
%% significant bit, starting at 1.
bitpos_first_set(N) ->
    L = [Bit || <<Bit:1>> <= binary:encode_unsigned(N)],
    length(L) - ffs(L, 0).
ffs([], Acc) ->
    Acc;
ffs([H|T], Acc) ->
    case H of
        0 -> ffs(T, Acc + 1);
        _ -> Acc
    end.

depth(#tree{version = -1}) ->
    0;
depth(#tree{version = V}) ->
    bitpos_first_set(V) + 1.

-spec mkleafhash(binary()) -> binary().
mkleafhash(Data) ->
    hash([<<"\x00">>, Data]).

-spec mkinnerhash(binary(), binary()) -> binary().
mkinnerhash(Hash1, Hash2) ->
    hash([<<"\x01">>, Hash1, Hash2]).

-spec hash(binary()) -> binary() | iolist().
hash(Data) ->
    crypto:hash(sha256, Data).

%%%%%%%%%%%%%%%%%%%%
%% Testing ht.
%% TODO: Move all these tests to a separate file in ../test. They're
%% only using external functions.
-define(TEST_VECTOR_LEAVES,
        ["", "\x00", "\x10", " !", "01", "@ABC", "PQRSTUVW", "`abcdefghijklmno"]).

%% FIXME: Don't start and stop the server manually all the time. EUnit
%% can help!
test_init(L) ->
    stop(),
    {ok, _Pid} = start_link(L).

%% @doc Build tree using add/2 and mth/2 and compare the resulting
%% tree hashes.
%% FIXME: Move outside.
add_test() ->
    lists:foreach(
      fun(X) -> L = lists:sublist(?TEST_VECTOR_LEAVES, X),
                test_init(L),
                ?assertEqual(mth(L), tree_hash()) end,
      random_entries(length(?TEST_VECTOR_LEAVES))).

%% FIXME: Move outside.
old_versions_test() ->
    test_init(?TEST_VECTOR_LEAVES),
    ?assertEqual(mth(?TEST_VECTOR_LEAVES), tree_hash()),
    lists:foreach(
      fun(X) -> ?assertEqual(mth(lists:sublist(?TEST_VECTOR_LEAVES, X)),
                             tree_hash(X - 1)) end,
      random_entries(length(?TEST_VECTOR_LEAVES))).

%% FIXME: Move outside.
old_versions_bigger_test() ->
    LEAVES = [<<X:32>> || X <- lists:seq(0, 64)], % 1024 is not unreasonable
    test_init(LEAVES),
    ?assertEqual(mth(LEAVES), tree_hash()),
    lists:foreach(
      fun(X) -> ?assertEqual(mth(lists:sublist(LEAVES, X)),
                             tree_hash(X - 1)) end,
      random_entries(length(LEAVES))).

%% Test vector from Googles C++ implementation, "Generated from
%% ReferenceMerklePath."
-define(TEST_VECTOR_PATHS,
        %% {leaf_index+1, version+1, path}
        [{0, 0, []},
         {1, 1, []},
         {1, 8,
          ["96a296d224f285c67bee93c30f8a309157f0daa35dc5b87e410b78630a09cfc7",
           "5f083f0a1a33ca076a95279832580db3e0ef4584bdff1f54c8a360f50de3031e",
           "6b47aaf29ee3c2af9af889bc1fb9254dabd31177f16232dd6aab035ca39bf6e4"]},
         {6, 8,
          ["bc1a0643b12e4d2d7c77918f44e0f4f79a838b6cf9ec5b5c283e1f4d88599e6b",
           "ca854ea128ed050b41b35ffc1b87b8eb2bde461e9e3b5596ece6b9d5975a0ae0",
           "d37ee418976dd95753c1c73862b9398fa2a2cf9b4ff0fdfe8b30cd95209614b7"]},
         {3, 3,
          ["fac54203e7cc696cf0dfcb42c92a1d9dbaf70ad9e621f4bd8d98662f00e3c125"]},
         {2, 5,
          ["6e340b9cffb37a989ca544e6bb780a2c78901d3fb33738768511a30617afa01d",
           "5f083f0a1a33ca076a95279832580db3e0ef4584bdff1f54c8a360f50de3031e",
           "bc1a0643b12e4d2d7c77918f44e0f4f79a838b6cf9ec5b5c283e1f4d88599e6b"]}]).

%% @doc Test paths on a single version 7 tree.
%% FIXME: Move outside.
path_test() ->
    test_init(?TEST_VECTOR_LEAVES),
    foreach(
      fun(N) ->
              Test = lists:nth(N, ?TEST_VECTOR_PATHS),
              ?assertEqual(
                 path_ref(element(1, Test) - 1,
                          lists:sublist(?TEST_VECTOR_LEAVES, element(2, Test))),
                 path(element(1, Test) - 1, element(2, Test) - 1))
      end,
      lists:seq(1, length(?TEST_VECTOR_PATHS))).

%% @doc Test path on minimal sized trees.
%% FIXME: Move outside.
path_inc_test() ->
    foreach(
      fun(N) ->
              Test = lists:nth(N, ?TEST_VECTOR_PATHS),
              Leaves = lists:sublist(?TEST_VECTOR_LEAVES, element(2, Test)),
              test_init(Leaves),
              ?assertEqual(
                 path_ref(element(1, Test) - 1, Leaves),
                 path(element(1, Test) - 1, element(2, Test) - 1))
      end,
      lists:seq(1, length(?TEST_VECTOR_PATHS))).

%%%%%%%%%%%%%%%%%%%%
%% Test helpers.
random_entries(N) ->
    [V || {_, V} <- lists:sort(
                      [{random:uniform(N), E} || E <- lists:seq(1, N)])].

%% @doc Return the Merkle Tree Head for the leaves in L. Reference
%% implementation for testing. Implements the algorithm in section 2.1
%% of RFC 6962.
-spec mth(list()) -> binary().
mth([]) ->
    hash(<<"">>);
mth([E]) ->
    hash([<<"\x00">>, E]);
mth(L) ->
    Split = 1 bsl (bitpos_first_set(length(L) - 1) - 1),
    {L1, L2} = lists:split(Split, L),
    hash([<<"\x01">>, mth(L1), mth(L2)]).

%% @doc Return the Merkle Audit Path from I to the root of the tree
%% with leaves L. Reference implementation for testing. Implements the
%% algorithm in section 2.1.1 of RFC 6962.
-spec path_ref(non_neg_integer(), list()) -> list().
path_ref(I, _) when I < 0 ->
    [];
path_ref(I, L) when I >= length(L) ->
    [];
path_ref(0, [_]) ->
    [];
path_ref(I, L) ->
    Split = 1 bsl (bitpos_first_set(length(L) - 1) - 1),
    {L1, L2} = lists:split(Split, L),
    case I of
        I when I < Split ->
            path_ref(I, L1) ++ [mth(L2)];
        _ ->
            path_ref(I - Split, L2) ++ [mth(L1)]
    end.

%%%%%%%%%%%%%%%%%%%%
%% Testing the test helpers. It's turtles all the way down.
-define(TEST_VECTOR_HASHES,
        ["6e340b9cffb37a989ca544e6bb780a2c78901d3fb33738768511a30617afa01d",
         "fac54203e7cc696cf0dfcb42c92a1d9dbaf70ad9e621f4bd8d98662f00e3c125",
         "aeb6bcfe274b70a14fb067a5e5578264db0fa9b51af5e0ba159158f329e06e77",
         "d37ee418976dd95753c1c73862b9398fa2a2cf9b4ff0fdfe8b30cd95209614b7",
         "4e3bbb1f7b478dcfe71fb631631519a3bca12c9aefca1612bfce4c13a86264d4",
         "76e67dadbcdf1e10e1b74ddc608abd2f98dfb16fbce75277b5232a127f2087ef",
         "ddb89be403809e325750d3d263cd78929c2942b7942a34b77e122c9594a74c8c",
         "5dc9da79a70659a9ad559cb701ded9a2ab9d823aad2f4960cfe370eff4604328"]).
mth_test() ->
    lists:foreach(
      fun(X) -> ?assertEqual(
		   hex:hexstr_to_bin(lists:nth(X, ?TEST_VECTOR_HASHES)),
                   mth(lists:sublist(?TEST_VECTOR_LEAVES, X)))
      end,
      lists:seq(1, length(?TEST_VECTOR_LEAVES))).

path_ref_test() ->
    foreach(
      fun(N) ->
              Test = lists:nth(N, ?TEST_VECTOR_PATHS),
              ?assertEqual(
                 [hex:hexstr_to_bin(X) || X <- element(3, Test)],
                 path_ref(element(1, Test) - 1,
                          lists:sublist(?TEST_VECTOR_LEAVES, element(2, Test))))
      end,
      lists:seq(1, length(?TEST_VECTOR_PATHS))).