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" "net" "sort" "strconv" "strings" ) type TcpforwardingService interface { GetTcpforwarding(ctx context.Context, req v1.GetForwardingRequest) (v1.TcpForwardingDataRequest, error) AddTcpForwarding(ctx context.Context, req *v1.TcpForwardingRequest) error EditTcpForwarding(ctx context.Context, req *v1.TcpForwardingRequest) error DeleteTcpForwarding(ctx context.Context, req v1.DeleteTcpForwardingRequest) error GetTcpForwardingAllIpsByHostId(ctx context.Context, req v1.GetForwardingRequest) ([]v1.TcpForwardingDataRequest, error) } func NewTcpforwardingService( service *Service, tcpforwardingRepository repository.TcpforwardingRepository, parser ParserService, required RequiredService, crawler CrawlerService, globalRep repository.GlobalLimitRepository, hostRep repository.HostRepository, wafformatter WafFormatterService, ) TcpforwardingService { return &tcpforwardingService{ Service: service, tcpforwardingRepository: tcpforwardingRepository, parser: parser, required: required, crawler: crawler, globalRep: globalRep, hostRep: hostRep, wafformatter: wafformatter, } } type tcpforwardingService struct { *Service tcpforwardingRepository repository.TcpforwardingRepository parser ParserService required RequiredService crawler CrawlerService globalRep repository.GlobalLimitRepository hostRep repository.HostRepository wafformatter WafFormatterService } func (s *tcpforwardingService) GetTcpforwarding(ctx context.Context, req v1.GetForwardingRequest) (v1.TcpForwardingDataRequest, error) { var tcpForwarding model.Tcpforwarding var backend model.TcpForwardingRule var err error g, gCtx := errgroup.WithContext(ctx) g.Go(func() error { res, e := s.tcpforwardingRepository.GetTcpforwarding(gCtx, int64(req.Id)) if e != nil { return fmt.Errorf("GetTcpforwarding failed: %w", e) } if res != nil { tcpForwarding = *res } return nil }) g.Go(func() error { res, e := s.tcpforwardingRepository.GetTcpForwardingIpsByID(gCtx, req.Id) if e != nil { return fmt.Errorf("GetTcpforwardingIps failed: %w", e) } if res != nil { backend = *res } return nil }) if err = g.Wait(); err != nil { return v1.TcpForwardingDataRequest{}, err } return v1.TcpForwardingDataRequest{ Id: tcpForwarding.Id, WafTcpId: tcpForwarding.WafTcpId, Tag: tcpForwarding.Tag, Port: tcpForwarding.Port, Comment: tcpForwarding.Comment, WafGatewayGroupId: tcpForwarding.WafGatewayGroupId, WafTcpLimitRuleId: tcpForwarding.TcpLimitRuleId, CcCount: tcpForwarding.CcCount, CcDuration: tcpForwarding.CcDuration, CcBlockCount: tcpForwarding.CcBlockCount, CcBlockDuration: tcpForwarding.CcBlockDuration, BackendProtocol: tcpForwarding.BackendProtocol, BackendTimeout: tcpForwarding.BackendTimeout, BackendList: backend.BackendList, AllowIpList: backend.AllowIpList, DenyIpList: backend.DenyIpList, AccessRule: backend.AccessRule, }, nil } func (s *tcpforwardingService) require(ctx context.Context,req v1.GlobalRequire) (v1.GlobalRequire, error) { res, err := s.wafformatter.require(ctx, req, "tcp") if err != nil { return v1.GlobalRequire{}, err } return res, nil } func (s *tcpforwardingService) buildWafFormData(req *v1.TcpForwardingDataSend, require v1.GlobalRequire) map[string]interface{} { return map[string]interface{}{ "waf_tcp_id": req.WafTcpId, "tag": require.Tag, "port": req.Port, "waf_gateway_group_id": require.WafGatewayGroupId, "waf_tcp_limit_id": require.LimitRuleId, "cc_count": req.CcCount, "cc_duration": req.CcDuration, "cc_block_count": req.CcBlockCount, "cc_block_duration": req.CcBlockDuration, "backend_protocol": req.BackendProtocol, "backend_timeout": req.BackendTimeout, "comment": req.Comment, "backend_list": req.BackendList, "allow_ip_list": req.AllowIpList, "deny_ip_list": req.DenyIpList, "access_rule": req.AccessRule, } } func (s *tcpforwardingService) buildTcpForwardingModel(req *v1.TcpForwardingDataRequest, ruleId int, require v1.GlobalRequire) *model.Tcpforwarding { return &model.Tcpforwarding{ HostId: require.HostId, WafTcpId: ruleId, Port: req.Port, Tag: require.Tag, Comment: req.Comment, TcpLimitRuleId: require.LimitRuleId, WafGatewayGroupId: require.WafGatewayGroupId, CcCount: req.CcCount, CcDuration: req.CcDuration, CcBlockCount: req.CcBlockCount, CcBlockDuration: req.CcBlockDuration, BackendProtocol: req.BackendProtocol, BackendTimeout: req.BackendTimeout, } } func (s *tcpforwardingService) buildTcpRuleModel(reqData *v1.TcpForwardingDataRequest, require v1.GlobalRequire, localDbId int) *model.TcpForwardingRule { return &model.TcpForwardingRule{ Uid: require.Uid, HostId: require.HostId, TcpId: localDbId, // 关联到本地数据库的主记录 ID BackendList: reqData.BackendList, AllowIpList: reqData.AllowIpList, DenyIpList: reqData.DenyIpList, AccessRule: reqData.AccessRule, } } func (s *tcpforwardingService) prepareWafData(ctx context.Context, req *v1.TcpForwardingRequest) (v1.GlobalRequire, map[string]interface{}, error) { // 1. 获取必要的全局信息 require, err := s.require(ctx, v1.GlobalRequire{ HostId: req.HostId, Uid: req.Uid, Comment: req.TcpForwardingData.Comment, }) if err != nil { return v1.GlobalRequire{}, nil, err } if require.WafGatewayGroupId == 0 || require.LimitRuleId == 0 { return v1.GlobalRequire{}, nil, fmt.Errorf("请先配置实例") } // 2. 将字符串切片拼接成字符串,用于 WAF API backendListStr := strings.Join(req.TcpForwardingData.BackendList, "\n") allowIpListStr := strings.Join(req.TcpForwardingData.AllowIpList, "\n") denyIpListStr := strings.Join(req.TcpForwardingData.DenyIpList, "\n") // 3. 创建用于构建 WAF 表单的数据结构 formDataBase := v1.TcpForwardingDataSend{ Tag: require.Tag, WafTcpId: req.TcpForwardingData.WafTcpId, WafGatewayGroupId: require.WafGatewayGroupId, WafTcpLimitRuleId: require.LimitRuleId, Port: req.TcpForwardingData.Port, CcCount: req.TcpForwardingData.CcCount, CcDuration: req.TcpForwardingData.CcDuration, CcBlockCount: req.TcpForwardingData.CcBlockCount, CcBlockDuration: req.TcpForwardingData.CcBlockDuration, BackendProtocol: req.TcpForwardingData.BackendProtocol, BackendTimeout: req.TcpForwardingData.BackendTimeout, BackendList: backendListStr, AllowIpList: allowIpListStr, DenyIpList: denyIpListStr, AccessRule: req.TcpForwardingData.AccessRule, Comment: req.TcpForwardingData.Comment, } // 4. 构建 WAF 表单数据映射 formData := s.buildWafFormData(&formDataBase, require) return require, formData, nil } func (s *tcpforwardingService) AddTcpForwarding(ctx context.Context, req *v1.TcpForwardingRequest) 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 } wafTcpId, err := s.wafformatter.sendFormData(ctx, "admin/info/waf_tcp/new", "admin/new/waf_tcp", formData) if err != nil { return err } // 异步任务:将IP添加到白名单 var ips []string if req.TcpForwardingData.BackendList != nil { for _, v := range req.TcpForwardingData.BackendList { ip, _, err := net.SplitHostPort(v) if err != nil { return err } ips = append(ips, ip) } ips = append(ips, req.TcpForwardingData.AllowIpList...) go s.wafformatter.PublishIpWhitelistTask(ips, "add") } tcpModel := s.buildTcpForwardingModel(&req.TcpForwardingData, wafTcpId, require) id, err := s.tcpforwardingRepository.AddTcpforwarding(ctx, tcpModel) if err != nil { return err } TcpRuleModel := s.buildTcpRuleModel(&req.TcpForwardingData, require, id) if _, err = s.tcpforwardingRepository.AddTcpforwardingIps(ctx, *TcpRuleModel); err != nil { return err } return nil } func (s *tcpforwardingService) EditTcpForwarding(ctx context.Context, req *v1.TcpForwardingRequest) error { WafTcpId, err := s.tcpforwardingRepository.GetTcpforwardingWafTcpIdById(ctx, req.TcpForwardingData.Id) if err != nil { return err } req.TcpForwardingData.WafTcpId = WafTcpId require, formData, err := s.prepareWafData(ctx, req) if err != nil { return err } _, err = s.wafformatter.sendFormData(ctx, "admin/info/waf_tcp/edit?&__goadmin_edit_pk="+strconv.Itoa(req.TcpForwardingData.WafTcpId), "admin/edit/waf_tcp", formData) if err != nil { return err } // 异步任务:将IP添加到白名单 ipData, err := s.tcpforwardingRepository.GetTcpForwardingIpsByID(ctx, req.TcpForwardingData.Id) if err != nil { return err } addedIps, removedIps, err := s.wafformatter.WashEditWafIp(ctx,req.TcpForwardingData.BackendList,req.TcpForwardingData.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") } tcpModel := s.buildTcpForwardingModel(&req.TcpForwardingData, req.TcpForwardingData.WafTcpId, require) tcpModel.Id = req.TcpForwardingData.Id if err = s.tcpforwardingRepository.EditTcpforwarding(ctx, tcpModel); err != nil { return err } TcpRuleModel := s.buildTcpRuleModel(&req.TcpForwardingData, require, req.TcpForwardingData.Id) if err = s.tcpforwardingRepository.EditTcpforwardingIps(ctx, *TcpRuleModel); err != nil { return err } return nil } func (s *tcpforwardingService) DeleteTcpForwarding(ctx context.Context, req v1.DeleteTcpForwardingRequest) error { for _, Id := range req.Ids { wafTcpId, err := s.tcpforwardingRepository.GetTcpforwardingWafTcpIdById(ctx, Id) if err != nil { return err } _, err = s.crawler.DeleteRule(ctx, wafTcpId, "admin/delete/waf_tcp?page=1&__pageSize=1000000&__sort=waf_tcp_id&__sort_type=desc") if err != nil { return err } // 删除白名单 var ips []string ipData, err := s.tcpforwardingRepository.GetTcpForwardingIpsByID(ctx, Id) if err != nil { return err } 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.tcpforwardingRepository.DeleteTcpforwarding(ctx, int64(Id)); err != nil { return err } if err = s.tcpforwardingRepository.DeleteTcpForwardingIpsById(ctx, Id); err != nil { return err } } return nil } func (s *tcpforwardingService) GetTcpForwardingAllIpsByHostId(ctx context.Context, req v1.GetForwardingRequest) ([]v1.TcpForwardingDataRequest, error) { type CombinedResult struct { Id int Forwarding *model.Tcpforwarding BackendRule *model.TcpForwardingRule Err error // 如果此ID的处理出错,则携带错误 } g,gCtx := errgroup.WithContext(ctx) ids, err := s.tcpforwardingRepository.GetTcpForwardingAllIdsByID(gCtx, req.HostId) if err != nil { return nil, fmt.Errorf("GetTcpForwardingAllIds failed: %w", err) } if len(ids) == 0 { return nil, nil } resChan := make(chan CombinedResult, len(ids)) g.Go(func() error { for _, idVal := range ids { currentID := idVal g.Go(func() error { var wf *model.Tcpforwarding var bk *model.TcpForwardingRule var localErr error wf, localErr = s.tcpforwardingRepository.GetTcpforwarding(gCtx, int64(currentID)) if localErr != nil { resChan <- CombinedResult{Id: currentID, Err: localErr} return localErr } bk, localErr = s.tcpforwardingRepository.GetTcpForwardingIpsByID(gCtx, currentID) if localErr != nil { resChan <- CombinedResult{Id: currentID, Err: localErr} return localErr } resChan <- CombinedResult{Id: currentID, Forwarding: wf, BackendRule: bk} return nil }) } return nil }) groupErr := g.Wait() close(resChan) if groupErr != nil { return nil, groupErr } res := make([]v1.TcpForwardingDataRequest, 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.TcpForwardingDataRequest{ Id: r.Forwarding.Id, Port: r.Forwarding.Port, CcCount: r.Forwarding.CcCount, CcDuration: r.Forwarding.CcDuration, CcBlockCount: r.Forwarding.CcBlockCount, CcBlockDuration: r.Forwarding.CcBlockDuration, BackendProtocol: r.Forwarding.BackendProtocol, BackendTimeout: r.Forwarding.BackendTimeout, 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 }