waf-platform/EdgeAPI/internal/rpc/services/httpdns/service_httpdns_rule.go

package httpdns

import (
	"context"
	"errors"
	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
	"github.com/TeaOSLab/EdgeAPI/internal/rpc/services"
	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
	"github.com/iwind/TeaGo/dbs"
)

// HTTPDNSRuleService HTTPDNS规则服务
type HTTPDNSRuleService struct {
	services.BaseService
	pb.UnimplementedHTTPDNSRuleServiceServer
}

func (this *HTTPDNSRuleService) CreateHTTPDNSCustomRule(ctx context.Context, req *pb.CreateHTTPDNSCustomRuleRequest) (*pb.CreateHTTPDNSCustomRuleResponse, error) {
	_, userId, err := this.ValidateAdminAndUser(ctx, true)
	if err != nil {
		return nil, err
	}
	if req.Rule == nil {
		return nil, errors.New("required 'rule'")
	}
	var ruleId int64
	err = this.RunTx(func(tx *dbs.Tx) error {
		domain, app, err := ensureDomainAccess(tx, req.Rule.DomainId, userId)
		if err != nil {
			return err
		}
		if domain == nil || app == nil {
			return errors.New("domain not found")
		}

		rule := &models.HTTPDNSCustomRule{
			AppId:         domain.AppId,
			DomainId:      uint32(req.Rule.DomainId),
			RuleName:      req.Rule.RuleName,
			LineScope:     req.Rule.LineScope,
			LineCarrier:   req.Rule.LineCarrier,
			LineRegion:    req.Rule.LineRegion,
			LineProvince:  req.Rule.LineProvince,
			LineContinent: req.Rule.LineContinent,
			LineCountry:   req.Rule.LineCountry,
			TTL:           req.Rule.Ttl,
			IsOn:          req.Rule.IsOn,
			Priority:      req.Rule.Priority,
		}
		ruleId, err = models.SharedHTTPDNSCustomRuleDAO.CreateRule(tx, rule)
		if err != nil {
			return err
		}
		for _, record := range req.Rule.Records {
			_, err := models.SharedHTTPDNSCustomRuleRecordDAO.CreateRecord(tx, ruleId, record.RecordType, record.RecordValue, record.Weight, record.Sort)
			if err != nil {
				return err
			}
		}
		return notifyHTTPDNSAppTasksByAppDbId(tx, int64(app.Id), models.HTTPDNSNodeTaskTypeRuleChanged)
	})
	if err != nil {
		return nil, err
	}
	return &pb.CreateHTTPDNSCustomRuleResponse{RuleId: ruleId}, nil
}

func (this *HTTPDNSRuleService) UpdateHTTPDNSCustomRule(ctx context.Context, req *pb.UpdateHTTPDNSCustomRuleRequest) (*pb.RPCSuccess, error) {
	_, userId, err := this.ValidateAdminAndUser(ctx, true)
	if err != nil {
		return nil, err
	}
	if req.Rule == nil || req.Rule.Id <= 0 {
		return nil, errors.New("invalid 'rule.id'")
	}
	err = this.RunTx(func(tx *dbs.Tx) error {
		oldRule, app, err := ensureRuleAccess(tx, req.Rule.Id, userId)
		if err != nil {
			return err
		}
		if oldRule == nil {
			return errors.New("rule not found")
		}

		rule := &models.HTTPDNSCustomRule{
			Id:            uint32(req.Rule.Id),
			RuleName:      req.Rule.RuleName,
			LineScope:     req.Rule.LineScope,
			LineCarrier:   req.Rule.LineCarrier,
			LineRegion:    req.Rule.LineRegion,
			LineProvince:  req.Rule.LineProvince,
			LineContinent: req.Rule.LineContinent,
			LineCountry:   req.Rule.LineCountry,
			TTL:           req.Rule.Ttl,
			IsOn:          req.Rule.IsOn,
			Priority:      req.Rule.Priority,
		}
		err = models.SharedHTTPDNSCustomRuleDAO.UpdateRule(tx, rule)
		if err != nil {
			return err
		}
		err = models.SharedHTTPDNSCustomRuleRecordDAO.DisableRecordsWithRuleId(tx, req.Rule.Id)
		if err != nil {
			return err
		}
		for _, record := range req.Rule.Records {
			_, err := models.SharedHTTPDNSCustomRuleRecordDAO.CreateRecord(tx, req.Rule.Id, record.RecordType, record.RecordValue, record.Weight, record.Sort)
			if err != nil {
				return err
			}
		}
		err = notifyHTTPDNSAppTasksByAppDbId(tx, int64(app.Id), models.HTTPDNSNodeTaskTypeRuleChanged)
		if err != nil {
			return err
		}

		targetAppDbId := int64(app.Id)
		return notifyHTTPDNSAppTasksByAppDbId(tx, targetAppDbId, models.HTTPDNSNodeTaskTypeRuleChanged)
	})
	if err != nil {
		return nil, err
	}
	return this.Success()
}

func (this *HTTPDNSRuleService) DeleteHTTPDNSCustomRule(ctx context.Context, req *pb.DeleteHTTPDNSCustomRuleRequest) (*pb.RPCSuccess, error) {
	_, userId, err := this.ValidateAdminAndUser(ctx, true)
	if err != nil {
		return nil, err
	}
	err = this.RunTx(func(tx *dbs.Tx) error {
		rule, app, err := ensureRuleAccess(tx, req.RuleId, userId)
		if err != nil {
			return err
		}
		if rule == nil {
			return nil
		}

		err = models.SharedHTTPDNSCustomRuleDAO.DisableRule(tx, req.RuleId)
		if err != nil {
			return err
		}
		return notifyHTTPDNSAppTasksByAppDbId(tx, int64(app.Id), models.HTTPDNSNodeTaskTypeRuleChanged)
	})
	if err != nil {
		return nil, err
	}
	return this.Success()
}

func (this *HTTPDNSRuleService) UpdateHTTPDNSCustomRuleStatus(ctx context.Context, req *pb.UpdateHTTPDNSCustomRuleStatusRequest) (*pb.RPCSuccess, error) {
	_, userId, err := this.ValidateAdminAndUser(ctx, true)
	if err != nil {
		return nil, err
	}
	err = this.RunTx(func(tx *dbs.Tx) error {
		rule, app, err := ensureRuleAccess(tx, req.RuleId, userId)
		if err != nil {
			return err
		}
		if rule == nil {
			return nil
		}

		err = models.SharedHTTPDNSCustomRuleDAO.UpdateRuleStatus(tx, req.RuleId, req.IsOn)
		if err != nil {
			return err
		}
		return notifyHTTPDNSAppTasksByAppDbId(tx, int64(app.Id), models.HTTPDNSNodeTaskTypeRuleChanged)
	})
	if err != nil {
		return nil, err
	}
	return this.Success()
}

func (this *HTTPDNSRuleService) ListHTTPDNSCustomRulesWithDomainId(ctx context.Context, req *pb.ListHTTPDNSCustomRulesWithDomainIdRequest) (*pb.ListHTTPDNSCustomRulesWithDomainIdResponse, error) {
	_, userId, validateErr := this.ValidateAdminAndUser(ctx, true)
	if validateErr != nil {
		if _, nodeErr := this.ValidateHTTPDNSNode(ctx); nodeErr != nil {
			return nil, validateErr
		}
	} else if userId > 0 {
		domain, _, err := ensureDomainAccess(this.NullTx(), req.DomainId, userId)
		if err != nil {
			return nil, err
		}
		if domain == nil {
			return &pb.ListHTTPDNSCustomRulesWithDomainIdResponse{}, nil
		}
	}
	rules, err := models.SharedHTTPDNSCustomRuleDAO.ListEnabledRulesWithDomainId(this.NullTx(), req.DomainId)
	if err != nil {
		return nil, err
	}
	var pbRules []*pb.HTTPDNSCustomRule
	for _, rule := range rules {
		records, err := models.SharedHTTPDNSCustomRuleRecordDAO.ListEnabledRecordsWithRuleId(this.NullTx(), int64(rule.Id))
		if err != nil {
			return nil, err
		}
		pbRules = append(pbRules, toPBRule(rule, records))
	}
	return &pb.ListHTTPDNSCustomRulesWithDomainIdResponse{Rules: pbRules}, nil
}