8 files changed, 387 insertions, 129 deletions

diff --git a/shard/lib/app/chat.ex b/shard/lib/app/chat.ex
index db2cb64..e61d648 100644
--- a/shard/lib/app/chat.ex
+++ b/shard/lib/app/chat.ex
@@ -28,7 +28,7 @@ defmodule SApp.Chat do
 
   defimpl Shard.Manifest, for: Manifest do
     def start(m) do
-      SApp.Chat.start_link(m.channel)
+      DynamicSupervisor.start_child(Shard.DynamicSupervisor, {SApp.Chat, m.channel})
     end
   end
 
@@ -69,6 +69,12 @@ defmodule SApp.Chat do
     end
   end
 
+  def find_proc(chan) do
+    manifest = %Manifest{channel: chan}
+    id = SData.term_hash manifest
+    Shard.Manager.find_proc id
+  end
+
   @doc """
   Implementation of the :manifest call that returns the chat room's manifest
   """
@@ -98,10 +104,11 @@ defmodule SApp.Chat do
   to send a message to the chat room. Puts the message in the store and syncs
   with all connected peers.
   """
-  def handle_cast({:chat_send, msg}, state) do
-    msgitem = {(System.os_time :seconds),
-               Shard.Identity.get_nickname(),
-               msg}
+  def handle_cast({:chat_send, pk, msg}, state) do
+    msgbin = SData.term_bin {(System.os_time :seconds), msg}
+    {:ok, sign} = Shard.Keys.sign_detached(pk, msgbin)
+    msgitem = {pk, msgbin, sign}
+
     prev_root = state.mst.root
     mst = MST.insert(state.mst, msgitem)
     state = %{state | mst: mst}
@@ -157,21 +164,26 @@ defmodule SApp.Chat do
           state
         else
           # Try adding the message
-          if prev_root == state.mst.root do
-            mst2 = MST.insert(state.mst, msgitem)
-            if mst2.root == new_root do
-              # This was the only message missing, we are happy!
-              state = %{state | mst: mst2}
-              Shard.Manager.save_state(state.id, mst2.root)
-              GenServer.cast(state.page_store, {:set_roots, [mst2.root]})
-              msg_callback(state, msgitem)
-              state
+          {pk, bin, sign} = msgitem
+          if Shard.Keys.verify(pk, bin, sign) == :ok do
+            if prev_root == state.mst.root do
+              mst2 = MST.insert(state.mst, msgitem)
+              if mst2.root == new_root do
+                # This was the only message missing, we are happy!
+                state = %{state | mst: mst2}
+                Shard.Manager.save_state(state.id, mst2.root)
+                GenServer.cast(state.page_store, {:set_roots, [mst2.root]})
+                msg_callback(state, msgitem)
+                state
+              else
+                Logger.warn("Invalid new root after inserting same message item!")
+                state
+              end
             else
-              # More messages are missing, start a full merge
               init_merge(state, new_root, peer_id)
             end
           else
-            init_merge(state, new_root, peer_id)
+            state
           end
         end
       {:root, new_root} ->
@@ -190,21 +202,33 @@ defmodule SApp.Chat do
 
   defp init_merge(state, new_root, source_peer) do
     # TODO: make the merge asynchronous
+    
+    Logger.info("Starting merge for #{inspect state.manifest}, merging root: #{new_root|>Base.encode16}")
+
     prev_last = for {x, true} <- MST.last(state.mst, nil, 100), into: MapSet.new, do: x
 
     mgmst = %{state.mst | root: new_root}
     mgmst = put_in(mgmst.store.prefer_ask, [source_peer])
     mst = MST.merge(state.mst, mgmst)
 
-    for {x, true} <- MST.last(mst, nil, 100) do
+    correct = for {x, true} <- MST.last(mst, nil, 100) do
       if not MapSet.member? prev_last, x do
         msg_callback(state, x)
+        {pk, bin, sign} = x
+        Shard.Keys.verify(pk, bin, sign)
+      else
+        true
       end
     end
 
-    GenServer.cast(state.page_store, {:set_roots, [mst.root]})
-    Shard.Manager.save_state(state.id, mst.root)
-    %{state | mst: mst}
+    if Enum.all? correct do
+      GenServer.cast(state.page_store, {:set_roots, [mst.root]})
+      Shard.Manager.save_state(state.id, mst.root)
+      %{state | mst: mst}
+    else
+      Logger.warn("Incorrect signatures somewhere while merging, dropping merged data")
+      state
+    end
   end
 
   def handle_info({:DOWN, _ref, :process, pid, _reason}, state) do
@@ -212,20 +236,23 @@ defmodule SApp.Chat do
     {:noreply, %{ state | subs: new_subs }}
   end
 
-  defp msg_callback(state, {ts, nick, msg}) do
+  defp msg_callback(state, {pk, msgbin, _sign}) do
+    {ts, msg} = SData.term_unbin msgbin
     for pid <- state.subs do
       if Process.alive?(pid) do
-        send(pid, {:chat_recv, state.channel, {ts, nick, msg}})
+        send(pid, {:chat_recv, state.channel, {ts, pk, msg}})
       end
     end
   end
 
-  defp msg_cmp({ts1, nick1, msg1}, {ts2, nick2, msg2}) do
+  defp msg_cmp({pk1, msgbin1, _sign1}, {pk2, msgbin2, _sign2}) do
+    {ts1, msg1} = SData.term_unbin msgbin1
+    {ts2, msg2} = SData.term_unbin msgbin2
     cond do
       ts1 > ts2 -> :after
       ts1 < ts2 -> :before
-      nick1 > nick2 -> :after
-      nick1 < nick2 -> :before
+      pk1 > pk2 -> :after
+      pk1 < pk2 -> :before
       msg1 > msg2 -> :after
       msg1 < msg2 -> :before
       true -> :duplicate
diff --git a/shard/lib/app/identity.ex b/shard/lib/app/identity.ex
new file mode 100644
index 0000000..7e97897
--- /dev/null
+++ b/shard/lib/app/identity.ex
@@ -0,0 +1,113 @@
+defmodule SApp.Identity do
+  use GenServer
+
+  require Logger
+
+  defmodule Manifest do
+    defstruct [:pk]
+  end
+
+  defmodule State do
+    defstruct [:info, :rev, :signed]
+  end
+
+  defimpl Shard.Manifest, for: Manifest do
+    def start(m) do
+      DynamicSupervisor.start_child(Shard.DynamicSupervisor, {SApp.Identity, m.pk})
+    end
+  end
+
+  def start_link(pk) do
+    GenServer.start_link(__MODULE__, pk)
+  end
+
+  def init(pk) do
+    manifest = %Manifest{pk: pk}
+    id = SData.term_hash manifest
+
+    case Shard.Manager.register(id, manifest, self()) do
+      :ok ->
+        Shard.Manager.dispatch_to(id, nil, self())
+        state = case Shard.Manager.load_state(id) do
+          nil ->
+            info = %{nick: default_nick(pk)}
+            SData.SignRev.new info
+          st ->
+            st
+        end
+        GenServer.cast(self(), :init_pull)
+        {:ok, %{pk: pk, id: id, state: state}}
+      :redundant ->
+        exit(:redundant)
+    end
+  end
+
+  def default_nick(pk) do
+    nick_suffix = pk
+                  |> binary_part(0, 4)
+                  |> Base.encode16
+                  |> String.downcase
+    "Anon" <> nick_suffix
+  end
+
+  def find_proc(pk) do
+    manifest = %Manifest{pk: pk}
+    id = SData.term_hash manifest
+    Shard.Manager.find_proc id
+  end
+
+  def handle_call(:manifest, _from, state) do
+    {:replyl, state.manifest, state}
+  end
+
+  def handle_call(:get_info, _from, state) do
+    {:reply, SData.SignRev.get(state.state), state}
+  end
+
+  def handle_call({:set_info, new_info}, _from, state) do
+    case SData.SignRev.set(state.state, new_info, state.pk) do
+      {:ok, st2} ->
+        Shard.Manager.save_state(state.id, st2)
+        state = put_in(state.state, st2)
+        bcast_state(state)
+        {:reply, :ok, state}
+      err ->
+        {:reply, err, state}
+    end
+  end
+
+  def handle_cast(:init_pull, state) do
+    for {_, pid, _, _} <- Shard.Manager.list_peers do
+      GenServer.cast(pid, {:send_msg, {:interested, [state.id]}})
+    end
+    {:noreply, state}
+  end
+
+  def handle_cast({:interested, peer_id}, state) do
+    Shard.Manager.send(peer_id, {state.id, nil, {:update, SData.SignRev.signed(state.state)}})
+  end
+
+  def handle_cast({:msg, peer_id, _shard_id, nil, msg}, state) do
+    state = case msg do
+      {:update, signed} ->
+        case SData.SignRev.merge(state.state, signed, state.pk) do
+          {true, st2} ->
+            Shard.Manager.save_state(state.id, st2)
+            state = put_in(state.state, st2)
+            bcast_state(state, [peer_id])
+            state
+          {false, _} ->
+            state
+        end
+    end
+    {:noreply, state}
+  end
+
+  def bcast_state(state, exclude \\ []) do
+    for peer_id <- Shard.Manager.get_shard_peers(state.id) do
+      if not Enum.member? exclude, peer_id do
+        Shard.Manager.send(peer_id, {state.id, nil, {:update, SData.SignRev.signed(state.state)}})
+      end
+    end
+  end
+end
diff --git a/shard/lib/application.ex b/shard/lib/application.ex
index 3e3a6ac..9c1577a 100644
--- a/shard/lib/application.ex
+++ b/shard/lib/application.ex
@@ -15,7 +15,7 @@ defmodule Shard.Application do
 
     # Define workers and child supervisors to be supervised
     children = [
-      Shard.Identity,
+      Shard.Keys,
       { DynamicSupervisor, strategy: :one_for_one, name: Shard.DynamicSupervisor },
       
       # Networking
diff --git a/shard/lib/cli/cli.ex b/shard/lib/cli/cli.ex
index bf3a555..3778b2d 100644
--- a/shard/lib/cli/cli.ex
+++ b/shard/lib/cli/cli.ex
@@ -3,22 +3,45 @@ defmodule SCLI do
   Small command line interface for the chat application
   """
 
+  defmodule State do
+    defstruct [:room_pid, :id_pid, :pk]
+  end
+
   def run() do
-    for {_chid, _manifest, chpid} <- Shard.Manager.list_shards do
+    for {_chid, %SApp.Chat.Manifest{}, chpid} <- Shard.Manager.list_shards do
       GenServer.cast(chpid, {:subscribe, self()})
     end
 
-    run(nil)
+    pk = case Shard.Keys.list_identities do
+      [pk1|_] -> pk1
+      _ ->
+        IO.puts "Generating a new identity..."
+        Shard.Keys.new_identity
+    end
+
+    run(%State{room_pid: nil, id_pid: nil, pk: pk})
   end
 
-  defp run(pid) do
+  defp run(state) do
     handle_messages()
 
-    nick = Shard.Identity.get_nickname
-    prompt = case pid do
+    id_pid = case state.id_pid do
+      nil -> SApp.Identity.find_proc(state.pk)
+      x -> x
+    end
+    state = put_in(state.id_pid, id_pid)
+
+    nick = case id_pid do
+      nil -> SApp.Identity.default_nick(state.pk)
+      _ ->
+        info = GenServer.call(id_pid, :get_info)
+        info.nick
+    end
+
+    prompt = case state.room_pid do
       nil -> "(no channel) #{nick}: "
       _ -> 
-        %SApp.Chat.Manifest{channel: chan} = GenServer.call(pid, :manifest)
+        %SApp.Chat.Manifest{channel: chan} = GenServer.call(state.room_pid, :manifest)
         "##{chan} #{nick}: "
     end
 
@@ -28,13 +51,13 @@ defmodule SCLI do
         nil
       String.slice(str, 0..0) == "/" ->
         command = str |> String.slice(1..-1) |> String.split(" ")
-        pid2 = handle_command(pid, command)
-        run(pid2)
+        state = handle_command(state, command)
+        run(state)
       true -> 
         if str != "" do
-          GenServer.cast(pid, {:chat_send, str})
+          GenServer.cast(state.room_pid, {:chat_send, state.pk, str})
         end
-        run(pid)
+        run(state)
     end
   end
 
@@ -50,57 +73,71 @@ defmodule SCLI do
     end
   end
 
-  defp handle_command(pid, ["connect", ipstr, portstr]) do
+  defp handle_command(state, ["connect", ipstr, portstr]) do
     {:ok, ip} = :inet.parse_address (to_charlist ipstr)
     {port, _} = Integer.parse portstr
     Shard.Manager.add_peer(ip, port)
-    pid
+    state
   end
 
-  defp handle_command(pid, ["list"]) do
+  defp handle_command(state, ["list"]) do
     IO.puts "List of known channels:"
 
-    for {_chid, manifest, _chpid} <- Shard.Manager.list_shards do
-      %SApp.Chat.Manifest{channel: chan} = manifest
+    for {_chid, %SApp.Chat.Manifest{channel: chan}, _chpid} <- Shard.Manager.list_shards do
       IO.puts "##{chan}"
     end
-    pid
+    state
   end
 
-  defp handle_command(pid, ["hist"]) do
-    if pid == nil do
+  defp handle_command(state, ["hist"]) do
+    if state.room_pid == nil do
       IO.puts "Not currently on a channel!"
     else
-      GenServer.call(pid, {:read_history, nil, 25})
-      |> Enum.each(fn {{ts, nick, msg}, true} ->
-        IO.puts "#{ts |> DateTime.from_unix! |> DateTime.to_iso8601}  <#{nick}> #{msg}"
+      GenServer.call(state.room_pid, {:read_history, nil, 25})
+      |> Enum.each(fn {{pk, msgbin, _sign}, true} ->
+        {ts, msg} = SData.term_unbin msgbin
+        nick = case SApp.Identity.find_proc pk do
+          nil ->
+            SApp.Identity.default_nick pk
+          pid ->
+            info = GenServer.call(pid, :get_info)
+            info.nick
+        end
+          IO.puts "#{ts |> DateTime.from_unix! |> DateTime.to_iso8601}  <#{nick} #{pk|>binary_part(0, 4)|>Base.encode16|>String.downcase}> #{msg}"
           end)
-      pid
     end
+    state
   end
 
-  defp handle_command(_pid, ["join", qchan]) do
-    list = for {_chid, manifest, chpid} <- Shard.Manager.list_shards,
-               %SApp.Chat.Manifest{channel: chan} = manifest,
-               do: {chan, chpid}
-    case List.keyfind(list, qchan, 0) do
+  defp handle_command(state, ["join", qchan]) do
+    pid = SApp.Chat.find_proc qchan
+    case pid do
       nil ->
-        {:ok, pid} = DynamicSupervisor.start_child(Shard.DynamicSupervisor, {SApp.Chat, qchan})
+        {:ok, pid} = Shard.Manifest.start %SApp.Chat.Manifest{channel: qchan}
         GenServer.cast(pid, {:subscribe, self()})
-        pid
-      {_, pid} ->
+        %{state | room_pid: pid}
+      pid ->
         IO.puts "Switching to ##{qchan}"
-        pid
+        %{state | room_pid: pid}
     end
   end
 
-  defp handle_command(pid, ["nick", nick]) do
-    Shard.Identity.set_nickname nick
-    pid
+  defp handle_command(state, ["nick", nick]) do
+    pid = case state.id_pid do
+      nil -> SApp.Identity.find_proc state.pk
+      x -> x
+    end
+    if pid == nil do
+      IO.puts "Sorry, we have a problem with the identity shard"
+    else
+      info = GenServer.call(pid, :get_info)
+      GenServer.call(pid, {:set_info, %{info | nick: nick}})
+    end
+    state
   end
 
-  defp handle_command(pid, _cmd) do
+  defp handle_command(state, _cmd) do
     IO.puts "Invalid command"
-    pid
+    state
   end
 end
diff --git a/shard/lib/identity.ex b/shard/lib/identity.ex
deleted file mode 100644
index 01b3c9e..0000000
--- a/shard/lib/identity.ex
+++ /dev/null
@@ -1,65 +0,0 @@
-defmodule Shard.Identity do
-  use Agent
-  require Salty.Sign.Ed25519, as: Sign
-  require Logger
-
-  @identity_db [Application.get_env(:shard, :data_path), "identity_db"] |> Path.join |> String.to_atom
-
-  def start_link(_) do
-    Agent.start_link(__MODULE__, :init, [], name: __MODULE__)
-  end
-
-  def init() do
-    :dets.start
-    {:ok, @identity_db} = :dets.open_file @identity_db, type: :set
-
-    case :dets.match @identity_db, :"$1" do
-      [] ->
-        Logger.info "Generating keypair..."
-        {pk, sk} = gen_keypair(Application.get_env(:shard, :peer_id_suffix))
-        nick_suffix = pk
-                      |> binary_part(0, 3)
-                      |> Base.encode16
-                      |> String.downcase
-        nick = "Anon" <> nick_suffix
-        :dets.insert @identity_db, {pk, sk, nick}
-        %{
-          keypair:  {pk, sk},
-          nickname: nick
-        }
-      [[{pk, sk, nick}] | _] ->
-        %{
-          keypair:  {pk, sk},
-          nickname: nick
-        }
-    end
-  end
-
-  defp gen_keypair(suffix, n \\ 0) do
-    {:ok, pk, sk} = Sign.keypair
-    if rem(n, 10000) == 0 do
-      Logger.info "#{n}... expected #{:math.pow(256, byte_size(suffix))}"
-    end
-    if :binary.longest_common_suffix([pk, suffix]) == byte_size(suffix) do
-      {pk, sk}
-    else
-      gen_keypair(suffix, n+1)
-    end
-  end
-
-  def get_keypair() do
-    Agent.get(__MODULE__, &(&1.keypair))
-  end
-
-  def get_nickname() do
-    Agent.get(__MODULE__, &(&1.nickname))
-  end
-
-  def set_nickname(newnick) do
-    Agent.update(__MODULE__, fn state ->
-      {pk, sk} = state.keypair
-      :dets.insert @identity_db, {pk, sk, newnick}
-      %{state | nickname: newnick}
-    end)
-  end
-end
diff --git a/shard/lib/keys.ex b/shard/lib/keys.ex
new file mode 100644
index 0000000..0dc3154
--- /dev/null
+++ b/shard/lib/keys.ex
@@ -0,0 +1,124 @@
+defmodule Shard.Keys do
+  @moduledoc"""
+    Module for saving private keys.
+  """
+
+  use Agent
+  require Salty.Sign.Ed25519, as: Sign
+  require Logger
+
+  @key_db [Application.get_env(:shard, :data_path), "key_db"] |> Path.join |> String.to_atom
+
+  def start_link(_) do
+    Agent.start_link(__MODULE__, :init, [], name: __MODULE__)
+  end
+
+  def init() do
+    :dets.start
+    {:ok, @key_db} = :dets.open_file(@key_db, [type: :set])
+
+    case :dets.lookup(@key_db, :peer) do
+      [] ->
+        Logger.info "Generating peer keypair..."
+        {pk, sk} = gen_keypair(Application.get_env(:shard, :peer_id_suffix))
+        :dets.insert @key_db, {:peer, pk, sk}
+        {pk, sk}
+      [{:peer, pk, sk}] ->
+        {pk, sk}
+    end
+  end
+
+  defp gen_keypair(suffix, n \\ 0) do
+    {:ok, pk, sk} = Sign.keypair
+    if rem(n, 10000) == 0 do
+      Logger.info "#{n}... expected #{:math.pow(256, byte_size(suffix))}"
+    end
+    if check_suffix(pk, suffix) do
+      {pk, sk}
+    else
+      gen_keypair(suffix, n+1)
+    end
+  end
+
+  defp check_suffix(pk, suffix) do
+    :binary.longest_common_suffix([pk, suffix]) == byte_size(suffix)
+  end
+
+  def get_peer_keypair() do
+    Agent.get(__MODULE__, &(&1))
+  end
+
+  @doc"""
+  Generate a new keypair for a user identity, and start an Identity Shard for it.
+  """
+  def new_identity() do
+    {pk, sk} = gen_keypair(Application.get_env(:shard, :identity_suffix))
+    :dets.insert @key_db, {pk, sk}
+    SApp.Identity.start_link(pk)
+    pk
+  end
+
+  @doc"""
+  List the public keys of all identities for which we have a secret key
+  """
+  def list_identities() do
+    for [{pk, _sk}] <- :dets.match(@key_db, :"$1"), do: pk
+  end
+
+  @doc"""
+    Lookup the secret key for a pk and sign a message with it.
+
+    Returns the input value alongside its signature.
+
+    Answer is {:ok, signed} if it worked, or :not_found if we didn't find the key.
+  """
+  def sign(pk, bin) do
+    case :dets.lookup @key_db, pk do
+      [{^pk, sk}] ->
+        Sign.sign(bin, sk)
+      _ -> {:error, :not_found}
+    end
+  end
+
+  @doc"""
+    Checks the signature appended to a signed message corresponds to a public key.
+
+    If correct, returns {:ok, original_message}
+  """
+  def open(pk, signed) do
+    if check_suffix(pk, Application.get_env(:shard, :identity_suffix)) do
+      Sign.open(signed, pk)
+    else
+      {:error, :invalid_pk_suffix}
+    end
+  end
+
+  @doc"""
+    Lookup the secret key for a pk and generate a detached signature for a message.
+
+    The original message is not returned.
+
+    Answer is {:ok, signature} if it worked, or :not_found if we didn't find the key.
+    
+  """
+  def sign_detached(pk, bin) do
+    case :dets.lookup @key_db, pk do
+      [{^pk, sk}] ->
+        Sign.sign_detached(bin, sk)
+      _ -> {:error, :not_found}
+    end
+  end
+
+  @doc"""
+    Verify a detached signature for a message
+
+    Returns :ok if the signature was correct.
+  """
+  def verify(pk, bin, sign) do
+    if check_suffix(pk, Application.get_env(:shard, :identity_suffix)) do
+      Sign.verify_detached(sign, bin, pk)
+    else
+      {:error, :invalid_pk_suffix}
+    end
+  end
+end
diff --git a/shard/lib/manager.ex b/shard/lib/manager.ex
index 57f2371..617378c 100644
--- a/shard/lib/manager.ex
+++ b/shard/lib/manager.ex
@@ -1,5 +1,17 @@
 defprotocol Shard.Manifest do
-  @doc "Start the corresponding Shard process"
+  @moduledoc"""
+  A shard manifest is a data structure that uniquely defines the identity of the shard.
+
+  The hash of the manifest is the unique identifier of that shard on the network.
+  
+  The Manifest protocol is a protocol implemented by the manifest structs for the
+  different shard types. It contains an operation start() that is able to launch the
+  correct process for this shard and connect to other peers that use it.
+  """
+
+  @doc"""
+  Start the corresponding Shard process
+  """
   def start(manifest)
 end
 
@@ -261,6 +273,16 @@ defmodule Shard.Manager do
   end
 
   @doc"""
+  Returns the pid for a shard if it exists
+  """
+  def find_proc(shard_id) do
+    case :dets.lookup(@shard_db, shard_id) do
+      [{^shard_id, _, pid}] -> pid
+      _ -> nil
+    end
+  end
+
+  @doc"""
   Register a process as the handler for shard packets for a given path.
   """
   def dispatch_to(shard_id, path, pid) do
diff --git a/shard/lib/net/tcpconn.ex b/shard/lib/net/tcpconn.ex
index 35f7ea5..543341a 100644
--- a/shard/lib/net/tcpconn.ex
+++ b/shard/lib/net/tcpconn.ex
@@ -20,7 +20,7 @@ defmodule SNet.TCPConn do
   def handle_cast(:handshake, state) do
     socket = state.socket
 
-    {srv_pkey, srv_skey} = Shard.Identity.get_keypair
+    {srv_pkey, srv_skey} = Shard.Keys.get_peer_keypair
     {:ok, sess_pkey, sess_skey} = Box.keypair
     {:ok, challenge} = Salty.Random.buf 32