Fix crash in inclusion() and add inclusion-and-leaf().

2 files changed, 41 insertions, 16 deletions

diff --git a/src/db.erl b/src/db.erl
index 25cfc5e..91e379e 100644
--- a/src/db.erl
+++ b/src/db.erl
@@ -72,25 +72,26 @@ stop() ->
 %%%%%%%%%%%%%%%%%%%%
 %% Public API.
 
+-spec add(binary()) -> ok.
 add(Entry) ->
     gen_server:call(?MODULE, {add, Entry}).
 
-%% @doc Find entry with entryhash=Hash.
-find(entryhash, Hash) ->
-    gen_server:call(?MODULE, {find, entryhash, Hash});
-%% @doc Find entry with mtlhash=Hash.
-find(mtlhash, Hash) ->
-    gen_server:call(?MODULE, {find, mtlhash, Hash}).
-
-dump(Table) ->
-    gen_server:call(?MODULE, {dump, Table}).
+%% @doc Find one entry.
+-spec find(entryhash | mtlhash | index, binary()) -> ht:mtl().
+find(Type, Hash) ->
+    gen_server:call(?MODULE, {find, Type, Hash}).
 
+-spec get_by_index(non_neg_integer(), non_neg_integer()) -> [ht:mtl()].
 get_by_index(Start, End) ->
     gen_server:call(?MODULE, {get_by_index, {Start, End}}).
 
+-spec get_by_index_sorted(non_neg_integer(), non_neg_integer()) -> [listht:mtl()].
 get_by_index_sorted(Start, End) ->
     gen_server:call(?MODULE, {get_by_index_sorted, {Start, End}}).
 
+%% Testing and debugging.
+dump(Table) ->
+    gen_server:call(?MODULE, {dump, Table}).
 
 %%%%%%%%%%%%%%%%%%%%
 %% gen_server callbacks.
@@ -108,18 +109,19 @@ terminate(_Reason, _State) ->
     io:format("~p terminating~n", [?MODULE]),
     ok.
 
-
 %%%%%%%%%%%%%%%%%%%%
 %% The meat.
 
 handle_call(stop, _From, State) ->
     {stop, normal, stopped, State};
+
 handle_call({add, Entry}, _From, State) ->
     F = fun() ->
                 mnesia:write(Entry)
         end,
     Res = mnesia:transaction(F),
     {reply, Res, State};
+
 handle_call({dump, Table}, _From, State) ->
     F = fun() ->
                 Q = qlc:q([E || E <- mnesia:table(Table)]),
@@ -127,6 +129,7 @@ handle_call({dump, Table}, _From, State) ->
         end,
     Res = mnesia:transaction(F),
     {reply, Res, State};
+
 handle_call({find, entryhash, Hash}, _From, State) ->
     {reply,
      find_entry(fun() -> mnesia:index_read(plop, Hash, #plop.entryhash) end),
@@ -135,9 +138,15 @@ handle_call({find, mtlhash, Hash}, _From, State) ->
     {reply,
      find_entry(fun() -> mnesia:index_read(plop, Hash, #plop.mtlhash) end),
      State};
+handle_call({find, index, Index}, _From, State) ->
+    {reply,
+     find_entry(fun() -> mnesia:read(plop, Index) end),
+     State};
+
 handle_call({get_by_index, {Start, End}}, _From, State) ->
     Res = [X || [_, X] <- select_index(Start, End)],
     {reply, Res, State};
+
 handle_call({get_by_index_sorted, {Start, End}}, _From, State) ->
     %% FIXME: RAM hog -- how bad is it?
     Res = [X || [_, X] <- lists:sort(select_index(Start, End))],
diff --git a/src/plop.erl b/src/plop.erl
index 8277175..abd7d87 100644
--- a/src/plop.erl
+++ b/src/plop.erl
@@ -26,7 +26,7 @@
 %% API.
 -export([start_link/2, stop/0]).
 -export([get_logid/0, serialise/1]).
--export([add/1, sth/0, get/2, consistency/2, inclusion/2]).
+-export([add/1, sth/0, get/2, consistency/2, inclusion/2, inclusion_and_leaf/2]).
 %% API for tests.
 -export([read_keyfile_rsa/2, read_keyfiles_ec/2]).
 -export([testing_get_pubkey/0]).
@@ -113,8 +113,14 @@ get(Start, End) ->
     gen_server:call(?MODULE, {get, {Start, End}}).
 consistency(TreeSizeFirst, TreeSizeSecond) ->
     gen_server:call(?MODULE, {consistency, {TreeSizeFirst, TreeSizeSecond}}).
+-spec inclusion(binary(), non_neg_integer()) ->
+                       {ok|notfound, plop_entry() | string()}.
 inclusion(Hash, TreeSize) ->
     gen_server:call(?MODULE, {inclusion, {Hash, TreeSize}}).
+-spec inclusion_and_leaf(non_neg_integer(), non_neg_integer()) ->
+                       {ok|notfound, plop_entry() | string()}.
+inclusion_and_leaf(Index, TreeSize) ->
+    gen_server:call(?MODULE, {inclusion_and_leaf, {Index, TreeSize}}).
 get_logid() ->
     gen_server:call(?MODULE, {get, logid}).
 testing_get_pubkey() ->
@@ -152,12 +158,22 @@ handle_call({consistency, {First, Second}}, _From, Plop) ->
     {reply, ht:consistency(First - 1, Second - 1), Plop};
 
 handle_call({inclusion, {Hash, TreeSize}}, _From, Plop) ->
-    {Index, Proof} = case db:find(mtlhash, Hash) of
-                         [] -> [];
-                         {plop, I, _EntryHash, _MTLHash, _MTL, _SPT} ->
-                             {I, ht:path(I, TreeSize - 1)}
+    R = case db:find(mtlhash, Hash) of
+            [] ->
+                {notfound, "Unknown hash"};     % FIXME: include Hash
+            {plop, I, _EntryHash, _MTLHash, _MTL, _SPT} ->
+                {ok, I, ht:path(I, TreeSize - 1)}
             end,
-    {reply, {Index, Proof}, Plop};
+    {reply, R, Plop};
+
+handle_call({inclusion_and_leaf, {Index, TreeSize}}, _From, Plop) ->
+    R = case db:find(index, Index) of
+            [] ->
+                {notfound, "Unknown index"};    % FIXME: include Index
+            {plop, I, _EntryHash, _MTLHash, MTL, _SPT} ->
+                {ok, MTL, ht:path(I, TreeSize - 1)}
+            end,
+    {reply, R, Plop};
 
 handle_call({test, pubkey}, _From,
             Plop = #state{pubkey = PK}) ->