From 874db799f872e869ffc50f90f45ef4315e1c0ea4 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Mon, 30 Aug 2021 18:57:11 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=81=E4=B8=9A=E8=AE=A4=E8=AF=81=E4=BF=A1?=
 =?UTF-8?q?=E6=81=AF=E4=B8=AD=E5=A2=9E=E5=8A=A0=E8=8A=82=E7=82=B9=E6=95=B0?=
 =?UTF-8?q?=E9=99=90=E5=88=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/const/vars.go         |  3 ++-
 internal/db/models/node_dao.go | 14 ++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/internal/const/vars.go b/internal/const/vars.go
index 8ba2cbbf..efac89ed 100644
--- a/internal/const/vars.go
+++ b/internal/const/vars.go
@@ -3,5 +3,6 @@
 package teaconst
 
 var (
-	IsPlus = false
+	IsPlus         = false
+	MaxNodes int32 = 0
 )
diff --git a/internal/db/models/node_dao.go b/internal/db/models/node_dao.go
index 0bc44041..c4569be8 100644
--- a/internal/db/models/node_dao.go
+++ b/internal/db/models/node_dao.go
@@ -2,6 +2,7 @@ package models
 
 import (
 	"encoding/json"
+	teaconst "github.com/TeaOSLab/EdgeAPI/internal/const"
 	"github.com/TeaOSLab/EdgeAPI/internal/db/models/dns"
 	"github.com/TeaOSLab/EdgeAPI/internal/errors"
 	"github.com/TeaOSLab/EdgeAPI/internal/utils"
@@ -118,6 +119,19 @@ func (this *NodeDAO) FindNodeName(tx *dbs.Tx, id int64) (string, error) {
 
 // CreateNode 创建节点
 func (this *NodeDAO) CreateNode(tx *dbs.Tx, adminId int64, name string, clusterId int64, groupId int64, regionId int64) (nodeId int64, err error) {
+	// 检查节点数量
+	if teaconst.MaxNodes > 0 {
+		count, err := this.Query(tx).
+			State(NodeStateEnabled).
+			Count()
+		if err != nil {
+			return 0, err
+		}
+		if int64(teaconst.MaxNodes) <= count {
+			return 0, errors.New("[企业版]超出最大节点数限制：" + types.String(teaconst.MaxNodes) + "，请购买更多配额")
+		}
+	}
+
 	uniqueId, err := this.GenUniqueId(tx)
 	if err != nil {
 		return 0, err