overlay: avoid one node joined multi times

cmd/overlay/overlay.go

@@ -58,6 +58,10 @@ func (c *Control) Run() error {
 		switch cmd.Action {
 		case "error":
 			logrus.Error(cmd.Data)
+		case "exit":
+			logrus.Error(cmd.Data)
+			c.Close()
+			return nil
 		case "connected":
 			c.latency = time.Since(c.latencyT1)
 			logrus.WithField("latency", c.latency).Info("started as an overlay node now")
@@ -69,6 +73,9 @@ func (c *Control) Run() error {
 }
 
 func (c *Control) Close() error {
+	c.controlConn.WriteControl(websocket.CloseMessage,
+		websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
+		time.Now().Add(time.Second))
 	return c.controlConn.Close()
 }
 

server/overlay/overlay.go

@@ -8,7 +8,6 @@ import (
 	"net"
 	"strings"
 	"sync"
-	"time"
 
 	"github.com/gorilla/websocket"
 	"github.com/rkonfj/toh/server/id"
@@ -36,6 +35,7 @@ type Session struct {
 }
 
 type Node struct {
+	id       string
 	key      string
 	publicIP string
 	router   *OverlayRouter
@@ -115,7 +115,6 @@ func (n *Node) Relay(sessionID, nonce string, data *websocket.Conn) {
 }
 
 func (n *Node) Close() error {
-	n.control.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))
 	n.control.Close()
 	return nil
 }
@@ -128,7 +127,7 @@ func (n *Node) Session(sessionID string) *Session {
 
 func (n *Node) ID() string {
 	if len(n.publicIP) > 0 {
-		return n.publicIP
+		return fmt.Sprintf("%s_%s", n.publicIP, n.id)
 	}
 	return n.control.RemoteAddr().String()
 }
@@ -184,18 +183,23 @@ func NewOverlayRouter() *OverlayRouter {
 	}
 }
 
-func (r *OverlayRouter) RegisterNode(key, nodeIP string, wsConn *websocket.Conn) {
+func (r *OverlayRouter) RegisterNode(key, nodeIP string, wsConn *websocket.Conn) error {
 	r.mut.Lock()
 	defer r.mut.Unlock()
+	if _, ok := r.connectedNodes[key]; ok {
+		return errors.New("node already joined the overlay network")
+	}
 	logrus.WithField("node", key).Debug("overlay node connected")
 	r.connectedNodes[key] = &Node{
+		id:       id.Generate(0),
 		key:      key,
 		publicIP: nodeIP,
 		router:   r,
 		control:  wsConn,
 		sessions: make(map[string]*Session),
 	}
 	go r.connectedNodes[key].runControlLoop()
+	return nil
 }
 
 func (r *OverlayRouter) UnregisterNode(key string) {

server/server.go

@@ -122,7 +122,11 @@ func (s *TohServer) handleOverlay(w http.ResponseWriter, r *http.Request) {
 		logrus.Error(err)
 		return
 	}
-	s.overlayRouter.RegisterNode(key, spec.RealIP(r), wsConn)
+	err = s.overlayRouter.RegisterNode(key, spec.RealIP(r), wsConn)
+	if err != nil {
+		wsConn.WriteJSON(overlay.ControlCommand{Action: "exit", Data: err.Error()})
+		return
+	}
 	wsConn.WriteJSON(overlay.ControlCommand{Action: "connected"})
 }