go_gameshield/internal/task/waf.go

package task

import (
	"context"
	"fmt"
	v1 "github.com/go-nunu/nunu-layout-advanced/api/v1"
	"github.com/go-nunu/nunu-layout-advanced/internal/model"
	"github.com/go-nunu/nunu-layout-advanced/internal/repository"
	"github.com/go-nunu/nunu-layout-advanced/internal/service"
	"github.com/hashicorp/go-multierror"
	"go.uber.org/zap"
	"sync"
	"time"
)

type WafTask interface {
	//获取到期时间小于3天的同步时间
	SynchronizationTime(ctx context.Context) error
	StopPlan(ctx context.Context) error
	RecoverStopPlan(ctx context.Context) error
}

func NewWafTask (
	webForWardingRep repository.WebForwardingRepository,
	tcpforwardingRep repository.TcpforwardingRepository,
	udpForWardingRep repository.UdpForWardingRepository,
	cdn service.CdnService,
	hostRep repository.HostRepository,
	globalLimitRep repository.GlobalLimitRepository,
	expiredRep repository.ExpiredRepository,
	task *Task,
	) WafTask{
	return &wafTask{
		Task: task,
		webForWardingRep: webForWardingRep,
		tcpforwardingRep: tcpforwardingRep,
		udpForWardingRep: udpForWardingRep,
		cdn: cdn,
		hostRep: hostRep,
		globalLimitRep: globalLimitRep,
		expiredRep: expiredRep,
	}
}
type wafTask struct {
	*Task
	webForWardingRep repository.WebForwardingRepository
	tcpforwardingRep repository.TcpforwardingRepository
	udpForWardingRep repository.UdpForWardingRepository
	cdn service.CdnService
	hostRep repository.HostRepository
	globalLimitRep repository.GlobalLimitRepository
	expiredRep repository.ExpiredRepository
}


const (
	// 1天后秒数
	OneDaysInSeconds = 1 * 24 * 60 * 60
	// 7天前秒数
	SevenDaysInSeconds = 7 * 24 * 60 * 60 * -1
)
// 获取cdn web id
func (t wafTask) GetCdnWebId(ctx context.Context,hostId []int) ([]int, error) {
	tcpIds, err := t.tcpforwardingRep.GetTcpAll(ctx, hostId)
	if err != nil {
		return nil, err
	}
	udpIds, err := t.udpForWardingRep.GetUdpAll(ctx, hostId)
	if err != nil {
		return nil, err
	}
	webIds, err := t.webForWardingRep.GetWebAll(ctx, hostId)
	if err != nil {
		return nil, err
	}
	var ids []int
	ids = append(ids, tcpIds...)
	ids = append(ids, udpIds...)
	ids = append(ids, webIds...)
	return ids, nil
}

// 启用/禁用 网站
func (t wafTask) BanServer(ctx context.Context, ids []int, isBan bool) error {
	var wg sync.WaitGroup
	errChan := make(chan error, len(ids))

	// 修正1：为每个 goroutine 增加 WaitGroup 的计数
	wg.Add(len(ids))

	for _, id := range ids {
		go func(id int) {
			// 修正2：确保每个 goroutine 在退出时都调用 Done()
			defer wg.Done()

			err := t.cdn.EditWebIsOn(ctx, int64(id), isBan)
			if err != nil {
				errChan <- err
				// 这里不需要 return，因为 defer wg.Done() 会在函数退出时执行
			}
		}(id)
	}

	// 现在 wg.Wait() 会正确地阻塞，直到所有 goroutine 都调用了 Done()
	wg.Wait()

	// 在所有 goroutine 都结束后，安全地关闭 channel
	close(errChan)

	var result error
	for err := range errChan {
		result = multierror.Append(result, err)  // 将多个 error 对象合并成一个单一的 error 对象
	}

	// 修正3：返回收集到的错误，而不是 nil
	return result
}



// 获取指定到期时间
func (t wafTask) GetAlmostExpiring(ctx context.Context,hostIds []int,addTime int64) ([]v1.GetAlmostExpireHostResponse,error) {
	// 3 天
	res, err := t.hostRep.GetAlmostExpired(ctx, hostIds, addTime)
	if err != nil {
		return nil,err
	}

	return res, nil
}


// 获取waf全局到期时间
func (t wafTask) GetGlobalAlmostExpiring(ctx context.Context,addTime int64) ([]model.GlobalLimit,error) {
	res, err := t.globalLimitRep.GetGlobalLimitAlmostExpired(ctx, addTime)
	if err != nil {
		return nil, err
	}
	return res, nil
}



// 修改全局续费
func (t wafTask) EditGlobalExpired(ctx context.Context, req []struct{
	hostId int
	expiredAt int64
}, state bool) error {
	var result *multierror.Error // 使用 multierror

	for _, v := range req {
		err := t.globalLimitRep.UpdateGlobalLimitByHostId(ctx, &model.GlobalLimit{
			HostId:    v.hostId,
			ExpiredAt: v.expiredAt,
			State:     state,
		})
		if err != nil {
			// 收集错误，而不是直接返回
			result = multierror.Append(result, err)
		}
	}

	// 返回所有收集到的错误
	return result.ErrorOrNil()
}



// 续费套餐
func (t wafTask) EnablePlan(ctx context.Context, req []struct{
	planId int
	expiredAt int64
}) error {
	var result *multierror.Error

	for _, v := range req {
		err := t.cdn.RenewPlan(ctx, v1.RenewalPlan{
			UserPlanId:  int64(v.planId),
			IsFree:      true,
			DayTo:       time.Unix(v.expiredAt, 0).Format("2006-01-02"),
			Period:      "monthly",
			CountPeriod: 1,
			PeriodDayTo: time.Unix(v.expiredAt, 0).Format("2006-01-02"),
		})
		if err != nil {
			result = multierror.Append(result, err)
		}
	}

	return result.ErrorOrNil()
}



// 续费操作
type RenewalRequest struct {
	HostId    int
	PlanId    int
	ExpiredAt int64
}

// 续费操作
func (t wafTask) EditExpired(ctx context.Context, reqs []RenewalRequest) error {
	// 如果请求为空，直接返回
	if len(reqs) == 0 {
		return nil
	}

	// 1. 准备用于更新 GlobalLimit 的数据
	var globalLimitUpdates []struct {
		hostId    int
		expiredAt int64
	}
	for _, req := range reqs {
		globalLimitUpdates = append(globalLimitUpdates, struct {
			hostId    int
			expiredAt int64
		}{req.HostId, req.ExpiredAt})
	}

	// 2. 准备用于续费套餐的数据
	var planRenewals []struct {
		planId    int
		expiredAt int64
	}
	for _, req := range reqs {
		planRenewals = append(planRenewals, struct {
			planId    int
			expiredAt int64
		}{req.PlanId, req.ExpiredAt})
	}

	var result *multierror.Error

	// 3. 执行更新，并收集错误
	if err := t.EditGlobalExpired(ctx, globalLimitUpdates, true); err != nil {
		result = multierror.Append(result, err)
	}

	if err := t.EnablePlan(ctx, planRenewals); err != nil {
		result = multierror.Append(result, err)
	}

	return result.ErrorOrNil()
}



// findMismatchedExpirations 检查 WAF 和 Host 的到期时间差异，并返回需要同步的请求。
func (t *wafTask) findMismatchedExpirations(ctx context.Context, wafLimits []model.GlobalLimit) ([]RenewalRequest, error) {
	if len(wafLimits) == 0 {
		return nil, nil
	}

	// 2. 将 WAF 数据组织成 Map
	wafExpiredMap := make(map[int]int64, len(wafLimits))
	wafPlanMap := make(map[int]int, len(wafLimits))
	var hostIds []int
	for _, limit := range wafLimits {
		hostIds = append(hostIds, limit.HostId)
		wafExpiredMap[limit.HostId] = limit.ExpiredAt
		wafPlanMap[limit.HostId] = limit.RuleId
	}

	// 3. 获取对应 Host 的到期时间
	hostExpirations, err := t.hostRep.GetExpireTimeByHostId(ctx, hostIds)
	if err != nil {
		return nil, fmt.Errorf("获取主机到期时间失败: %w", err)
	}
	hostExpiredMap := make(map[int]int64, len(hostExpirations))
	for _, h := range hostExpirations {
		hostExpiredMap[h.HostId] = h.ExpiredAt
	}

	// 4. 找出时间不一致的记录
	var renewalRequests []RenewalRequest
	for hostId, wafExpiredTime := range wafExpiredMap {
		hostTime, ok := hostExpiredMap[hostId]

		// 如果 Host 时间与 WAF 时间不一致，则需要同步
		if !ok || hostTime != wafExpiredTime {
			planId, planOk := wafPlanMap[hostId]
			if !planOk {
				t.logger.Warn("数据不一致：在waf_limits中找不到hostId对应的套餐ID", zap.Int("hostId", hostId))
				continue
			}
			renewalRequests = append(renewalRequests, RenewalRequest{
				HostId:    hostId,
				ExpiredAt: hostTime, // 以 host 表的时间为准
				PlanId:    planId,
			})
		}
	}

	return renewalRequests, nil
}


//获取同步到期时间小于1天的套餐

func (t *wafTask) SynchronizationTime(ctx context.Context) error {
	// 1. 获取 WAF 全局配置中即将到期（小于3天）的数据
	wafLimits, err := t.GetGlobalAlmostExpiring(ctx, OneDaysInSeconds)
	if err != nil {
		return fmt.Errorf("获取全局到期配置失败: %w", err)
	}

	// 2. 找出需要同步的数据
	renewalRequests, err := t.findMismatchedExpirations(ctx, wafLimits)
	if err != nil {
		return err // 错误已在辅助函数中包装
	}

	// 3. 如果有需要同步的数据，执行续费操作
	if len(renewalRequests) > 0 {
		t.logger.Info("发现记录需要同步到期时间。", zap.Int("数量", len(renewalRequests)))
		return t.EditExpired(ctx, renewalRequests)
	}

	return nil
}



// 获取到期的进行关闭套餐操作
func (t *wafTask) StopPlan(ctx context.Context) error {
	// 1. 获取 WAF 全局配置中已经到期的数据
	// 使用 time.Now().Unix() 表示获取所有 expired_at <= 当前时间的记录
	wafLimits, err := t.globalLimitRep.GetGlobalLimitAlmostExpired(ctx, 0)
	if err != nil {
		return fmt.Errorf("获取全局到期配置失败: %w", err)
	}
	if len(wafLimits) == 0 {
		return nil // 没有到期的，任务完成
	}

	// 2. （可选，但推荐）先同步任何时间不一致的数据，确保状态准确
	renewalRequests, err := t.findMismatchedExpirations(ctx, wafLimits)
	if err != nil {
		t.logger.Error("在关闭套餐前，同步时间失败", zap.Error(err))
		// 根据业务决定是否要继续，这里我们选择继续，但记录错误
	}
	if len(renewalRequests) > 0 {
		t.logger.Info("关闭套餐前，发现并同步不一致的时间记录", zap.Int("数量", len(renewalRequests)))
		if err := t.EditExpired(ctx, renewalRequests); err != nil {
			t.logger.Error("同步不一致的时间记录失败", zap.Error(err))
		}
	}

	// 3. 筛选出尚未被关闭的套餐
	var plansToClose []model.GlobalLimit
	for _, limit := range wafLimits {
		isClosed, err := t.expiredRep.IsPlanClosed(ctx, int64(limit.HostId))
		if err != nil {
			t.logger.Error("检查Redis中套餐关闭状态失败", zap.Int("hostId", limit.HostId), zap.Error(err))
			continue // 跳过这个，处理下一个
		}
		if !isClosed {
			plansToClose = append(plansToClose, limit)
		}
	}

	if len(plansToClose) == 0 {
		t.logger.Info("没有新的到期套餐需要关闭")
		return nil
	}

	// 4. 对筛选出的套餐执行关闭操作
	t.logger.Info("开始关闭新的到期WAF服务", zap.Int("数量", len(plansToClose)))
	var allErrors *multierror.Error

	var webIds []int
	for _, limit := range plansToClose {
		webIds = append(webIds, limit.HostId)
	}

	if err := t.BanServer(ctx, webIds, false); err != nil {
		allErrors = multierror.Append(allErrors, fmt.Errorf("关闭hostId %v 的服务失败: %w", webIds, err))
	} else {
		// 服务关闭成功后，将这些套餐信息添加到 Redis
		var expiredInfos []repository.ExpiredInfo
		for _, limit := range plansToClose {
			expiredInfos = append(expiredInfos, repository.ExpiredInfo{
				HostID: int64(limit.HostId),
				Expiry: time.Unix(limit.ExpiredAt, 0),
			})
		}

		if len(expiredInfos) > 0 {
			if err := t.expiredRep.AddClosePlans(ctx, expiredInfos...); err != nil {
				allErrors = multierror.Append(allErrors, fmt.Errorf("添加已关闭套餐信息到Redis失败: %w", err))
			}
		}
	}

	return allErrors.ErrorOrNil()
}
//对于到期7天内续费的产品需要进行恢复操作

func (t *wafTask) RecoverStopPlan(ctx context.Context) error {
	// 1. 获取所有已过期（expired_at < now）但状态仍为 true 的 WAF 记录
	// StopPlan 任务会禁用这些服务，但不会改变它们的 state
	wafLimits, err := t.globalLimitRep.GetGlobalLimitAlmostExpired(ctx, SevenDaysInSeconds) // addTime=0 表示获取所有当前时间之前到期的
	if err != nil {
		return fmt.Errorf("获取过期WAF配置失败: %w", err)
	}
	if len(wafLimits) == 0 {
		t.logger.Info("没有已过期且需要检查恢复的服务")
		return nil
	}

	// 2. 检查这些记录对应的 host 是否已续费
	// findMismatchedExpirations 会比较 waf.expired_at 和 host.nextduedate
	renewalRequests, err := t.findMismatchedExpirations(ctx, wafLimits)
	if err != nil {
		return fmt.Errorf("检查续费状态失败: %w", err)
	}

	if len(renewalRequests) == 0 {
		t.logger.Info("在已过期的服务中，没有发现已续费且需要恢复的服务")
		return nil
	}

	// 3. 对已续费的服务执行恢复操作
	t.logger.Info("发现已续费、需要恢复的WAF服务", zap.Int("数量", len(renewalRequests)))
	var allErrors *multierror.Error

	var webIds []int
	for _, req := range renewalRequests {
		webIds = append(webIds, req.HostId)
	}

	if err := t.BanServer(ctx, webIds, true); err != nil {
		allErrors = multierror.Append(allErrors, fmt.Errorf("恢复hostId %v: 启用服务失败: %w", webIds, err))
	} else {
		// 服务恢复成功后，从 Redis 中移除这些套餐的关闭记录
		planIds := make([]int64, len(webIds))
		for i, id := range webIds {
			planIds[i] = int64(id)
		}
		if err := t.expiredRep.RemoveClosePlanIds(ctx, planIds...); err != nil {
			allErrors = multierror.Append(allErrors, fmt.Errorf("从Redis移除已恢复的套餐失败: %w", err))
		}
	}

	if len(renewalRequests) > 0 {
		// 统一执行续费和数据库更新操作
		if err := t.EditExpired(ctx, renewalRequests); err != nil {
			allErrors = multierror.Append(allErrors, fmt.Errorf("批量更新已恢复服务的数据库状态失败: %w", err))
		}
	}


	return allErrors.ErrorOrNil()
}

//对于大于7天的药进行数据情侣操作