waf-platform/EdgeUser/internal/web/actions/default/ns/nsutils/utils.go

// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .

package nsutils

import (
	"context"
	"encoding/json"
	"errors"
	"github.com/TeaOSLab/EdgeCommon/pkg/dnsconfigs"
	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
	"github.com/TeaOSLab/EdgeCommon/pkg/systemconfigs"
	"github.com/TeaOSLab/EdgeUser/internal/rpc"
	timeutil "github.com/iwind/TeaGo/utils/time"
	"strings"
)

// FindUserPlanConfig 用户配置
// 需要保证非 error 的情况下，一定会返回一个不为空的 NSPlanConfig
func FindUserPlanConfig(ctx context.Context) (*dnsconfigs.NSPlanConfig, error) {
	rpcClient, err := rpc.SharedRPC()
	if err != nil {
		return nil, err
	}

	// 从套餐中读取
	userPlanResp, err := rpcClient.NSUserPlanRPC().FindNSUserPlan(ctx, &pb.FindNSUserPlanRequest{})
	if err != nil {
		return nil, err
	}

	if userPlanResp.NsUserPlan != nil &&
		len(userPlanResp.NsUserPlan.DayTo) > 0 &&
		userPlanResp.NsUserPlan.DayTo >= timeutil.Format("Ymd") /** 在有效期内 **/ {
		var planId = userPlanResp.NsUserPlan.NsPlanId
		if planId > 0 {
			planResp, err := rpcClient.NSPlanRPC().FindNSPlan(ctx, &pb.FindNSPlanRequest{NsPlanId: planId})
			if err != nil {
				return nil, err
			}
			if planResp.NsPlan != nil &&
				len(planResp.NsPlan.ConfigJSON) > 0 &&
				planResp.NsPlan.IsOn {
				var config = dnsconfigs.DefaultNSUserPlanConfig()
				err = json.Unmarshal(planResp.NsPlan.ConfigJSON, config)
				if err != nil {
					return nil, errors.New("decode plan config failed '" + err.Error() + "'")
				}
				return config, nil
			}
		}
	}

	// 从用户设置中读取
	userConfigResp, err := rpcClient.SysSettingRPC().ReadSysSetting(ctx, &pb.ReadSysSettingRequest{Code: systemconfigs.SettingCodeNSUserConfig})
	if err != nil {
		return nil, err
	}
	if len(userConfigResp.ValueJSON) > 0 {
		var config = dnsconfigs.NewNSUserConfig()
		err = json.Unmarshal(userConfigResp.ValueJSON, config)
		if err != nil {
			return nil, err
		}
		if config.DefaultPlanConfig != nil {
			return config.DefaultPlanConfig, nil
		}
	}

	return dnsconfigs.DefaultNSUserPlanConfig(), nil
}

// FindBasicPlan 查找基础套餐
func FindBasicPlan(ctx context.Context) (*dnsconfigs.NSPlanConfig, error) {
	rpcClient, err := rpc.SharedRPC()
	if err != nil {
		return nil, err
	}

	// 从用户设置中读取
	userConfigResp, err := rpcClient.SysSettingRPC().ReadSysSetting(ctx, &pb.ReadSysSettingRequest{Code: systemconfigs.SettingCodeNSUserConfig})
	if err != nil {
		return nil, err
	}
	if len(userConfigResp.ValueJSON) > 0 {
		var config = dnsconfigs.NewNSUserConfig()
		err = json.Unmarshal(userConfigResp.ValueJSON, config)
		if err != nil {
			return nil, err
		}
		if config.DefaultPlanConfig != nil {
			return config.DefaultPlanConfig, nil
		}
	}

	return dnsconfigs.DefaultNSUserPlanConfig(), nil
}

// CheckDomainsQuote 检查域名限额
func CheckDomainsQuote(ctx context.Context, countNew int32) (maxDomains int32, ok bool, err error) {
	config, err := FindUserPlanConfig(ctx)
	if err != nil {
		return 0, false, err
	}
	maxDomains = config.MaxDomains
	if maxDomains <= 0 {
		return 0, true, nil
	}

	rpcClient, err := rpc.SharedRPC()
	if err != nil {
		return maxDomains, false, err
	}

	countDomainsResp, err := rpcClient.NSDomainRPC().CountAllNSDomains(ctx, &pb.CountAllNSDomainsRequest{})
	if err != nil {
		return maxDomains, false, err
	}

	if countDomainsResp.Count+int64(countNew) > int64(maxDomains) {
		return maxDomains, false, nil
	}
	return maxDomains, true, nil
}

// CheckRecordsQuota 检查记录限额
func CheckRecordsQuota(ctx context.Context, domainId int64, countNew int32) (maxRecords int32, ok bool, err error) {
	config, err := FindUserPlanConfig(ctx)
	if err != nil {
		return 0, false, err
	}
	maxRecords = config.MaxRecordsPerDomain
	if maxRecords <= 0 {
		return 0, true, nil
	}

	rpcClient, err := rpc.SharedRPC()
	if err != nil {
		return maxRecords, false, err
	}

	countRecordsResp, err := rpcClient.NSRecordRPC().CountAllNSRecords(ctx, &pb.CountAllNSRecordsRequest{
		NsDomainId: domainId,
	})
	if err != nil {
		return maxRecords, false, err
	}

	if countRecordsResp.Count+int64(countNew) > int64(maxRecords) {
		return maxRecords, false, nil
	}
	return maxRecords, true, nil
}

// CheckLoadBalanceRecordsQuota 检查负载均衡限额
func CheckLoadBalanceRecordsQuota(ctx context.Context, domainId int64, recordName string, recordType string, countNew int32) (maxRecords int32, ok bool, err error) {
	recordName = strings.ToLower(recordName)
	recordType = strings.ToUpper(recordType)

	config, err := FindUserPlanConfig(ctx)
	if err != nil {
		return 0, false, err
	}
	maxRecords = config.MaxLoadBalanceRecordsPerRecord
	if maxRecords <= 0 {
		return 0, true, nil
	}

	rpcClient, err := rpc.SharedRPC()
	if err != nil {
		return maxRecords, false, err
	}

	countRecordsResp, err := rpcClient.NSRecordRPC().CountAllNSRecordsWithName(ctx, &pb.CountAllNSRecordsWithNameRequest{
		NsDomainId: domainId,
		Name:       recordName,
		Type:       recordType,
	})
	if err != nil {
		return maxRecords, false, err
	}

	if countRecordsResp.Count+int64(countNew) > int64(maxRecords) {
		return maxRecords, false, nil
	}
	return maxRecords, true, nil
}

// FindMinTTL 获取最小的TTL
func FindMinTTL(ctx context.Context) (int32, error) {
	config, err := FindUserPlanConfig(ctx)
	if err != nil {
		return 0, err
	}
	return config.MinTTL, nil
}

func FindDefaultTTL(ctx context.Context) (int32, error) {
	config, err := FindUserPlanConfig(ctx)
	if err != nil {
		return 0, err
	}
	var minTTL = config.MinTTL

	if minTTL > 600 {
		return minTTL, nil
	}
	return 600, nil
}

// CheckRoutesQuota 检查线路限额
func CheckRoutesQuota(ctx context.Context, countNew int32) (maxCustomRoutes int32, ok bool, err error) {
	config, err := FindUserPlanConfig(ctx)
	if err != nil {
		return 0, false, err
	}
	maxCustomRoutes = config.MaxCustomRoutes
	if maxCustomRoutes <= 0 {
		return 0, true, nil
	}

	rpcClient, err := rpc.SharedRPC()
	if err != nil {
		return maxCustomRoutes, false, err
	}

	countRoutesResp, err := rpcClient.NSRouteRPC().CountAllNSRoutes(ctx, &pb.CountAllNSRoutesRequest{})
	if err != nil {
		return maxCustomRoutes, false, err
	}

	if countRoutesResp.Count+int64(countNew) > int64(maxCustomRoutes) {
		return maxCustomRoutes, false, nil
	}
	return maxCustomRoutes, true, nil
}