go_gameshield/internal/service/webforwarding.go

package service

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/pkg/rabbitmq"
	amqp "github.com/rabbitmq/amqp091-go"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
	"sort"
	"strconv"
	"strings"
)

type WebForwardingService interface {
	GetWebForwarding(ctx context.Context, req v1.GetForwardingRequest) (v1.WebForwardingDataRequest, error)
	GetWebForwardingWafWebAllIps(ctx context.Context, req v1.GetForwardingRequest) ([]v1.WebForwardingDataRequest, error)
	AddWebForwarding(ctx context.Context, req *v1.WebForwardingRequest) error
	EditWebForwarding(ctx context.Context, req *v1.WebForwardingRequest) error
	DeleteWebForwarding(ctx context.Context, Ids []int) error
}

func NewWebForwardingService(
	service *Service,
	required RequiredService,
	webForwardingRepository repository.WebForwardingRepository,
	crawler CrawlerService,
	parser ParserService,
	wafformatter WafFormatterService,
	aoDun AoDunService,
	mq *rabbitmq.RabbitMQ,
) WebForwardingService {
	return &webForwardingService{
		Service:                 service,
		webForwardingRepository: webForwardingRepository,
		required:                required,
		parser:                  parser,
		crawler:                 crawler,
		wafformatter:            wafformatter,
		aoDun:                   aoDun,
		mq:                      mq,
	}
}

type webForwardingService struct {
	*Service
	webForwardingRepository repository.WebForwardingRepository
	required                RequiredService
	parser                  ParserService
	crawler                 CrawlerService
	wafformatter            WafFormatterService
	aoDun                   AoDunService
	mq                      *rabbitmq.RabbitMQ
}

func (s *webForwardingService) require(ctx context.Context,req v1.GlobalRequire) (v1.GlobalRequire, error) {
	var err error
	var res v1.GlobalRequire
	g, gCtx := errgroup.WithContext(ctx)

	g.Go(func() error {
		result, e := s.wafformatter.require(gCtx, req, "web")
		if e != nil {
			return e
		}
		res = result
		return nil
	})

	g.Go(func() error {
		e := s.wafformatter.validateWafDomainCount(gCtx, req)
		if e != nil {
			return e
		}
		return nil
	})
	if err 		= g.Wait(); err != nil {
		return v1.GlobalRequire{}, err
	}
	return res, nil
}

func (s *webForwardingService) GetWebForwarding(ctx context.Context, req v1.GetForwardingRequest) (v1.WebForwardingDataRequest, error) {
	var webForwarding model.WebForwarding
	var backend model.WebForwardingRule
	g, gCtx := errgroup.WithContext(ctx)
	g.Go(func() error {
		res, e := s.webForwardingRepository.GetWebForwarding(gCtx, int64(req.Id))
		if e != nil {
			// 直接返回错误，errgroup 会捕获它
			return fmt.Errorf("GetWebForwarding failed: %w", e)
		}
		if res != nil {
			webForwarding = *res
		}
		return nil
	})

	g.Go(func() error {
		res, e := s.webForwardingRepository.GetWebForwardingIpsByID(ctx, req.Id)
		if e != nil {
			return fmt.Errorf("GetWebForwardingByID failed: %w", e)
		}
		if res != nil {
			backend = *res
		}
		return nil
	})

	if err := g.Wait(); err != nil {
		return v1.WebForwardingDataRequest{}, err
	}

	return v1.WebForwardingDataRequest{
		Id:                  webForwarding.Id,
		WafWebId:            webForwarding.WafWebId,
		Tag:                 webForwarding.Tag,
		Port:                webForwarding.Port,
		Domain:              webForwarding.Domain,
		CustomHost:          webForwarding.CustomHost,
		WafWebLimitId:      webForwarding.WebLimitRuleId,
		WafGatewayGroupId:   webForwarding.WafGatewayGroupId,
		CcCount:             webForwarding.CcCount,
		CcDuration:          webForwarding.CcDuration,
		CcBlockCount:        webForwarding.CcBlockCount,
		CcBlockDuration:     webForwarding.CcBlockDuration,
		Cc4xxCount:          webForwarding.Cc4xxCount,
		Cc4xxDuration:       webForwarding.Cc4xxDuration,
		Cc4xxBlockCount:     webForwarding.Cc4xxBlockCount,
		Cc4xxBlockDuration:  webForwarding.Cc4xxBlockDuration,
		Cc5xxCount:          webForwarding.Cc5xxCount,
		Cc5xxDuration:       webForwarding.Cc5xxDuration,
		Cc5xxBlockCount:     webForwarding.Cc5xxBlockCount,
		Cc5xxBlockDuration:  webForwarding.Cc5xxBlockDuration,
		IsHttps:             webForwarding.IsHttps,
		Comment:             webForwarding.Comment,
		BackendList:         backend.BackendList,
		AllowIpList:         backend.AllowIpList,
		DenyIpList:          backend.DenyIpList,
		AccessRule:          backend.AccessRule,
		HttpsKey: 			webForwarding.HttpsKey,
		HttpsCert: 			webForwarding.HttpsCert,
	}, nil
}

// buildWafFormData 辅助函数，用于构建通用的 formData
func (s *webForwardingService) buildWafFormData(req *v1.WebForwardingDataSend, require v1.GlobalRequire) map[string]interface{} {
	// 将BackendList序列化为JSON字符串
	backendJSON, err := json.MarshalIndent(req.BackendList, "", " ")
	var backendStr interface{}
	if err != nil {
		// 如果序列化失败，使用空数组
		backendStr = "[]"
	} else {
		// 成功序列化后，使用JSON字符串
		backendStr = string(backendJSON)
	}

	return map[string]interface{}{
		"waf_web_id":            req.WafWebId,
		"tag":                   require.Tag,
		"port":                  req.Port,
		"domain":                req.Domain,
		"custom_host":           req.CustomHost,
		"waf_gateway_group_id":  require.WafGatewayGroupId,
		"waf_web_limit_id":      require.LimitRuleId,
		"cc_count":              req.CcCount,
		"cc_duration":           req.CcDuration,
		"cc_block_count":        req.CcBlockCount,
		"cc_block_duration":     req.CcBlockDuration,
		"cc_4xx_count":          req.Cc4xxCount,
		"cc_4xx_duration":       req.Cc4xxDuration,
		"cc_4xx_block_count":    req.Cc4xxBlockCount,
		"cc_4xx_block_duration": req.Cc4xxBlockDuration,
		"cc_5xx_count":          req.Cc5xxCount,
		"cc_5xx_duration":       req.Cc5xxDuration,
		"cc_5xx_block_count":    req.Cc5xxBlockCount,
		"cc_5xx_block_duration": req.Cc5xxBlockDuration,
		"backend":          	 backendStr,
		"allow_ip_list":         req.AllowIpList,
		"deny_ip_list":          req.DenyIpList,
		"access_rule":           req.AccessRule,
		"is_https":              req.IsHttps,
		"comment":               req.Comment,
		"https_cert":            req.HttpsCert,
		"https_key":             req.HttpsKey,
	}
}

// buildWebForwardingModel 辅助函数，用于构建通用的 WebForwarding 模型
// ruleId 是从 WAF 系统获取的 ID
func (s *webForwardingService) buildWebForwardingModel(req *v1.WebForwardingDataRequest,ruleId int, require v1.GlobalRequire) *model.WebForwarding {
	return &model.WebForwarding{
		HostId: 			require.HostId,
		WafWebId:           ruleId,
		Tag:                require.Tag,
		Port:               req.Port,
		Domain:             req.Domain,
		CustomHost:         req.CustomHost,
		WafGatewayGroupId:  require.WafGatewayGroupId,
		WebLimitRuleId:     require.LimitRuleId,
		CcCount:            req.CcCount,
		CcDuration:         req.CcDuration,
		CcBlockCount:       req.CcBlockCount,
		CcBlockDuration:    req.CcBlockDuration,
		Cc4xxCount:         req.Cc4xxCount,
		Cc4xxDuration:      req.Cc4xxDuration,
		Cc4xxBlockCount:    req.Cc4xxBlockCount,
		Cc4xxBlockDuration: req.Cc4xxBlockDuration,
		Cc5xxCount:         req.Cc5xxCount,
		Cc5xxDuration:      req.Cc5xxDuration,
		Cc5xxBlockCount:    req.Cc5xxBlockCount,
		Cc5xxBlockDuration: req.Cc5xxBlockDuration,
		IsHttps:            req.IsHttps,
		Comment:            req.Comment,
		HttpsCert:          req.HttpsCert,
		HttpsKey:           req.HttpsKey,
	}
}

func (s *webForwardingService) buildWebRuleModel(reqData *v1.WebForwardingDataRequest, require v1.GlobalRequire, localDbId int) *model.WebForwardingRule {
	return &model.WebForwardingRule{
		Uid:         require.Uid,
		HostId:      require.HostId,
		WebId:       localDbId, // 关联到本地数据库的主记录 ID
		BackendList: reqData.BackendList,
		AllowIpList: reqData.AllowIpList,
		DenyIpList:  reqData.DenyIpList,
		AccessRule:  reqData.AccessRule,
	}
}

func (s *webForwardingService) prepareWafData(ctx context.Context, req *v1.WebForwardingRequest) (v1.GlobalRequire, map[string]interface{}, error) {
	// 1. 获取必要的全局信息
	require, err := s.require(ctx, v1.GlobalRequire{
		HostId:  req.HostId,
		Uid:     req.Uid,
		Comment: req.WebForwardingData.Comment,
		Domain:  req.WebForwardingData.Domain,
	})
	if err != nil {
		return v1.GlobalRequire{}, nil, err
	}
	if require.WafGatewayGroupId == 0 || require.LimitRuleId == 0 {
		return v1.GlobalRequire{}, nil, fmt.Errorf("请先配置实例")
	}

	// 2. 将字符串切片拼接成字符串，用于 WAF API
	allowIpListStr := strings.Join(req.WebForwardingData.AllowIpList, "\n")
	denyIpListStr := strings.Join(req.WebForwardingData.DenyIpList, "\n")

	PortInt, err := strconv.Atoi(req.WebForwardingData.Port)
	if err != nil {
		return v1.GlobalRequire{}, nil, err
	}
	// 3. 创建用于构建 WAF 表单的数据结构
	formDataBase := v1.WebForwardingDataSend{
		Tag:                 require.Tag,
		WafWebId:            req.WebForwardingData.WafWebId,
		WafGatewayGroupId:   require.WafGatewayGroupId,
		WafWebLimitId:   require.LimitRuleId,
		Port:                PortInt,
		Domain:              req.WebForwardingData.Domain,
		CustomHost:          req.WebForwardingData.CustomHost,
		CcCount:             req.WebForwardingData.CcCount,
		CcDuration:          req.WebForwardingData.CcDuration,
		CcBlockCount:        req.WebForwardingData.CcBlockCount,
		CcBlockDuration:     req.WebForwardingData.CcBlockDuration,
		Cc4xxCount:          req.WebForwardingData.Cc4xxCount,
		Cc4xxDuration:       req.WebForwardingData.Cc4xxDuration,
		Cc4xxBlockCount:     req.WebForwardingData.Cc4xxBlockCount,
		Cc4xxBlockDuration:  req.WebForwardingData.Cc4xxBlockDuration,
		Cc5xxCount:          req.WebForwardingData.Cc5xxCount,
		Cc5xxDuration:       req.WebForwardingData.Cc5xxDuration,
		Cc5xxBlockCount:     req.WebForwardingData.Cc5xxBlockCount,
		Cc5xxBlockDuration:  req.WebForwardingData.Cc5xxBlockDuration,
		IsHttps:             req.WebForwardingData.IsHttps,
		BackendList:         req.WebForwardingData.BackendList,
		AllowIpList:         allowIpListStr,
		DenyIpList:          denyIpListStr,
		AccessRule:          req.WebForwardingData.AccessRule,
		Comment:             req.WebForwardingData.Comment,
		HttpsCert:           req.WebForwardingData.HttpsCert,
		HttpsKey:            req.WebForwardingData.HttpsKey,
	}

	// 4. 构建 WAF 表单数据映射
	formData := s.buildWafFormData(&formDataBase, require)

	return require, formData, nil
}

func (s *webForwardingService) AddWebForwarding(ctx context.Context, req *v1.WebForwardingRequest) 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
	}
	wafWebId, err := s.wafformatter.sendFormData(ctx, "admin/info/waf_web/new", "admin/new/waf_web", formData)
	if err != nil {
		return err
	}
	// 异步任务：将域名添加到白名单
	go s.publishDomainWhitelistTask(req.WebForwardingData.Domain, "add")

	webModel := s.buildWebForwardingModel(&req.WebForwardingData, wafWebId, require)

	id, err := s.webForwardingRepository.AddWebForwarding(ctx, webModel)
	if err != nil {
		return err
	}
	webRuleModel := s.buildWebRuleModel(&req.WebForwardingData, require, id)
	//var ips []v1.IpInfo
	//for _, v := range req.WebForwardingData.AllowIpList {
	//	ips = append(ips, v1.IpInfo{
	//		FType:      "allow",
	//		FStartIp:   v,
	//		FEndIp:     v,
	//		FRemark:    "宁波高防IP过白",
	//		FServerIp:  "",
	//	})
	//}
	//err = s.aoDun.AddDomainWhiteList(ctx, req.WebForwardingData.Domain)
	if _, err = s.webForwardingRepository.AddWebForwardingIps(ctx, *webRuleModel); err != nil {
		return err
	}
	return nil
}

func (s *webForwardingService) EditWebForwarding(ctx context.Context, req *v1.WebForwardingRequest) error {
	WafWebId, err := s.webForwardingRepository.GetWebForwardingWafWebIdById(ctx, req.WebForwardingData.Id)
	if err != nil {
		return err
	}
	req.WebForwardingData.WafWebId = WafWebId
	require, formData, err := s.prepareWafData(ctx, req)
	if err != nil {
		return err
	}

	_, err = s.wafformatter.sendFormData(ctx, "admin/info/waf_web/edit?&__goadmin_edit_pk="+strconv.Itoa(req.WebForwardingData.WafWebId), "admin/edit/waf_web", formData)
	if err != nil {
		return err
	}
	webModel := s.buildWebForwardingModel(&req.WebForwardingData, req.WebForwardingData.WafWebId, require)
	webModel.Id = req.WebForwardingData.Id
	if err = s.webForwardingRepository.EditWebForwarding(ctx, webModel); err != nil {
		return err
	}
	webRuleModel := s.buildWebRuleModel(&req.WebForwardingData, require, req.WebForwardingData.Id)
	if err = s.webForwardingRepository.EditWebForwardingIps(ctx, *webRuleModel); err != nil {
		return err
	}
	return nil
}

func (s *webForwardingService) DeleteWebForwarding(ctx context.Context, Ids []int) error {
	for _, Id := range Ids {
		wafWebId, err := s.webForwardingRepository.GetWebForwardingWafWebIdById(ctx, Id)
		if err != nil {
			return err
		}
		_, err = s.crawler.DeleteRule(ctx, wafWebId, "admin/delete/waf_web?page=1&__pageSize=10&__sort=waf_web_id&__sort_type=desc")
		if err != nil {
			return err
		}
		if err = s.webForwardingRepository.DeleteWebForwarding(ctx, int64(Id)); err != nil {
			return err
		}
		if err = s.webForwardingRepository.DeleteWebForwardingIpsById(ctx, Id); err != nil {
			return err
		}
	}

	return nil
}

func (s *webForwardingService) GetWebForwardingWafWebAllIps(ctx context.Context, req v1.GetForwardingRequest) ([]v1.WebForwardingDataRequest, error) {
	type CombinedResult struct {
		Id          int
		Forwarding  *model.WebForwarding
		BackendRule *model.WebForwardingRule
		Err         error // 如果此ID的处理出错，则携带错误
	}
	g, gCtx := errgroup.WithContext(ctx)
	ids, err := s.webForwardingRepository.GetWebForwardingWafWebAllIds(gCtx, req.HostId)
	if err != nil {
		return nil, fmt.Errorf("GetWebForwardingWafWebAllIds failed: %w", err)
	}

	if len(ids) == 0 {
		return nil, nil // 没有ID，直接返回空切片
	}

	// 创建一个通道来接收每个ID的处理结果
	// 通道缓冲区大小设为ID数量，这样发送者不会因为接收者慢而阻塞（在所有goroutine都启动后）
	resultsChan := make(chan CombinedResult, len(ids))

	for _, idVal := range ids {
		currentID := idVal // 捕获循环变量
		g.Go(func() error {
			var wf *model.WebForwarding
			var bk *model.WebForwardingRule
			var localErr error

			// 1. 获取 WebForwarding 信息
			wf, localErr = s.webForwardingRepository.GetWebForwarding(gCtx, int64(currentID))
			if localErr != nil {
				// 发送错误到通道，并由 errgroup 捕获
				// errgroup 会处理第一个非nil错误，并取消其他 goroutine
				resultsChan <- CombinedResult{Id: currentID, Err: fmt.Errorf("GetWebForwarding for id %d failed: %w", currentID, localErr)}
				return localErr // 返回错误给 errgroup
			}
			if wf == nil { // 正常情况下，如果没错误，wf不应为nil，但防御性检查
				localErr = fmt.Errorf("GetWebForwarding for id %d returned nil data without error", currentID)
				resultsChan <- CombinedResult{Id: currentID, Err: localErr}
				return localErr
			}


			// 2. 获取 Backend IP 信息
			// 注意：这里我们允许 GetWebForwardingIpsByID 可能返回 nil 数据（例如没有规则）而不是错误
			// 如果它也可能返回错误，则处理方式与上面类似
			bk, localErr = s.webForwardingRepository.GetWebForwardingIpsByID(gCtx, currentID)
			if localErr != nil {
				// 如果获取IP信息失败是一个致命错误，则也应返回错误
				// 如果允许部分成功（比如有WebForwarding但没有IP信息），则可以不将此视为errgroup的错误
				// 这里假设它也是一个需要errgroup捕获的错误
				resultsChan <- CombinedResult{Id: currentID, Forwarding: wf, Err: fmt.Errorf("GetWebForwardingIpsByID for id %d failed: %w", currentID, localErr)}
				return localErr // 返回错误给 errgroup
			}
			// bk 可能是 nil 如果没有错误且没有规则，这取决于业务逻辑

			// 发送成功的结果到通道
			resultsChan <- CombinedResult{Id: currentID, Forwarding: wf, BackendRule: bk}
			return nil // 此goroutine成功
		})
	}

	// 等待所有goroutine完成
	groupErr := g.Wait()

	// 关闭通道，表示所有发送者都已完成
	// 这一步很重要，这样下面的 range 循环才能正常结束
	close(resultsChan)

	// 如果 errgroup 捕获到任何错误，优先返回该错误
	if groupErr != nil {
		// 虽然errgroup已经出错了，但通道中可能已经有一些结果（来自出错前成功或出错的goroutine）
		// 我们需要排空通道以避免goroutine泄漏（如果它们在发送时阻塞）
		// 但由于我们优先返回groupErr，这些结果将被丢弃。
		// 在这种设计下，通常任何一个子任务失败都会导致整个操作失败。
		return nil, groupErr
	}

	// 如果没有错误，收集所有成功的结果
	finalResults := make([]v1.WebForwardingDataRequest, 0, len(ids))
	for res := range resultsChan {
		// 再次检查通道中的错误，尽管 errgroup 应该已经捕获了
		// 但这是一种更细致的错误处理，以防万一有goroutine在errgroup.Wait()前发送了错误但未被errgroup捕获
		// （理论上，如果goroutine返回了错误，errgroup会处理）
		// 主要目的是处理 res.forwarding 为 nil 的情况 (如果上面允许不返回错误)
		if res.Err != nil {
			// 如果到这里还有错误，说明逻辑可能有问题，或者我们决定忽略某些类型的子错误
			// 在此示例中，因为 g.Wait() 没有错误，所以这里的 res.err 应该是nil
			// 如果不是，那么可能是goroutine在return nil前发送了带有错误的res。
			// 严格来说，如果errgroup没有错误，这里res.err也应该是nil
			// 但以防万一，我们可以记录日志
			return nil, fmt.Errorf("received error from goroutine for ID %d: %w", res.Id, res.Err)
		}
		if res.Forwarding == nil {
			return nil, fmt.Errorf("received nil forwarding from goroutine for ID %d", res.Id)
		}


		dataReq := v1.WebForwardingDataRequest{
			Id:                  res.Forwarding.Id,
			Port:                res.Forwarding.Port,
			Domain:              res.Forwarding.Domain,
			CustomHost:          res.Forwarding.CustomHost,
			CcCount:             res.Forwarding.CcCount,
			CcDuration:          res.Forwarding.CcDuration,
			CcBlockCount:        res.Forwarding.CcBlockCount,
			CcBlockDuration:     res.Forwarding.CcBlockDuration,
			Cc4xxCount:          res.Forwarding.Cc4xxCount,
			Cc4xxDuration:       res.Forwarding.Cc4xxDuration,
			Cc4xxBlockCount:     res.Forwarding.Cc4xxBlockCount,
			Cc4xxBlockDuration:  res.Forwarding.Cc4xxBlockDuration,
			Cc5xxCount:          res.Forwarding.Cc5xxCount,
			Cc5xxDuration:       res.Forwarding.Cc5xxDuration,
			Cc5xxBlockCount:     res.Forwarding.Cc5xxBlockCount,
			Cc5xxBlockDuration:  res.Forwarding.Cc5xxBlockDuration,
			IsHttps:             res.Forwarding.IsHttps,
			Comment:             res.Forwarding.Comment,
			HttpsKey:            res.Forwarding.HttpsKey,
			HttpsCert:           res.Forwarding.HttpsCert,
		}

		if res.BackendRule != nil { // 只有当 BackendRule 存在时才填充相关字段
			dataReq.BackendList = res.BackendRule.BackendList
			dataReq.AllowIpList = res.BackendRule.AllowIpList
			dataReq.DenyIpList = res.BackendRule.DenyIpList
			dataReq.AccessRule = res.BackendRule.AccessRule
		}
		finalResults = append(finalResults, dataReq)
	}

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

// publishDomainWhitelistTask is a helper function to publish domain whitelist tasks to RabbitMQ.
// It can handle different actions like "add" or "del".
func (s *webForwardingService) publishDomainWhitelistTask(domain, action string) {
	// Define message payload, including the action
	type domainTaskPayload struct {
		Domain string `json:"domain"`
		Action string `json:"action"`
	}
	payload := domainTaskPayload{
		Domain: domain,
		Action: action,
	}

	// Serialize the message
	msgBody, err := json.Marshal(payload)
	if err != nil {
		s.logger.Error("Failed to serialize domain whitelist task message", zap.Error(err), zap.String("domain", domain), zap.String("action", action))
		return
	}

	// Get task configuration
	taskCfg, ok := s.mq.GetTaskConfig("domain_whitelist")
	if !ok {
		s.logger.Error("Failed to get 'domain_whitelist' task configuration")
		return
	}

	// Construct the routing key dynamically based on the action
	routingKey := fmt.Sprintf("whitelist.domain.%s", action)

	// Construct the amqp.Publishing message
	publishingMsg := amqp.Publishing{
		ContentType:  "application/json",
		Body:         msgBody,
		DeliveryMode: amqp.Persistent, // Persistent message
	}

	// Publish the message
	err = s.mq.PublishWithCh(taskCfg.Exchange, routingKey, publishingMsg)
	if err != nil {
		s.logger.Error("Failed to publish domain whitelist task to MQ", zap.Error(err), zap.String("domain", domain), zap.String("action", action))
	} else {
		s.logger.Info("Successfully published domain whitelist task to MQ", zap.String("domain", domain), zap.String("action", action))
	}
}