Просмотр исходного кода

associate enrollment key tags to node

abhishek9686 11 месяцев назад
Родитель
Сommit
7dffa98884
3 измененных файлов с 15 добавлено и 11 удалено
  1. 10 2
      auth/host_session.go
  2. 2 9
      controllers/enrollmentkeys.go
  3. 3 0
      logic/nodes.go

+ 10 - 2
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())

+ 2 - 9
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)
 }

+ 3 - 0
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