Set tags as part of handleAuthKeyCommon

protocol_common.go

@@ -345,6 +345,7 @@ func (h *Headscale) handleAuthKeyCommon(
 		machine.NodeKey = nodeKey
 		machine.AuthKeyID = uint(pak.ID)
 		err := h.RefreshMachine(machine, registerRequest.Expiry)
+
 		if err != nil {
 			log.Error().
 				Caller().
@@ -355,6 +356,25 @@ func (h *Headscale) handleAuthKeyCommon(
 
 			return
 		}
+
+		aclTags := pak.toProto().AclTags
+		if len(aclTags) > 0 {
+			// This conditional preserves the existing behaviour, although SaaS would reset the tags on auth-key login
+			err = h.SetTags(machine, aclTags)
+		}
+
+		if err != nil {
+			log.Error().
+				Caller().
+				Bool("noise", machineKey.IsZero()).
+				Str("machine", machine.Hostname).
+				Strs("aclTags", aclTags).
+				Err(err).
+				Msg("Failed to set tags after refreshing machine")
+
+			return
+		}
+
 	} else {
 		now := time.Now().UTC()
 
@@ -380,6 +400,7 @@ func (h *Headscale) handleAuthKeyCommon(
 			NodeKey:        nodeKey,
 			LastSeen:       &now,
 			AuthKeyID:      uint(pak.ID),
+			ForcedTags:     pak.toProto().AclTags,
 		}
 
 		machine, err = h.RegisterMachine(