Adds grpc/cli support for preauthkey tags

preauth_keys.go

@@ -29,17 +29,26 @@ type PreAuthKey struct {
 	Reusable    bool
 	Ephemeral   bool `gorm:"default:false"`
 	Used        bool `gorm:"default:false"`
+	AclTags     []PreAuthKeyAclTag
 
 	CreatedAt  *time.Time
 	Expiration *time.Time
 }
 
+// PreAuthKeyAclTag describes an autmatic tag applied to a node when registered with the associated PreAuthKey
+type PreAuthKeyAclTag struct {
+	ID           uint64 `gorm:"primary_key"`
+	PreAuthKeyID uint64
+	Tag          string
+}
+
 // CreatePreAuthKey creates a new PreAuthKey in a namespace, and returns it.
 func (h *Headscale) CreatePreAuthKey(
 	namespaceName string,
 	reusable bool,
 	ephemeral bool,
 	expiration *time.Time,
+	aclTags []string,
 ) (*PreAuthKey, error) {
 	namespace, err := h.GetNamespace(namespaceName)
 	if err != nil {
@@ -62,8 +71,26 @@ func (h *Headscale) CreatePreAuthKey(
 		Expiration:  expiration,
 	}
 
-	if err := h.db.Save(&key).Error; err != nil {
-		return nil, fmt.Errorf("failed to create key in the database: %w", err)
+	err = h.db.Transaction(func(db *gorm.DB) error {
+		if err := db.Save(&key).Error; err != nil {
+			return fmt.Errorf("failed to create key in the database: %w", err)
+		}
+
+		if len(aclTags) > 0 {
+			for _, tag := range aclTags {
+				if err := db.Save(&PreAuthKeyAclTag{PreAuthKeyID: key.ID, Tag: tag}).Error; err != nil {
+					return fmt.Errorf(
+						"failed to create key tag in the database: %w",
+						err,
+					)
+				}
+			}
+		}
+		return nil
+	})
+
+	if err != nil {
+		return nil, err
 	}
 
 	return &key, nil
@@ -77,7 +104,7 @@ func (h *Headscale) ListPreAuthKeys(namespaceName string) ([]PreAuthKey, error)
 	}
 
 	keys := []PreAuthKey{}
-	if err := h.db.Preload("Namespace").Where(&PreAuthKey{NamespaceID: namespace.ID}).Find(&keys).Error; err != nil {
+	if err := h.db.Preload("Namespace").Preload("AclTags").Where(&PreAuthKey{NamespaceID: namespace.ID}).Find(&keys).Error; err != nil {
 		return nil, err
 	}
 
@@ -101,11 +128,17 @@ func (h *Headscale) GetPreAuthKey(namespace string, key string) (*PreAuthKey, er
 // DestroyPreAuthKey destroys a preauthkey. Returns error if the PreAuthKey
 // does not exist.
 func (h *Headscale) DestroyPreAuthKey(pak PreAuthKey) error {
-	if result := h.db.Unscoped().Delete(pak); result.Error != nil {
-		return result.Error
-	}
+	return h.db.Transaction(func(db *gorm.DB) error {
+		if result := db.Unscoped().Delete(PreAuthKeyAclTag{PreAuthKeyID: pak.ID}); result.Error != nil {
+			return result.Error
+		}
 
-	return nil
+		if result := db.Unscoped().Delete(pak); result.Error != nil {
+			return result.Error
+		}
+
+		return nil
+	})
 }
 
 // MarkExpirePreAuthKey marks a PreAuthKey as expired.
@@ -131,7 +164,7 @@ func (h *Headscale) UsePreAuthKey(k *PreAuthKey) error {
 // If returns no error and a PreAuthKey, it can be used.
 func (h *Headscale) checkKeyValidity(k string) (*PreAuthKey, error) {
 	pak := PreAuthKey{}
-	if result := h.db.Preload("Namespace").First(&pak, "key = ?", k); errors.Is(
+	if result := h.db.Preload("Namespace").Preload("AclTags").First(&pak, "key = ?", k); errors.Is(
 		result.Error,
 		gorm.ErrRecordNotFound,
 	) {
@@ -176,6 +209,7 @@ func (key *PreAuthKey) toProto() *v1.PreAuthKey {
 		Ephemeral: key.Ephemeral,
 		Reusable:  key.Reusable,
 		Used:      key.Used,
+		AclTags:   make([]string, len(key.AclTags)),
 	}
 
 	if key.Expiration != nil {
@@ -186,5 +220,11 @@ func (key *PreAuthKey) toProto() *v1.PreAuthKey {
 		protoKey.CreatedAt = timestamppb.New(*key.CreatedAt)
 	}
 
+	if len(key.AclTags) > 0 {
+		for idx := range key.AclTags {
+			protoKey.AclTags[idx] = key.AclTags[0].Tag
+		}
+	}
+
 	return &protoKey
 }