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 true bc of
%%% "placeholders", FIXME: doc)
%%% 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([get_incl/2, get_cons/2]).
-export([mktree/1]).
-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}).
get_incl(I, V) ->
    gen_server:call(?MODULE, {inclusion, I, V}).
get_cons(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} = tree_hash(State, State#tree.version),
    {reply, Hash, NewState};
handle_call({tree_hash, Version}, _From, State) ->
    {NewState, Hash} = tree_hash(State, Version),
    {reply, Hash, NewState};
handle_call({inclusion, _Index, _Version}, _From, State) ->
    {reply, nyi, State};
handle_call({consistency, _Version1, _Version2}, _From, State) ->
    {reply, nyi, State}.

%%%%%%%%%%%%%%%%%%%%
%% Private.
-spec get_hash(tree(), non_neg_integer(), tuple()) -> {tree(), binary()}.
get_hash(Tree, Version, IR) ->
    NewTree = update(Tree, Version),
    Hash = ts:retrieve_hash(NewTree#tree.store, IR), % TODO: Honour version!
    %%io:format("DEBUG: get_hash(Tree=~p,~nVersion=~p,~nIR=~p):~nNewTree=~p~n", [Tree, Version, IR, NewTree]),
    {NewTree, Hash}.

-spec tree_hash(tree(), integer()) -> {tree(), binary()}.
tree_hash(Tree, -1) ->
    {Tree, hash(<<"">>)};
tree_hash(Tree, Version) ->
    get_hash(Tree, Version, {0, depth(Tree) - 1}).

%% @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).

%% @doc Calculate hashes in Tree up to and including leaf 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}.

-spec update_layer(tree(), non_neg_integer(), non_neg_integer(),
                   non_neg_integer()) -> tree().
update_layer(Tree, _Layer, _ICur, 0) ->
    Tree;
update_layer(Tree, Layer, ICur, ILast) ->
    NewTree = Tree#tree{store = update_parent(Tree#tree.store, Layer,
                                              strip_bits_bottom(ICur, 1), ILast)},
    %%io:format("DEBUG: update_layer(~p, ~p, ~p, ~p
    update_layer(NewTree, Layer + 1, parent(ICur), parent(ILast)).

-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 ->
    case right_leaf_p(ILast) of
        true -> S;
        _ -> ts:append(S, Layer + 1, ts:retrieve_hash(S, {ILast, Layer}))
    end;
update_parent(S, Layer, I, ILast) ->
    %%io:format("DEBUG: update_parent(R=~p, I=~p, ILast=~p): parent(I)=~p~n", [Layer, I, ILast, parent(I)]),
    S2 = ts:store(S, {parent(I), Layer + 1},
                  mkinnerhash(ts:retrieve_hash(S, {I, Layer}),
                              ts:retrieve_hash(S, {I + 1, Layer}))),
    update_parent(S2, Layer, I + 2, ILast).

parent(I) ->
    I bsr 1.

-spec right_leaf_p(integer()) -> boolean().
right_leaf_p(Index) ->
    case Index band 1 of
        1 -> true;
        _ -> false
    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.
-define(TEST_VECTOR_LEAVES,
        ["", "\x00", "\x10", " !", "01", "@ABC", "PQRSTUVW", "`abcdefghijklmno"]).
-define(TEST_VECTOR_HASHES,
        ["6e340b9cffb37a989ca544e6bb780a2c78901d3fb33738768511a30617afa01d",
         "fac54203e7cc696cf0dfcb42c92a1d9dbaf70ad9e621f4bd8d98662f00e3c125",
         "aeb6bcfe274b70a14fb067a5e5578264db0fa9b51af5e0ba159158f329e06e77",
         "d37ee418976dd95753c1c73862b9398fa2a2cf9b4ff0fdfe8b30cd95209614b7",
         "4e3bbb1f7b478dcfe71fb631631519a3bca12c9aefca1612bfce4c13a86264d4",
         "76e67dadbcdf1e10e1b74ddc608abd2f98dfb16fbce75277b5232a127f2087ef",
         "ddb89be403809e325750d3d263cd78929c2942b7942a34b77e122c9594a74c8c",
         "5dc9da79a70659a9ad559cb701ded9a2ab9d823aad2f4960cfe370eff4604328"]).

%% @doc Build tree using add/2 and mth/2 and compare the resulting
%% tree hashes.
%% FIXME: Don't start and stop the server here!
add_test() ->
    lists:foreach(
      fun(X) -> L = lists:sublist(?TEST_VECTOR_LEAVES, X),
                {ok, _Pid} = start_link(L),
                mktree(L),
                ?assertEqual(mth(L), tree_hash()),
                stop() end,
      random_entries(length(?TEST_VECTOR_LEAVES))).

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

update_test_OFF() ->
    mktree(?TEST_VECTOR_LEAVES),
    lists:foreach(
      fun(Test) ->
              L = lists:sublist(?TEST_VECTOR_LEAVES, Test),
              ?assertEqual(mth(L),
                           tree_hash(lists:nth(Test, ?TEST_VECTOR_HASHES))) end,
      random_entries(length(?TEST_VECTOR_LEAVES))).

%% verify_old_versions_test() ->
%%     {ok, _Pid} = start_link(db:size() - 1).

%%%%%%%%%%%%%%%%%%%%
%% Testing the test helpers.
mth_test() ->
    lists:foreach(
      fun(X) -> ?assertEqual(
		   mth(lists:sublist(?TEST_VECTOR_LEAVES, X)),
		   hex:hexstr_to_bin(lists:nth(X, ?TEST_VECTOR_HASHES)))
      end,
      lists:seq(1, length(?TEST_VECTOR_LEAVES))).

%%%%%%%%%%%%%%%%%%%%
%% Test helpers.

mktree(L) ->
    mktree(new(), L).
mktree(Tree, []) ->
    Tree;
mktree(Tree, [H|T]) ->
    mktree(add(Tree, H), T).

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