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" "strconv" "strings" ) 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) require(ctx context.Context,req v1.GlobalRequire) (v1.GlobalRequire, error) { res, err := s.wafformatter.require(ctx, req, "udp") if err != nil { return v1.GlobalRequire{}, err } return res, nil } 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, "tag": require.Tag, "port": req.Port, "waf_gateway_group_id": require.WafGatewayGroupId, "waf_udp_limit_id": require.LimitRuleId, "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, Tag: require.Tag, Port: req.Port, WafGatewayGroupId: require.WafGatewayGroupId, UdpLimitRuleId: require.LimitRuleId, 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 require, formData, 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 } wafUdpId, err := s.wafformatter.sendFormData(ctx, "admin/info/waf_udp/new", "admin/new/waf_udp", formData) if err != nil { return err } 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 } 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 } 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 }