From 7dffa98884162b4d1ce8fd45ce9abaa5693274f6 Mon Sep 17 00:00:00 2001
From: abhishek9686 <abhi281342@gmail.com>
Date: Sun, 22 Sep 2024 14:06:17 +0400
Subject: [PATCH] associate enrollment key tags to node

---
 auth/host_session.go          | 12 ++++++++++--
 controllers/enrollmentkeys.go | 11 ++---------
 logic/nodes.go                |  3 +++
 3 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/auth/host_session.go b/auth/host_session.go
index 62e9d438..7a392924 100644
--- a/auth/host_session.go
+++ b/auth/host_session.go
@@ -222,7 +222,7 @@ func SessionHandler(conn *websocket.Conn) {
 		if err = conn.WriteMessage(messageType, reponseData); err != nil {
 			logger.Log(0, "error during message writing:", err.Error())
 		}
-		go CheckNetRegAndHostUpdate(netsToAdd[:], &result.Host, uuid.Nil)
+		go CheckNetRegAndHostUpdate(netsToAdd[:], &result.Host, uuid.Nil, []models.TagID{})
 	case <-timeout: // the read from req.answerCh has timed out
 		logger.Log(0, "timeout signal recv,exiting oauth socket conn")
 		break
@@ -236,7 +236,7 @@ func SessionHandler(conn *websocket.Conn) {
 }
 
 // CheckNetRegAndHostUpdate - run through networks and send a host update
-func CheckNetRegAndHostUpdate(networks []string, h *models.Host, relayNodeId uuid.UUID) {
+func CheckNetRegAndHostUpdate(networks []string, h *models.Host, relayNodeId uuid.UUID, tags []models.TagID) {
 	// publish host update through MQ
 	for i := range networks {
 		network := networks[i]
@@ -246,6 +246,14 @@ func CheckNetRegAndHostUpdate(networks []string, h *models.Host, relayNodeId uui
 				logger.Log(0, "failed to add host to network:", h.ID.String(), h.Name, network, err.Error())
 				continue
 			}
+			if len(tags) > 0 {
+				newNode.Tags = make(map[models.TagID]struct{})
+				for _, tagI := range tags {
+					newNode.Tags[tagI] = struct{}{}
+				}
+				logic.UpsertNode(newNode)
+			}
+
 			if relayNodeId != uuid.Nil && !newNode.IsRelayed {
 				// check if relay node exists and acting as relay
 				relaynode, err := logic.GetNodeByID(relayNodeId.String())
diff --git a/controllers/enrollmentkeys.go b/controllers/enrollmentkeys.go
index 4ae31e68..9d7fbe43 100644
--- a/controllers/enrollmentkeys.go
+++ b/controllers/enrollmentkeys.go
@@ -308,10 +308,7 @@ func handleHostRegister(w http.ResponseWriter, r *http.Request) {
 				return
 			}
 		}
-		newHost.Tags = make(map[models.TagID]struct{})
-		for _, tagI := range enrollmentKey.Groups {
-			newHost.Tags[tagI] = struct{}{}
-		}
+
 		if err = logic.CreateHost(&newHost); err != nil {
 			logger.Log(
 				0,
@@ -342,10 +339,6 @@ func handleHostRegister(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 		logic.UpdateHostFromClient(&newHost, currHost)
-		currHost.Tags = make(map[models.TagID]struct{})
-		for _, tagI := range enrollmentKey.Groups {
-			currHost.Tags[tagI] = struct{}{}
-		}
 		err = logic.UpsertHost(currHost)
 		if err != nil {
 			slog.Error("failed to update host", "id", currHost.ID, "error", err)
@@ -364,5 +357,5 @@ func handleHostRegister(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusOK)
 	json.NewEncoder(w).Encode(&response)
 	// notify host of changes, peer and node updates
-	go auth.CheckNetRegAndHostUpdate(enrollmentKey.Networks, &newHost, enrollmentKey.Relay)
+	go auth.CheckNetRegAndHostUpdate(enrollmentKey.Networks, &newHost, enrollmentKey.Relay, enrollmentKey.Groups)
 }
diff --git a/logic/nodes.go b/logic/nodes.go
index 8cb451fe..af59f6e9 100644
--- a/logic/nodes.go
+++ b/logic/nodes.go
@@ -420,6 +420,9 @@ func SetNodeDefaults(node *models.Node, resetConnected bool) {
 		node.SetDefaultConnected()
 	}
 	node.SetExpirationDateTime()
+	if node.Tags == nil {
+		node.Tags = make(map[models.TagID]struct{})
+	}
 }
 
 // GetRecordKey - get record key