1 files changed, 40 insertions, 18 deletions
diff --git a/shard/lib/net/manager.ex b/shard/lib/net/manager.ex
index 17d6e06..75307ee 100644
--- a/shard/lib/net/manager.ex
+++ b/shard/lib/net/manager.ex
@@ -22,8 +22,8 @@ defmodule SNet.Manager do
     {:ok, nil} 
   end
 
-  def handle_call({:add_peer, peer_info}, _from, state) do
-    pid = add_peer_internal(peer_info)
+  def handle_call({:add_peer, peer_info, auth}, _from, state) do
+    pid = add_peer_internal(peer_info, auth)
     {:reply, pid, state}
   end
 
@@ -34,8 +34,8 @@ defmodule SNet.Manager do
   end
 
   def handle_call({:peer_up, pid, peer_info, auth}, _from, state) do
-    case :ets.match(:connections,  {peer_info, :_, auth}) do
-      [{_, pid2, _}] when pid2 != pid ->
+    case :ets.match(:connections,  {peer_info, :'$1', auth}) do
+      [[pid2]|_] when pid2 != pid ->
         {:reply, :redundant, state}
       _ ->
         :ets.insert(:connections, {peer_info, pid, auth})
@@ -48,8 +48,8 @@ defmodule SNet.Manager do
     end
   end
 
-  def handle_cast({:connect_and_send, peer_info, msg}, state) do
-    pid = add_peer_internal(peer_info)
+  def handle_cast({:connect_and_send, peer_info, auth, msg}, state) do
+    pid = add_peer_internal(peer_info, auth)
     GenServer.cast(pid, {:send_msg, msg})
     {:noreply, state}
   end
@@ -59,15 +59,18 @@ defmodule SNet.Manager do
     {:noreply, state}
   end
 
-  defp add_peer_internal(peer_info) do
-    case :ets.lookup(:connections, peer_info) do
-      [{_, pid, _}|_] ->
-        pid
-      [] ->
-        my_port = Application.get_env(:shard, :port)
-        {:ok, pid} = SNet.TCPConn.start_link(%{connect_to: peer_info, my_port: my_port, auth: nil})
-        :ets.insert(:connections, {peer_info, pid, nil})
-        pid
+  defp add_peer_internal(peer_info, auth) do
+    if SNet.Addr.is_local? peer_info do
+      nil
+    else
+      case :ets.match(:connections, {peer_info, :'$1', (if auth != nil do auth else :_ end)}) do
+        [[pid]|_] -> pid
+        [] -> 
+          my_port = Application.get_env(:shard, :port)
+          {:ok, pid} = SNet.TCPConn.start_link(%{connect_to: peer_info, my_port: my_port, auth: auth})
+          :ets.insert(:connections, {peer_info, pid, nil})
+          pid
+      end
     end
   end
 
@@ -78,8 +81,8 @@ defmodule SNet.Manager do
   @doc"""
   Connect to a peer specified by ip address and port
   """
-  def add_peer(peer_info) do
-    GenServer.call(__MODULE__, {:add_peer, peer_info})
+  def add_peer(peer_info, auth \\ nil) do
+    GenServer.call(__MODULE__, {:add_peer, peer_info, auth})
   end
 
   @doc"""
@@ -97,6 +100,16 @@ defmodule SNet.Manager do
   end
 
   @doc"""
+  Return the list of connections to a given peer that match a given auth spec
+  """
+  def get_auth_connections_to(peer_info, my_auth, his_auth) do
+    for {^peer_info, pid, %SNet.Auth{my_pk: my_pk, his_pk: his_pk}} <- :ets.lookup(:connections, peer_info),
+      my_pk == my_auth or my_pk in my_auth,
+      his_pk == his_auth or his_pk in his_auth,
+    do: pid
+  end
+
+  @doc"""
   Send message to a peer specified by peer info.
   Opens a connection if necessary.
   """
@@ -105,7 +118,16 @@ defmodule SNet.Manager do
       [{^peer_info, pid, _auth}|_] ->
         GenServer.cast(pid, {:send_msg, msg})
       [] ->
-        GenServer.cast(__MODULE__, {:connect_and_send, peer_info, msg})
+        GenServer.cast(__MODULE__, {:connect_and_send, peer_info, nil, msg})
+    end
+  end
+
+  def send_auth(peer_info, auth, msg) do
+    case :ets.match(:connections, {peer_info, :'$1', auth}) do
+      [[pid]|_] ->
+        GenServer.cast(pid, {:send_msg, msg})
+      [] ->
+        GenServer.cast(__MODULE__, {:connect_and_send, peer_info, auth, msg})
     end
   end