3 files changed, 23 insertions, 8 deletions
diff --git a/shard/lib/net/group.ex b/shard/lib/net/group.ex
index f09d174..7086c2d 100644
--- a/shard/lib/net/group.ex
+++ b/shard/lib/net/group.ex
@@ -33,15 +33,18 @@ defmodule SNet.PubShardGroup do
   defstruct [:id]
 
   defimpl SNet.Group do
-    def init_lookup(%SNet.PubShardGroup{id: id}, _notify_to) do
+    def init_lookup(%SNet.PubShardGroup{id: id}, notify_to) do
       # For now: ask all currently connected peers and connect to new peers we know of
       spawn fn ->
         for {_, pid, _} <- SNet.Manager.list_connections do
-          GenServer.cast(pid, {:send_msg, {:interested, [id]}})
+          GenServer.cast(notify_to, {:peer_connected, pid})
         end
         for peer_info <- Shard.Manager.get_shard_peers id do
           if SNet.Manager.get_connections_to peer_info == [] do
-            SNet.Manager.add_peer(peer_info)      # TODO callback when connected
+            SNet.Manager.add_peer(peer_info,
+              callback: fn pid ->
+                GenServer.cast(notify_to, {:peer_connected, pid})
+              end)
           end
         end
       end
@@ -96,8 +99,11 @@ defmodule SNet.PrivGroup do
           info = GenServer.call(pid, :get_info)
           if Map.has_key?(info, :peer_info) do
             for pi <- info.peer_info do
-              SNet.Manager.add_peer(pi, %SNet.Auth{my_pk: my_pk, his_pk: pk})
-              # no callback here, we don't know if connect was successful
+              SNet.Manager.add_peer(pi,
+                auth: %SNet.Auth{my_pk: my_pk, his_pk: pk},
+                callback: fn pid ->
+                  GenServer.cast(notify_to, {:peer_connected, pid})
+                end)
             end
           end
         end
diff --git a/shard/lib/net/manager.ex b/shard/lib/net/manager.ex
index 75307ee..624d7e6 100644
--- a/shard/lib/net/manager.ex
+++ b/shard/lib/net/manager.ex
@@ -22,8 +22,11 @@ defmodule SNet.Manager do
     {:ok, nil} 
   end
 
-  def handle_call({:add_peer, peer_info, auth}, _from, state) do
+  def handle_call({:add_peer, peer_info, auth, callback}, _from, state) do
     pid = add_peer_internal(peer_info, auth)
+    if callback != nil do
+      GenServer.cast(pid, {:callback, callback})
+    end
     {:reply, pid, state}
   end
 
@@ -38,6 +41,7 @@ defmodule SNet.Manager do
       [[pid2]|_] when pid2 != pid ->
         {:reply, :redundant, state}
       _ ->
+        :ets.match_delete(:connections, {peer_info, pid, :_})
         :ets.insert(:connections, {peer_info, pid, auth})
 
         # Send interested message for all our shards
@@ -81,8 +85,8 @@ defmodule SNet.Manager do
   @doc"""
   Connect to a peer specified by ip address and port
   """
-  def add_peer(peer_info, auth \\ nil) do
-    GenServer.call(__MODULE__, {:add_peer, peer_info, auth})
+  def add_peer(peer_info, opts \\ []) do
+    GenServer.call(__MODULE__, {:add_peer, peer_info, opts[:auth], opts[:callback]})
   end
 
   @doc"""
diff --git a/shard/lib/net/tcpconn.ex b/shard/lib/net/tcpconn.ex
index 25dc839..bd169aa 100644
--- a/shard/lib/net/tcpconn.ex
+++ b/shard/lib/net/tcpconn.ex
@@ -284,6 +284,11 @@ defmodule SNet.TCPConn do
     end
   end
 
+  def handle_cast({:callback, cb}, state) do
+    cb.(self())
+    {:noreply, state}
+  end
+
   def handle_cast({:send_msg, msg}, state) do
     msgbin = :erlang.term_to_binary msg
     enc = :enacl.secretbox(msgbin, state.nonce_send, state.secret_send)