go_gameshield/internal/service/udpforwarding.go

package service

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"
	"golang.org/x/sync/errgroup"
	"sort"
)

type UdpForWardingService interface {
	GetUdpForWarding(ctx context.Context,req v1.GetForwardingRequest) (v1.UdpForwardingDataRequest, error)
	AddUdpForwarding(ctx context.Context, req *v1.UdpForwardingRequest) error
	EditUdpForwarding(ctx context.Context, req *v1.UdpForwardingRequest) error
	DeleteUdpForwarding(ctx context.Context, Ids []int) error
	GetUdpForwardingWafUdpAllIps(ctx context.Context, req v1.GetForwardingRequest) ([]v1.UdpForwardingDataRequest, error)
}

func NewUdpForWardingService(
	service *Service,
	udpForWardingRepository repository.UdpForWardingRepository,
	required RequiredService,
	parser ParserService,
	crawler CrawlerService,
	globalRep repository.GlobalLimitRepository,
	hostRep repository.HostRepository,
	wafformatter WafFormatterService,
) UdpForWardingService {
	return &udpForWardingService{
		Service:                 service,
		udpForWardingRepository: udpForWardingRepository,
		required:                required,
		parser:                  parser,
		crawler:                 crawler,
		globalRep:               globalRep,
		hostRep:                 hostRep,
		wafformatter:            wafformatter,
	}
}

type udpForWardingService struct {
	*Service
	udpForWardingRepository repository.UdpForWardingRepository
	required                RequiredService
	parser                  ParserService
	crawler                 CrawlerService
	globalRep               repository.GlobalLimitRepository
	hostRep      repository.HostRepository
	wafformatter WafFormatterService
}


func (s *udpForWardingService) GetUdpForWarding(ctx context.Context,req v1.GetForwardingRequest) (v1.UdpForwardingDataRequest, error) {
	var udpForWarding model.UdpForWarding
	var backend model.UdpForwardingRule
	var err error
	g, gCtx := errgroup.WithContext(ctx)
	g.Go(func() error {
		res, e := s.udpForWardingRepository.GetUdpForWarding(gCtx, int64(req.Id))
		if e != nil {
			return fmt.Errorf("GetUdpForWarding failed: %w", e)
		}
		if res != nil {
			udpForWarding = *res
		}
		return nil
	})
	g.Go(func() error {
		res, e := s.udpForWardingRepository.GetUdpForwardingIpsByID(gCtx, req.Id)
		if e != nil {
			return fmt.Errorf("GetUdpForWardingByID failed: %w", e)
		}
		if res != nil {
			backend = *res
		}
		return nil
	})
	if err = g.Wait(); err != nil {
		return v1.UdpForwardingDataRequest{}, err
	}

	return v1.UdpForwardingDataRequest{
		Id:                 udpForWarding.Id,
		WafUdpId:           udpForWarding.WafUdpId,
		Tag:                udpForWarding.Tag,
		Port:               udpForWarding.Port,
		WafGatewayGroupId:  udpForWarding.WafGatewayGroupId,
		WafUdpLimitId:      udpForWarding.UdpLimitRuleId,
		CcPacketCount:      udpForWarding.CcPacketCount,
		CcPacketDuration:   udpForWarding.CcPacketDuration,
		CcPacketBlockCount: udpForWarding.CcPacketBlockCount,
		CcPacketBlockDuration: udpForWarding.CcPacketBlockDuration,
		CcCount:            udpForWarding.CcCount,
		CcDuration:         udpForWarding.CcDuration,
		CcBlockCount:       udpForWarding.CcBlockCount,
		CcBlockDuration:    udpForWarding.CcBlockDuration,
		SessionTimeout:     udpForWarding.SessionTimeout,
		BackendList:         backend.BackendList,
		AllowIpList:        backend.AllowIpList,
		DenyIpList:         backend.DenyIpList,
		AccessRule:         backend.AccessRule,
		Comment:            udpForWarding.Comment,

	}, nil
}

func (s *udpForWardingService) buildWafFormData(req *v1.UdpForwardingDataSend, require v1.GlobalRequire) map[string]interface{} {
	return map[string]interface{}{
		"waf_udp_id":           req.WafUdpId,
		"port":                 req.Port,
		"cc_packet_count":      req.CcPacketCount,
		"cc_packet_duration":   req.CcPacketDuration,
		"cc_packet_block_count": req.CcBlockCount,
		"cc_packet_block_duration": req.CcBlockDuration,
		"cc_count":             req.CcCount,
		"cc_duration":          req.CcDuration,
		"cc_block_count":       req.CcBlockCount,
		"cc_block_duration":    req.CcBlockDuration,
		"session_timeout":      req.SessionTimeout,
		"backend_list":         req.BackendList,
		"allow_ip_list":        req.AllowIpList,
		"deny_ip_list":         req.DenyIpList,
		"access_rule":          req.AccessRule,
		"comment":              req.Comment,
	}
}

func (s *udpForWardingService) buildUdpForwardingModel(req *v1.UdpForwardingDataRequest, ruleId int, require v1.GlobalRequire) *model.UdpForWarding {
	return &model.UdpForWarding{
		HostId: 			  require.HostId,
		WafUdpId:             ruleId,
		Port:                 req.Port,
		CcPacketCount:        req.CcPacketCount,
		CcPacketDuration:     req.CcPacketDuration,
		CcPacketBlockCount:   req.CcBlockCount,
		CcPacketBlockDuration: req.CcBlockDuration,
		CcCount:              req.CcCount,
		CcDuration:           req.CcDuration,
		CcBlockCount:         req.CcBlockCount,
		CcBlockDuration:      req.CcBlockDuration,
		SessionTimeout:       req.SessionTimeout,
		Comment:              req.Comment,
	}
}

func (s *udpForWardingService) buildUdpRuleModel(reqData *v1.UdpForwardingDataRequest, require v1.GlobalRequire, localDbId int) *model.UdpForwardingRule {
	return &model.UdpForwardingRule{
		Uid:         require.Uid,
		HostId:      require.HostId,
		UdpId:       localDbId, // 关联到本地数据库的主记录 ID
		BackendList: reqData.BackendList,
		AllowIpList: reqData.AllowIpList,
		DenyIpList:  reqData.DenyIpList,
		AccessRule:  reqData.AccessRule,
	}
}

func (s *udpForWardingService) prepareWafData(ctx context.Context, req *v1.UdpForwardingRequest) (v1.GlobalRequire, map[string]interface{}, error) {
	// 1. 获取必要的全局信息
	//require, err := s.require(ctx, v1.GlobalRequire{
	//	HostId:  req.HostId,
	//	Uid:     req.Uid,
	//	Comment: req.UdpForwardingData.Comment,
	//})
	//if err != nil {
	//	return v1.GlobalRequire{}, nil, err
	//}
	//
	//if require.LimitRuleId == 0 || require.WafGatewayGroupId == 0 {
	//	return v1.GlobalRequire{}, nil, fmt.Errorf("请先配置实例")
	//}
	//
	//// 2. 将字符串切片拼接成字符串，用于 WAF API
	//backendListStr := strings.Join(req.UdpForwardingData.BackendList, "\n")
	//allowIpListStr := strings.Join(req.UdpForwardingData.AllowIpList, "\n")
	//denyIpListStr := strings.Join(req.UdpForwardingData.DenyIpList, "\n")
	//
	//PortInt, err := strconv.Atoi(req.UdpForwardingData.Port)
	//if err != nil {
	//	return v1.GlobalRequire{}, nil, err
	//}
	//
	//// 3. 创建用于构建 WAF 表单的数据结构
	//formDataBase := v1.UdpForwardingDataSend{
	//	Tag:               require.Tag,
	//	WafUdpId:          req.UdpForwardingData.WafUdpId,
	//	WafGatewayGroupId: require.WafGatewayGroupId,
	//	WafUdpLimitId: require.LimitRuleId,
	//	Port:              PortInt,
	//	CcPacketCount:     req.UdpForwardingData.CcPacketCount,
	//	CcPacketDuration:  req.UdpForwardingData.CcPacketDuration,
	//	CcPacketBlockCount: req.UdpForwardingData.CcBlockCount,
	//	CcPacketBlockDuration: req.UdpForwardingData.CcBlockDuration,
	//	CcCount:           req.UdpForwardingData.CcCount,
	//	CcDuration:        req.UdpForwardingData.CcDuration,
	//	CcBlockCount:      req.UdpForwardingData.CcBlockCount,
	//	CcBlockDuration:   req.UdpForwardingData.CcBlockDuration,
	//	SessionTimeout:    req.UdpForwardingData.SessionTimeout,
	//	BackendList:       backendListStr,
	//	AllowIpList:       allowIpListStr,
	//	DenyIpList:        denyIpListStr,
	//	AccessRule:        req.UdpForwardingData.AccessRule,
	//	Comment:           req.UdpForwardingData.Comment,
	//}
	//
	//// 4. 构建 WAF 表单数据映射
	//formData := s.buildWafFormData(&formDataBase, require)

	return v1.GlobalRequire{}, nil, nil
}

func (s *udpForWardingService) AddUdpForwarding(ctx context.Context, req *v1.UdpForwardingRequest) error {
	//require, formData, err := s.prepareWafData(ctx, req)
	//if err != nil {
	//	return err
	//}
	//err = s.wafformatter.validateWafPortCount(ctx, require.HostId)
	//if err != nil {
	//	return err
	//}
	//
	//gatewayIps, _, err := s.wafformatter.GetIp(ctx, require.WafGatewayGroupId)
	//if err != nil {
	//	return err
	//}
	//// 异步任务：将IP添加到白名单
	//var ips []string
	//if req.UdpForwardingData.BackendList != nil {
	//	for _, v := range req.UdpForwardingData.BackendList {
	//		ip, _, err := net.SplitHostPort(v)
	//		if err != nil {
	//			return err
	//		}
	//		ips = append(ips, ip)
	//	}
	//	go s.wafformatter.PublishIpWhitelistTask(ips, "add","")
	//}
	//var accessRuleIps []string
	//if req.UdpForwardingData.AllowIpList != nil {
	//	for _, v := range gatewayIps {
	//		for _, ip := range req.UdpForwardingData.AllowIpList {
	//			if net.ParseIP(ip) != nil {
	//				accessRuleIps = append(accessRuleIps, ip)
	//			}
	//		}
	//		go s.wafformatter.PublishIpWhitelistTask(accessRuleIps, "add",v)
	//	}
	//}
	//
	//
	//udpModel := s.buildUdpForwardingModel(&req.UdpForwardingData, wafUdpId, require)
	//
	//id, err := s.udpForWardingRepository.AddUdpForwarding(ctx, udpModel)
	//if err != nil {
	//	return err
	//}
	//udpRuleModel := s.buildUdpRuleModel(&req.UdpForwardingData, require, id)
	//if _, err = s.udpForWardingRepository.AddUdpForwardingIps(ctx, *udpRuleModel); err != nil {
	//	return err
	//}
	return nil
}

func (s *udpForWardingService) EditUdpForwarding(ctx context.Context, req *v1.UdpForwardingRequest) error {
	//WafUdpId, err := s.udpForWardingRepository.GetUdpForwardingWafUdpIdById(ctx, req.UdpForwardingData.Id)
	//if err != nil {
	//	return err
	//}
	//req.UdpForwardingData.WafUdpId = WafUdpId
	//require, formData, err := s.prepareWafData(ctx, req)
	//if err != nil {
	//	return err
	//}
	//_, err = s.wafformatter.sendFormData(ctx, "admin/info/waf_udp/edit?&__goadmin_edit_pk="+strconv.Itoa(req.UdpForwardingData.WafUdpId), "admin/edit/waf_udp", formData)
	//if err != nil {
	//	return err
	//}
	//
	//// 异步任务：将IP添加到白名单
	//gatewayIps, _, err := s.wafformatter.GetIp(ctx, require.WafGatewayGroupId)
	//if err != nil {
	//	return err
	//}
	//ipData, err := s.udpForWardingRepository.GetUdpForwardingIpsByID(ctx, req.UdpForwardingData.Id)
	//if err != nil {
	//	return err
	//}
	//addedIps, removedIps, addedAllowIps, removedAllowIps, err := s.wafformatter.WashEditWafIp(ctx,req.UdpForwardingData.BackendList,req.UdpForwardingData.AllowIpList,ipData.BackendList,ipData.AllowIpList)
	//if err != nil {
	//	return err
	//}
	//if len(addedIps) > 0 {
	//	go s.wafformatter.PublishIpWhitelistTask(addedIps, "add","")
	//}
	//if len(removedIps) > 0 {
	//	go s.wafformatter.PublishIpWhitelistTask(removedIps, "del","")
	//}
	//
	//if len(addedAllowIps) > 0 {
	//	for _, v := range gatewayIps {
	//		go s.wafformatter.PublishIpWhitelistTask(addedAllowIps, "add",v)
	//	}
	//}
	//if len(removedAllowIps) > 0 {
	//	for _, v := range gatewayIps {
	//		go s.wafformatter.PublishIpWhitelistTask(removedAllowIps, "del",v)
	//	}
	//
	//}
	//
	//
	//udpModel := s.buildUdpForwardingModel(&req.UdpForwardingData, req.UdpForwardingData.WafUdpId, require)
	//udpModel.Id = req.UdpForwardingData.Id
	//if err = s.udpForWardingRepository.EditUdpForwarding(ctx, udpModel); err != nil {
	//	return err
	//}
	//udpRuleModel := s.buildUdpRuleModel(&req.UdpForwardingData, require, req.UdpForwardingData.Id)
	//if err = s.udpForWardingRepository.EditUdpForwardingIps(ctx, *udpRuleModel); err != nil {
	//	return err
	//}
	return nil
}

func (s *udpForWardingService) DeleteUdpForwarding(ctx context.Context, Ids []int) error {
	for _, id := range Ids {
		wafUdpId, err := s.udpForWardingRepository.GetUdpForwardingWafUdpIdById(ctx, id)
		if err != nil {
			return err
		}
		_, err = s.crawler.DeleteRule(ctx, wafUdpId, "admin/delete/waf_udp?page=1&__pageSize=10&__sort=waf_udp_id&__sort_type=desc")
		if err != nil {
			return err
		}

		// 异步任务：将IP添加到白名单
		ipData, err := s.udpForWardingRepository.GetUdpForwardingIpsByID(ctx, id)
		if err != nil {
			return err
		}
		var ips []string
		ips, err = s.wafformatter.WashDeleteWafIp(ctx, ipData.BackendList, ipData.AllowIpList)
		if err != nil {
			return err
		}
		if len(ips) > 0 {
			go s.wafformatter.PublishIpWhitelistTask(ips, "del","")
		}


		if err = s.udpForWardingRepository.DeleteUdpForwarding(ctx, int64(id)); err != nil {
			return err
		}

		if err = s.udpForWardingRepository.DeleteUdpForwardingIpsById(ctx, id); err != nil {
			return err
		}
	}
	return nil
}

func (s *udpForWardingService) GetUdpForwardingWafUdpAllIps(ctx context.Context, req v1.GetForwardingRequest) ([]v1.UdpForwardingDataRequest, error) {
	type CombinedResult struct {
		Id          int
		Forwarding  *model.UdpForWarding
		BackendRule *model.UdpForwardingRule
		Err         error // 如果此ID的处理出错，则携带错误
	}

	g,gCtx := errgroup.WithContext(ctx)
	ids, err := s.udpForWardingRepository.GetUdpForwardingWafUdpAllIds(gCtx, req.HostId)
	if err != nil {
		return nil, fmt.Errorf("GetUdpForwardingWafUdpAllIds failed: %w", err)
	}
	if len(ids) == 0 {
		return nil, nil
	}
	resChan := make(chan CombinedResult, len(ids))

	for _, idVal := range ids {
		currentID := idVal
		g.Go(func() error {
			var wf *model.UdpForWarding
			var bk *model.UdpForwardingRule
			var localErr error
			wf, localErr = s.udpForWardingRepository.GetUdpForWarding(gCtx, int64(currentID))
			if localErr != nil {
				resChan <- CombinedResult{Id: currentID, Err: localErr}
				return localErr
			}
			bk, localErr = s.udpForWardingRepository.GetUdpForwardingIpsByID(gCtx, currentID)
			if localErr != nil {
				resChan <- CombinedResult{Id: currentID, Err: localErr}
				return localErr
			}
			resChan <- CombinedResult{Id: currentID, Forwarding: wf, BackendRule: bk}
			return nil
		})
	}
	groupErr := g.Wait()
	close(resChan)
	if groupErr != nil {
		return nil, groupErr
	}
	res := make([]v1.UdpForwardingDataRequest, 0, len(ids))
	for r := range resChan {
		if r.Err != nil {
			return nil, fmt.Errorf("received error from goroutine for ID %d: %w", r.Id, r.Err)
		}
		if r.Forwarding == nil  {
			return nil, fmt.Errorf("received nil forwarding from goroutine for ID %d", r.Id)
		}

		dataReq := v1.UdpForwardingDataRequest{
			Id: r.Forwarding.Id,
			Port: r.Forwarding.Port,
			CcPacketCount:     r.Forwarding.CcPacketCount,
			CcPacketDuration:  r.Forwarding.CcPacketDuration,
			CcPacketBlockCount: r.Forwarding.CcPacketBlockCount,
			CcPacketBlockDuration: r.Forwarding.CcPacketBlockDuration,
			CcCount:           r.Forwarding.CcCount,
			CcDuration:        r.Forwarding.CcDuration,
			CcBlockCount:      r.Forwarding.CcBlockCount,
			CcBlockDuration:   r.Forwarding.CcBlockDuration,
			SessionTimeout:    r.Forwarding.SessionTimeout,
			Comment:           r.Forwarding.Comment,
		}

		if r.BackendRule != nil {
			dataReq.BackendList = r.BackendRule.BackendList
			dataReq.AllowIpList = r.BackendRule.AllowIpList
			dataReq.DenyIpList = r.BackendRule.DenyIpList
			dataReq.AccessRule = r.BackendRule.AccessRule
		}
		res = append(res, dataReq)
	}

	sort.Slice(res, func(i, j int) bool {
		return res[i].Id > res[j].Id
	})

	return res, nil
}