From 66bbb5bc6a310d22950f62e6d190d3e4671fa89d Mon Sep 17 00:00:00 2001
From: rkonfj <rkonfj@gmail.com>
Date: Sun, 7 Jan 2024 12:34:50 +0800
Subject: [PATCH] overlay: avoid one node joined multi times

---
 cmd/overlay/overlay.go    |  7 +++++++
 server/overlay/overlay.go | 12 ++++++++----
 server/server.go          |  6 +++++-
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/cmd/overlay/overlay.go b/cmd/overlay/overlay.go
index 8eb8b35..9d7ab6f 100644
--- a/cmd/overlay/overlay.go
+++ b/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()
 }
 
diff --git a/server/overlay/overlay.go b/server/overlay/overlay.go
index c50b71a..5ce666a 100644
--- a/server/overlay/overlay.go
+++ b/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,11 +183,15 @@ 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,
@@ -196,6 +199,7 @@ func (r *OverlayRouter) RegisterNode(key, nodeIP string, wsConn *websocket.Conn)
 		sessions: make(map[string]*Session),
 	}
 	go r.connectedNodes[key].runControlLoop()
+	return nil
 }
 
 func (r *OverlayRouter) UnregisterNode(key string) {
diff --git a/server/server.go b/server/server.go
index be10852..67a6749 100644
--- a/server/server.go
+++ b/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"})
 }