package waf import ( "context" "encoding/json" "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/repository/api/waf" "github.com/go-nunu/nunu-layout-advanced/internal/service" "github.com/go-nunu/nunu-layout-advanced/internal/service/api/flexCdn" "golang.org/x/sync/errgroup" "maps" "net" "sort" ) type UdpForWardingService interface { GetUdpForWarding(ctx context.Context,req v1.GetForwardingRequest) (v1.UdpForwardingDataRequest, error) AddUdpForwarding(ctx context.Context, req *v1.UdpForwardingRequest) (int, error) EditUdpForwarding(ctx context.Context, req *v1.UdpForwardingRequest) error DeleteUdpForwarding(ctx context.Context, req v1.DeleteUdpForwardingRequest) error GetUdpForwardingWafUdpAllIps(ctx context.Context, req v1.GetForwardingRequest) ([]v1.UdpForwardingDataRequest, error) } func NewUdpForWardingService( service *service.Service, udpForWardingRepository waf.UdpForWardingRepository, required service.RequiredService, parser service.ParserService, crawler service.CrawlerService, globalRep waf.GlobalLimitRepository, hostRep repository.HostRepository, wafformatter WafFormatterService, cdn flexCdn.CdnService, proxy flexCdn.ProxyService, ) UdpForWardingService { return &udpForWardingService{ Service: service, udpForWardingRepository: udpForWardingRepository, required: required, parser: parser, crawler: crawler, globalRep: globalRep, hostRep: hostRep, wafformatter: wafformatter, cdn: cdn, proxy: proxy, } } type udpForWardingService struct { *service.Service udpForWardingRepository waf.UdpForWardingRepository required service.RequiredService parser service.ParserService crawler service.CrawlerService globalRep waf.GlobalLimitRepository hostRep repository.HostRepository wafformatter WafFormatterService cdn flexCdn.CdnService proxy flexCdn.ProxyService } 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, Port: udpForWarding.Port, BackendList: backend.BackendList, Comment: udpForWarding.Comment, Proxy: udpForWarding.Proxy, }, nil } func (s *udpForWardingService) buildUdpForwardingModel(req *v1.UdpForwardingDataRequest, ruleId int, require RequireResponse) *model.UdpForWarding { return &model.UdpForWarding{ HostId: require.HostId, CdnWebId: ruleId, Port: req.Port, Comment: req.Comment, Proxy: req.Proxy, } } func (s *udpForWardingService) buildUdpRuleModel(reqData *v1.UdpForwardingDataRequest, require RequireResponse, localDbId int, cdnOriginIds map[string]int64) *model.UdpForwardingRule { return &model.UdpForwardingRule{ Uid: require.Uid, HostId: require.HostId, UdpId: localDbId, // 关联到本地数据库的主记录 ID CdnOriginIds: cdnOriginIds, BackendList: reqData.BackendList, } } func (s *udpForWardingService) prepareWafData(ctx context.Context, req *v1.UdpForwardingRequest) (RequireResponse, v1.WebsiteSend, error) { require, err := s.wafformatter.Require(ctx, v1.GlobalRequire{ HostId: req.HostId, Uid: req.Uid, Comment: req.UdpForwardingData.Comment, }) if err != nil { return RequireResponse{}, v1.WebsiteSend{}, err } if require.Uid == 0 { return RequireResponse{}, v1.WebsiteSend{}, fmt.Errorf("请先配置实例") } var jsonData v1.TypeJSON jsonData.IsOn = true for _, v := range require.GatewayIps { jsonData.Listen = append(jsonData.Listen, v1.Listen{ Protocol: "udp", Host: v, Port: req.UdpForwardingData.Port, }) } byteData, err := json.Marshal(jsonData) if err != nil { return RequireResponse{}, v1.WebsiteSend{}, err } formData := v1.WebsiteSend{ UserId: int64(require.CdnUid), Type: "udpProxy", Name: require.Tag, Description: req.UdpForwardingData.Comment, UdpJSON: byteData, ServerGroupIds: []int64{int64(require.GroupId)}, NodeClusterId: 2, } return require, formData, nil } func (s *udpForWardingService) AddUdpForwarding(ctx context.Context, req *v1.UdpForwardingRequest) (int, error) { require, formData, err := s.prepareWafData(ctx, req) if err != nil { return 0, err } err = s.wafformatter.validateWafPortCount(ctx, require.HostId) if err != nil { return 0, err } // 验证端口重复 err = s.wafformatter.VerifyPort(ctx, "udp", int64(req.UdpForwardingData.Id), req.UdpForwardingData.Port, int64(require.HostId), "") if err != nil { return 0, err } udpId, err := s.cdn.CreateWebsite(ctx, formData) if err != nil { return 0, err } // 添加源站 cdnOriginIds := make(map[string]int64) for _, v := range req.UdpForwardingData.BackendList { id, err := s.wafformatter.AddOrigin(ctx, v1.WebJson{ ApiType: "udp", BackendList: v, Comment: req.UdpForwardingData.Comment, }) if err != nil { return 0, err } cdnOriginIds[v] = id } // 添加源站到网站 for _, v := range cdnOriginIds { err = s.cdn.AddServerOrigin(ctx, udpId, v) if err != nil { return 0, err } } // 开启proxy if req.UdpForwardingData.Proxy { err = s.proxy.EditProxy(ctx,udpId, v1.ProxyProtocolJSON{ IsOn: true, Version: 1, }) if err != nil { return 0, err } } udpModel := s.buildUdpForwardingModel(&req.UdpForwardingData, int(udpId), require) id, err := s.udpForWardingRepository.AddUdpForwarding(ctx, udpModel) if err != nil { return 0, err } udpRuleModel := s.buildUdpRuleModel(&req.UdpForwardingData, require, id, cdnOriginIds) if _, err = s.udpForWardingRepository.AddUdpForwardingIps(ctx, *udpRuleModel); err != nil { return 0, 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 0, err } ips = append(ips, ip) } go s.wafformatter.PublishIpWhitelistTask(ips, "add","","white") } return id, nil } func (s *udpForWardingService) EditUdpForwarding(ctx context.Context, req *v1.UdpForwardingRequest) error { require, formData, err := s.prepareWafData(ctx, req) if err != nil { return err } oldData, err := s.udpForWardingRepository.GetUdpForWarding(ctx, int64(req.UdpForwardingData.Id)) if err != nil { return err } // 验证端口重复 if oldData.Port != req.UdpForwardingData.Port { err = s.wafformatter.VerifyPort(ctx, "udp", int64(req.UdpForwardingData.Id), req.UdpForwardingData.Port, int64(require.HostId), "") if err != nil { return err } } //修改网站端口 if oldData.Port != req.UdpForwardingData.Port { err = s.cdn.EditServerType(ctx, v1.EditWebsite{ Id: int64(oldData.CdnWebId), TypeJSON: formData.TcpJSON, }, "tcp") if err != nil { return err } } //修改网站名字 if oldData.Comment != req.UdpForwardingData.Comment { nodeId, err := s.globalRep.GetNodeId(ctx, oldData.CdnWebId) if err != nil { return err } err = s.cdn.EditServerBasic(ctx, int64(oldData.CdnWebId), require.Tag, nodeId) if err != nil { return err } } //修改Proxy if oldData.Proxy != req.UdpForwardingData.Proxy { err = s.proxy.EditProxy(ctx, int64(oldData.CdnWebId), v1.ProxyProtocolJSON{ IsOn: req.UdpForwardingData.Proxy, Version: 1, }) if err != nil { return err } } // 异步任务:将IP添加到白名单 ipData, err := s.udpForWardingRepository.GetUdpForwardingIpsByID(ctx, req.UdpForwardingData.Id) if err != nil { return err } addedIps, removedIps, err := s.wafformatter.WashEditWafIp(ctx,req.UdpForwardingData.BackendList,ipData.BackendList) if err != nil { return err } if len(addedIps) > 0 { go s.wafformatter.PublishIpWhitelistTask(addedIps, "add","","white") } if len(removedIps) > 0 { ipsToDelist, err := s.wafformatter.WashDelIps(ctx, removedIps) if err != nil { return err } // 4. 如果有需要处理的IP,则批量发布一次任务 if len(ipsToDelist) > 0 { go s.wafformatter.PublishIpWhitelistTask(ipsToDelist, "del", "0", "white") } } //修改源站 addOrigins, delOrigins := s.wafformatter.findIpDifferences(ipData.BackendList, req.UdpForwardingData.BackendList) addedIds := make(map[string]int64) for _, v := range addOrigins { id, err := s.wafformatter.AddOrigin(ctx,v1.WebJson{ ApiType: "udp", BackendList: v, Comment: req.UdpForwardingData.Comment, }) if err != nil { return err } addedIds[v] = id } for _, v := range addedIds { err = s.cdn.AddServerOrigin(ctx, int64(oldData.CdnWebId), v) if err != nil { return err } } maps.Copy(ipData.CdnOriginIds, addedIds) for k, v := range ipData.CdnOriginIds { for _, ip := range delOrigins { if k == ip { err = s.cdn.DelServerOrigin(ctx, int64(oldData.CdnWebId), v) if err != nil { return err } delete(ipData.CdnOriginIds, k) } } } udpModel := s.buildUdpForwardingModel(&req.UdpForwardingData, oldData.CdnWebId, 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, ipData.CdnOriginIds) if err = s.udpForWardingRepository.EditUdpForwardingIps(ctx, *udpRuleModel); err != nil { return err } return nil } func (s *udpForWardingService) DeleteUdpForwarding(ctx context.Context, req v1.DeleteUdpForwardingRequest) error { for _, id := range req.Ids { oldData, err := s.udpForWardingRepository.GetUdpForWarding(ctx, int64(id)) if err != nil { return err } if oldData.HostId != req.HostId { return fmt.Errorf("用户权限不足") } err = s.cdn.DelServer(ctx, int64(oldData.CdnWebId)) if err != nil { return err } // 异步任务:将IP添加到白名单 ipData, err := s.udpForWardingRepository.GetUdpForwardingIpsByID(ctx, id) if err != nil { return err } var ips []string if ipData != nil && len(ipData.BackendList) > 0 { ips, err = s.wafformatter.WashDeleteWafIp(ctx, ipData.BackendList) if err != nil { return err } } if len(ips) > 0 { ipsToDelist, err := s.wafformatter.WashDelIps(ctx, ips) if err != nil { return err } // 4. 如果有需要处理的IP,则批量发布一次任务 if len(ipsToDelist) > 0 { go s.wafformatter.PublishIpWhitelistTask(ipsToDelist, "del", "0", "white") } } 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, Comment: r.Forwarding.Comment, Proxy: r.Forwarding.Proxy, } if r.BackendRule != nil { dataReq.BackendList = r.BackendRule.BackendList } res = append(res, dataReq) } sort.Slice(res, func(i, j int) bool { return res[i].Id > res[j].Id }) return res, nil }