WAF记录ServerId

internal/db/models/http_firewall_policy_dao.go

@@ -91,9 +91,10 @@ func (this *HTTPFirewallPolicyDAO) FindAllEnabledFirewallPolicies(tx *dbs.Tx) (r
 }
 
 // 创建策略
-func (this *HTTPFirewallPolicyDAO) CreateFirewallPolicy(tx *dbs.Tx, userId int64, isOn bool, name string, description string, inboundJSON []byte, outboundJSON []byte) (int64, error) {
+func (this *HTTPFirewallPolicyDAO) CreateFirewallPolicy(tx *dbs.Tx, userId int64, serverId int64, isOn bool, name string, description string, inboundJSON []byte, outboundJSON []byte) (int64, error) {
 	op := NewHTTPFirewallPolicyOperator()
 	op.UserId = userId
+	op.ServerId = serverId
 	op.State = HTTPFirewallPolicyStateEnabled
 	op.IsOn = isOn
 	op.Name = name
@@ -177,6 +178,7 @@ func (this *HTTPFirewallPolicyDAO) CountAllEnabledFirewallPolicies(tx *dbs.Tx) (
 	return this.Query(tx).
 		State(HTTPFirewallPolicyStateEnabled).
 		Attr("userId", 0).
+		Attr("serverId", 0).
 		Count()
 }
 
@@ -185,6 +187,7 @@ func (this *HTTPFirewallPolicyDAO) ListEnabledFirewallPolicies(tx *dbs.Tx, offse
 	_, err = this.Query(tx).
 		State(HTTPFirewallPolicyStateEnabled).
 		Attr("userId", 0).
+		Attr("serverId", 0).
 		Offset(offset).
 		Limit(size).
 		DescPk().
@@ -324,6 +327,15 @@ func (this *HTTPFirewallPolicyDAO) FindEnabledFirewallPolicyIdWithRuleGroupId(tx
 		FindInt64Col(0)
 }
 
+// 设置某个策略所属的服务ID
+func (this *HTTPFirewallPolicyDAO) UpdateFirewallPolicyServerId(tx *dbs.Tx, policyId int64, serverId int64) error {
+	_, err :=  this.Query(tx).
+		Pk(policyId).
+		Set("serverId", serverId).
+		Update()
+	return err
+}
+
 // 通知更新
 func (this *HTTPFirewallPolicyDAO) NotifyUpdate(tx *dbs.Tx, policyId int64) error {
 	webIds, err := SharedHTTPWebDAO.FindAllWebIdsWithHTTPFirewallPolicyId(tx, policyId)

internal/db/models/http_firewall_policy_model.go

@@ -6,6 +6,7 @@ type HTTPFirewallPolicy struct {
 	TemplateId   uint32 `field:"templateId"`   // 模版ID
 	AdminId      uint32 `field:"adminId"`      // 管理员ID
 	UserId       uint32 `field:"userId"`       // 用户ID
+	ServerId     uint32 `field:"serverId"`     // 服务ID
 	State        uint8  `field:"state"`        // 状态
 	CreatedAt    uint64 `field:"createdAt"`    // 创建时间
 	IsOn         uint8  `field:"isOn"`         // 是否启用
@@ -21,6 +22,7 @@ type HTTPFirewallPolicyOperator struct {
 	TemplateId   interface{} // 模版ID
 	AdminId      interface{} // 管理员ID
 	UserId       interface{} // 用户ID
+	ServerId     interface{} // 服务ID
 	State        interface{} // 状态
 	CreatedAt    interface{} // 创建时间
 	IsOn         interface{} // 是否启用

internal/rpc/services/service_http_firewall_policy.go

@@ -56,7 +56,7 @@ func (this *HTTPFirewallPolicyService) CreateHTTPFirewallPolicy(ctx context.Cont
 
 	tx := this.NullTx()
 
-	policyId, err := models.SharedHTTPFirewallPolicyDAO.CreateFirewallPolicy(tx, userId, req.IsOn, req.Name, req.Description, nil, nil)
+	policyId, err := models.SharedHTTPFirewallPolicyDAO.CreateFirewallPolicy(tx, userId, req.ServerId, req.IsOn, req.Name, req.Description, nil, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -122,9 +122,18 @@ func (this *HTTPFirewallPolicyService) CreateEmptyHTTPFirewallPolicy(ctx context
 		return nil, err
 	}
 
+	if userId > 0 {
+		if req.ServerId > 0 {
+			err = models.SharedServerDAO.CheckUserServer(nil, userId, req.ServerId)
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
+
 	tx := this.NullTx()
 
-	policyId, err := models.SharedHTTPFirewallPolicyDAO.CreateFirewallPolicy(tx, userId, req.IsOn, req.Name, req.Description, nil, nil)
+	policyId, err := models.SharedHTTPFirewallPolicyDAO.CreateFirewallPolicy(tx, userId, req.ServerId, req.IsOn, req.Name, req.Description, nil, nil)
 	if err != nil {
 		return nil, err
 	}

internal/setup/sql_data.go

@@ -2,8 +2,10 @@ package setup
 
 import (
 	"github.com/TeaOSLab/EdgeAPI/internal/acme"
+	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
 	"github.com/TeaOSLab/EdgeAPI/internal/errors"
 	"github.com/iwind/TeaGo/dbs"
+	"github.com/iwind/TeaGo/lists"
 	"github.com/iwind/TeaGo/rands"
 	"github.com/iwind/TeaGo/types"
 	stringutil "github.com/iwind/TeaGo/utils/string"
@@ -24,6 +26,9 @@ var upgradeFuncs = []*upgradeVersion{
 	{
 		"0.0.6", upgradeV0_0_6,
 	},
+	{
+		"0.0.9", upgradeV0_0_9,
+	},
 }
 
 // 升级SQL数据
@@ -136,3 +141,42 @@ func upgradeV0_0_6(db *dbs.DB) error {
 
 	return nil
 }
+
+// v0.0.9
+func upgradeV0_0_9(db *dbs.DB) error {
+	// firewall policies
+	var tx *dbs.Tx
+	dbs.NotifyReady()
+	policies, err := models.NewHTTPFirewallPolicyDAO().FindAllEnabledFirewallPolicies(tx)
+	if err != nil {
+		return err
+	}
+	for _, policy := range policies {
+		if policy.ServerId > 0 {
+			continue
+		}
+		policyId := int64(policy.Id)
+		webIds, err := models.NewHTTPWebDAO().FindAllWebIdsWithHTTPFirewallPolicyId(tx, policyId)
+		if err != nil {
+			return err
+		}
+		serverIds := []int64{}
+		for _, webId := range webIds {
+			serverId, err := models.NewServerDAO().FindEnabledServerIdWithWebId(tx, webId)
+			if err != nil {
+				return err
+			}
+			if serverId > 0 && !lists.ContainsInt64(serverIds, serverId) {
+				serverIds = append(serverIds, serverId)
+			}
+		}
+		if len(serverIds) == 1 {
+			err = models.NewHTTPFirewallPolicyDAO().UpdateFirewallPolicyServerId(tx, policyId, serverIds[0])
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}