实现集群CC防护策略设置

2025-11-03 23:20:26 +08:00 · 2023-05-23 19:16:30 +08:00
parent c7001dfb76
commit aa716c03da
6 changed files with 209 additions and 2 deletions
--- a/internal/db/models/node_cluster_dao.go
+++ b/internal/db/models/node_cluster_dao.go
@@ -996,7 +996,7 @@ func (this *NodeClusterDAO) FindClusterBasicInfo(tx *dbs.Tx, clusterId int64, ca
 	cluster, err := this.Query(tx).
 		Pk(clusterId).
 		State(NodeClusterStateEnabled).
-		Result("id", "name", "timeZone", "nodeMaxThreads", "cachePolicyId", "httpFirewallPolicyId", "autoOpenPorts", "webp", "uam", "httpPages", "isOn", "ddosProtection", "clock", "globalServerConfig", "autoInstallNftables").
+		Result("id", "name", "timeZone", "nodeMaxThreads", "cachePolicyId", "httpFirewallPolicyId", "autoOpenPorts", "webp", "uam", "cc", "httpPages", "isOn", "ddosProtection", "clock", "globalServerConfig", "autoInstallNftables").
 		Find()
 	if err != nil || cluster == nil {
 		return nil, err
@@ -1125,6 +1125,65 @@ func (this *NodeClusterDAO) FindClusterUAMPolicy(tx *dbs.Tx, clusterId int64, ca
 	return policy, nil
 }

+// UpdateClusterHTTPCCPolicy 修改CC策略设置
+func (this *NodeClusterDAO) UpdateClusterHTTPCCPolicy(tx *dbs.Tx, clusterId int64, httpCCPolicy *nodeconfigs.HTTPCCPolicy) error {
+	if httpCCPolicy == nil {
+		err := this.Query(tx).
+			Pk(clusterId).
+			Set("cc", dbs.SQL("null")).
+			UpdateQuickly()
+		if err != nil {
+			return err
+		}
+
+		return this.NotifyHTTPCCUpdate(tx, clusterId)
+	}
+
+	httpCCPolicyJSON, err := json.Marshal(httpCCPolicy)
+	if err != nil {
+		return err
+	}
+	err = this.Query(tx).
+		Pk(clusterId).
+		Set("cc", httpCCPolicyJSON).
+		UpdateQuickly()
+	if err != nil {
+		return err
+	}
+
+	return this.NotifyHTTPCCUpdate(tx, clusterId)
+}
+
+// FindClusterHTTPCCPolicy 查询CC策略设置
+func (this *NodeClusterDAO) FindClusterHTTPCCPolicy(tx *dbs.Tx, clusterId int64, cacheMap *utils.CacheMap) (*nodeconfigs.HTTPCCPolicy, error) {
+	var cacheKey = this.Table + ":FindClusterHTTPCCPolicy:" + types.String(clusterId)
+	if cacheMap != nil {
+		cache, ok := cacheMap.Get(cacheKey)
+		if ok {
+			return cache.(*nodeconfigs.HTTPCCPolicy), nil
+		}
+	}
+
+	httpCCJSON, err := this.Query(tx).
+		Pk(clusterId).
+		Result("cc").
+		FindJSONCol()
+	if err != nil {
+		return nil, err
+	}
+
+	if IsNull(httpCCJSON) {
+		return nodeconfigs.NewHTTPCCPolicy(), nil
+	}
+
+	var policy = nodeconfigs.NewHTTPCCPolicy()
+	err = json.Unmarshal(httpCCJSON, policy)
+	if err != nil {
+		return nil, err
+	}
+	return policy, nil
+}
+
 // UpdateClusterHTTPPagesPolicy 修改自定义页面设置
 func (this *NodeClusterDAO) UpdateClusterHTTPPagesPolicy(tx *dbs.Tx, clusterId int64, httpPagesPolicy *nodeconfigs.HTTPPagesPolicy) error {
 	if httpPagesPolicy == nil {
@@ -1298,6 +1357,11 @@ func (this *NodeClusterDAO) NotifyUAMUpdate(tx *dbs.Tx, clusterId int64) error {
 	return SharedNodeTaskDAO.CreateClusterTask(tx, nodeconfigs.NodeRoleNode, clusterId, 0, 0, NodeTaskTypeUAMPolicyChanged)
 }

+// NotifyHTTPCCUpdate 通知HTTP CC更新
+func (this *NodeClusterDAO) NotifyHTTPCCUpdate(tx *dbs.Tx, clusterId int64) error {
+	return SharedNodeTaskDAO.CreateClusterTask(tx, nodeconfigs.NodeRoleNode, clusterId, 0, 0, NodeTaskTypeHTTPCCPolicyChanged)
+}
+
 // NotifyHTTPPagesPolicyUpdate 通知HTTP Pages更新
 func (this *NodeClusterDAO) NotifyHTTPPagesPolicyUpdate(tx *dbs.Tx, clusterId int64) error {
 	return SharedNodeTaskDAO.CreateClusterTask(tx, nodeconfigs.NodeRoleNode, clusterId, 0, 0, NodeTaskTypeHTTPPagesPolicyChanged)
--- a/internal/db/models/node_cluster_model.go
+++ b/internal/db/models/node_cluster_model.go
@@ -42,6 +42,7 @@ type NodeCluster struct {
 	AutoInstallNftables  bool     `field:"autoInstallNftables"`  // 自动安装nftables
 	IsAD                 bool     `field:"isAD"`                 // 是否为高防集群
 	HttpPages            dbs.JSON `field:"httpPages"`            // 自定义页面设置
+	Cc                   dbs.JSON `field:"cc"`                   // CC设置
 }

 type NodeClusterOperator struct {
@@ -83,6 +84,7 @@ type NodeClusterOperator struct {
 	AutoInstallNftables  any // 自动安装nftables
 	IsAD                 any // 是否为高防集群
 	HttpPages            any // 自定义页面设置
+	Cc                   any // CC设置
 }

 func NewNodeClusterOperator() *NodeClusterOperator {
--- a/internal/db/models/node_dao.go
+++ b/internal/db/models/node_dao.go
@@ -1085,6 +1085,7 @@ func (this *NodeDAO) ComposeNodeConfig(tx *dbs.Tx, nodeId int64, dataMap *shared
 	var clusterIndex = 0
 	config.WebPImagePolicies = map[int64]*nodeconfigs.WebPImagePolicy{}
 	config.UAMPolicies = map[int64]*nodeconfigs.UAMPolicy{}
+	config.HTTPCCPolicies = map[int64]*nodeconfigs.HTTPCCPolicy{}
 	config.HTTPPagesPolicies = map[int64]*nodeconfigs.HTTPPagesPolicy{}
 	var allowIPMaps = map[string]bool{}
 	for _, clusterId := range clusterIds {
@@ -1179,7 +1180,25 @@ func (this *NodeDAO) ComposeNodeConfig(tx *dbs.Tx, nodeId int64, dataMap *shared
 			config.UAMPolicies[clusterId] = uamPolicy
 		}

-		// HTTP Pages
+		// HTTP CC Policy
+		if IsNotNull(nodeCluster.Cc) {
+			var ccPolicy = nodeconfigs.NewHTTPCCPolicy()
+			err = json.Unmarshal(nodeCluster.Cc, ccPolicy)
+			if err != nil {
+				return nil, err
+			}
+
+			// 集成默认设置
+			for i := 0; i < len(serverconfigs.DefaultHTTPCCThresholds); i ++ {
+				if i < len(ccPolicy.Thresholds) {
+					ccPolicy.Thresholds[i].MergeIfEmpty(serverconfigs.DefaultHTTPCCThresholds[i])
+				}
+			}
+
+			config.HTTPCCPolicies[clusterId] = ccPolicy
+		}
+
+		// HTTP Pages Policy
 		if IsNotNull(nodeCluster.HttpPages) {
 			var httpPagesPolicy = nodeconfigs.NewHTTPPagesPolicy()
 			err = json.Unmarshal(nodeCluster.HttpPages, httpPagesPolicy)
--- a/internal/db/models/node_task_dao.go
+++ b/internal/db/models/node_task_dao.go
@@ -26,6 +26,7 @@ const (
 	NodeTaskTypeUserServersStateChanged   NodeTaskType = "userServersStateChanged"   // 用户服务状态变化
 	NodeTaskTypeUAMPolicyChanged          NodeTaskType = "uamPolicyChanged"          // UAM策略变化
 	NodeTaskTypeHTTPPagesPolicyChanged    NodeTaskType = "httpPagesPolicyChanged"    // 自定义页面变化
+	NodeTaskTypeHTTPCCPolicyChanged       NodeTaskType = "httpCCPolicyChanged"       // CC策略变化
 	NodeTaskTypeUpdatingServers           NodeTaskType = "updatingServers"           // 更新一组服务

 	// NS相关
--- a/internal/rpc/services/service_node.go
+++ b/internal/rpc/services/service_node.go
@@ -2294,6 +2294,50 @@ func (this *NodeService) FindNodeUAMPolicies(ctx context.Context, req *pb.FindNo
 	}, nil
 }

+// FindNodeHTTPCCPolicies 查找节点的HTTP CC策略
+func (this *NodeService) FindNodeHTTPCCPolicies(ctx context.Context, req *pb.FindNodeHTTPCCPoliciesRequest) (*pb.FindNodeHTTPCCPoliciesResponse, error) {
+	nodeId, err := this.ValidateNode(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	var tx = this.NullTx()
+	clusterIds, err := models.SharedNodeDAO.FindEnabledAndOnNodeClusterIds(tx, nodeId)
+	if err != nil {
+		return nil, err
+	}
+
+	var pbPolicies = []*pb.FindNodeHTTPCCPoliciesResponse_HTTPCCPolicy{}
+	for _, clusterId := range clusterIds {
+		policy, err := models.SharedNodeClusterDAO.FindClusterHTTPCCPolicy(tx, clusterId, nil)
+		if err != nil {
+			return nil, err
+		}
+		if policy == nil {
+			continue
+		}
+
+		// 集成默认设置
+		for i := 0; i < len(serverconfigs.DefaultHTTPCCThresholds); i ++ {
+			if i < len(policy.Thresholds) {
+				policy.Thresholds[i].MergeIfEmpty(serverconfigs.DefaultHTTPCCThresholds[i])
+			}
+		}
+
+		policyJSON, err := json.Marshal(policy)
+		if err != nil {
+			return nil, err
+		}
+		pbPolicies = append(pbPolicies, &pb.FindNodeHTTPCCPoliciesResponse_HTTPCCPolicy{
+			NodeClusterId:    clusterId,
+			HttpCCPolicyJSON: policyJSON,
+		})
+	}
+	return &pb.FindNodeHTTPCCPoliciesResponse{
+		HttpCCPolicies: pbPolicies,
+	}, nil
+}
+
 // FindNodeHTTPPagesPolicies 查找节点的自定义页面策略
 func (this *NodeService) FindNodeHTTPPagesPolicies(ctx context.Context, req *pb.FindNodeHTTPPagesPoliciesRequest) (*pb.FindNodeHTTPPagesPoliciesResponse, error) {
 	nodeId, err := this.ValidateNode(ctx)
--- a/internal/rpc/services/service_node_cluster.go
+++ b/internal/rpc/services/service_node_cluster.go
@@ -3,6 +3,7 @@ package services
 import (
 	"context"
 	"encoding/json"
+	teaconst "github.com/TeaOSLab/EdgeAPI/internal/const"
 	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
 	"github.com/TeaOSLab/EdgeAPI/internal/db/models/dns"
 	"github.com/TeaOSLab/EdgeAPI/internal/db/models/dns/dnsutils"
@@ -1140,6 +1141,18 @@ func (this *NodeClusterService) FindEnabledNodeClusterConfigInfo(ctx context.Con
 		result.UamIsOn = nodeconfigs.DefaultUAMPolicy.IsOn
 	}

+	// HTTP CC
+	if models.IsNotNull(cluster.Cc) {
+		var httpCCPolicy = nodeconfigs.NewHTTPCCPolicy()
+		err = json.Unmarshal(cluster.Cc, httpCCPolicy)
+		if err != nil {
+			return nil, err
+		}
+		result.HttpCCIsOn = httpCCPolicy.IsOn
+	} else {
+		result.HttpCCIsOn = nodeconfigs.NewHTTPCCPolicy().IsOn
+	}
+
 	// system service
 	if models.IsNotNull(cluster.SystemServices) {
 		var servicesMap = map[string]maps.Map{}
@@ -1235,6 +1248,10 @@ func (this *NodeClusterService) UpdateNodeClusterWebPPolicy(ctx context.Context,

 // FindEnabledNodeClusterUAMPolicy 读取集群UAM策略
 func (this *NodeClusterService) FindEnabledNodeClusterUAMPolicy(ctx context.Context, req *pb.FindEnabledNodeClusterUAMPolicyRequest) (*pb.FindEnabledNodeClusterUAMPolicyResponse, error) {
+	if !teaconst.IsPlus {
+		return nil, this.NotImplementedYet()
+	}
+
 	_, _, err := this.ValidateAdminAndUser(ctx, false)
 	if err != nil {
 		return nil, err
@@ -1256,6 +1273,10 @@ func (this *NodeClusterService) FindEnabledNodeClusterUAMPolicy(ctx context.Cont

 // UpdateNodeClusterUAMPolicy 设置集群的UAM策略
 func (this *NodeClusterService) UpdateNodeClusterUAMPolicy(ctx context.Context, req *pb.UpdateNodeClusterUAMPolicyRequest) (*pb.RPCSuccess, error) {
+	if !teaconst.IsPlus {
+		return nil, this.NotImplementedYet()
+	}
+
 	_, err := this.ValidateAdmin(ctx)
 	if err != nil {
 		return nil, err
@@ -1280,6 +1301,62 @@ func (this *NodeClusterService) UpdateNodeClusterUAMPolicy(ctx context.Context,
 	return this.Success()
 }

+
+// FindEnabledNodeClusterHTTPCCPolicy 读取集群HTTP CC策略
+func (this *NodeClusterService) FindEnabledNodeClusterHTTPCCPolicy(ctx context.Context, req *pb.FindEnabledNodeClusterHTTPCCPolicyRequest) (*pb.FindEnabledNodeClusterHTTPCCPolicyResponse, error) {
+	if !teaconst.IsPlus {
+		return nil, this.NotImplementedYet()
+	}
+
+	_, _, err := this.ValidateAdminAndUser(ctx, false)
+	if err != nil {
+		return nil, err
+	}
+
+	var tx = this.NullTx()
+	httpCCPolicy, err := models.SharedNodeClusterDAO.FindClusterHTTPCCPolicy(tx, req.NodeClusterId, nil)
+	if err != nil {
+		return nil, err
+	}
+	httpCCPolicyJSON, err := json.Marshal(httpCCPolicy)
+	if err != nil {
+		return nil, err
+	}
+	return &pb.FindEnabledNodeClusterHTTPCCPolicyResponse{
+		HttpCCPolicyJSON: httpCCPolicyJSON,
+	}, nil
+}
+
+// UpdateNodeClusterHTTPCCPolicy 设置集群的HTTP CC策略
+func (this *NodeClusterService) UpdateNodeClusterHTTPCCPolicy(ctx context.Context, req *pb.UpdateNodeClusterHTTPCCPolicyRequest) (*pb.RPCSuccess, error) {
+	if !teaconst.IsPlus {
+		return nil, this.NotImplementedYet()
+	}
+
+	_, err := this.ValidateAdmin(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	var httpCCPolicy = nodeconfigs.NewHTTPCCPolicy()
+	err = json.Unmarshal(req.HttpCCPolicyJSON, httpCCPolicy)
+	if err != nil {
+		return nil, err
+	}
+
+	err = httpCCPolicy.Init()
+	if err != nil {
+		return nil, errors.New("validate http cc policy failed: " + err.Error())
+	}
+
+	var tx = this.NullTx()
+	err = models.SharedNodeClusterDAO.UpdateClusterHTTPCCPolicy(tx, req.NodeClusterId, httpCCPolicy)
+	if err != nil {
+		return nil, err
+	}
+	return this.Success()
+}
+
 // FindNodeClusterDDoSProtection 获取集群的DDoS设置
 func (this *NodeClusterService) FindNodeClusterDDoSProtection(ctx context.Context, req *pb.FindNodeClusterDDoSProtectionRequest) (*pb.FindNodeClusterDDoSProtectionResponse, error) {
 	_, err := this.ValidateAdmin(ctx)