go_gameshield/internal/service/tcpforwarding.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"
	"golang.org/x/sync/errgroup"
	"maps"
	"net"
	"sort"
)

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,
	cdn CdnService,
) TcpforwardingService {
	return &tcpforwardingService{
		Service:                 service,
		tcpforwardingRepository: tcpforwardingRepository,
		parser:                  parser,
		required:                required,
		crawler:                 crawler,
		globalRep:               globalRep,
		hostRep:                 hostRep,
		wafformatter:            wafformatter,
		cdn:                     cdn,
	}
}

type tcpforwardingService struct {
	*Service
	tcpforwardingRepository repository.TcpforwardingRepository
	parser                  ParserService
	required                RequiredService
	crawler                 CrawlerService
	globalRep               repository.GlobalLimitRepository
	hostRep      repository.HostRepository
	wafformatter WafFormatterService
	cdn CdnService
}

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,
		Port:             tcpForwarding.Port,
		Comment:          tcpForwarding.Comment,
		BackendList:       backend.BackendList,
	}, nil
}


func (s *tcpforwardingService) buildTcpForwardingModel(req *v1.TcpForwardingDataRequest, ruleId int, require RequireResponse) *model.Tcpforwarding {
	return &model.Tcpforwarding{
		HostId:  require.HostId,
		CdnWebId: ruleId,
		Port:    req.Port,
		Comment: req.Comment,
	}
}

func (s *tcpforwardingService) buildTcpRuleModel(reqData *v1.TcpForwardingDataRequest, require RequireResponse, localDbId int, cdnOriginIds map[string]int64) *model.TcpForwardingRule {
	return &model.TcpForwardingRule{
		Uid:         require.Uid,
		HostId:      require.HostId,
		TcpId:       localDbId, // 关联到本地数据库的主记录 ID
		CdnOriginIds: cdnOriginIds,
		BackendList: reqData.BackendList,
	}
}

func (s *tcpforwardingService) prepareWafData(ctx context.Context, req *v1.TcpForwardingRequest) (RequireResponse, v1.Website, error) {
	// 1. 获取必要的全局信息
	require, err := s.wafformatter.Require(ctx, v1.GlobalRequire{
		HostId:  req.HostId,
		Uid:     req.Uid,
		Comment: req.TcpForwardingData.Comment,
	})
	if err != nil {
		return RequireResponse{}, v1.Website{}, err
	}
	if require.GatewayGroupId == 0 || require.Uid == 0 {
		return RequireResponse{}, v1.Website{}, fmt.Errorf("请先配置实例")
	}
	var jsonData v1.TypeJSON
	jsonData.IsOn = true
	for _, v := range require.GatewayIps {
		jsonData.Listen = append(jsonData.Listen, v1.Listen{
			Protocol: "tcp",
			Host:     v,
			Port: 		req.TcpForwardingData.Port,
		})
	}

	byteData, err := json.Marshal(jsonData)
	if err != nil {
		return RequireResponse{}, v1.Website{}, err
	}

	formData := v1.Website{
		UserId:         int64(require.CdnUid),
		Type:           "tcpProxy",
		Name:           require.Tag,
		Description:    req.TcpForwardingData.Comment,
		TcpJSON:        byteData,
		ServerGroupIds: []int64{int64(require.GroupId)},
		UserPlanId: int64(require.RuleId),
		NodeClusterId:  1,
	}
	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
	}

	tcpId, err := s.cdn.CreateWebsite(ctx, formData)
	if err != nil {
		return err
	}

	// 添加源站
	cdnOriginIds := make(map[string]int64)
	for _, v := range req.TcpForwardingData.BackendList{
		id, err := s.wafformatter.AddOrigin(ctx, v1.WebJson{
			ApiType:     "tcp",
			BackendList: v,
			Comment:     req.TcpForwardingData.Comment,
		})
		if err != nil {
			return err
		}
		cdnOriginIds[v] = id
	}

	// 添加源站到网站
	for _, v := range cdnOriginIds {
		err = s.cdn.AddServerOrigin(ctx, tcpId, v)
		if err != nil {
			return err
		}
	}



	tcpModel := s.buildTcpForwardingModel(&req.TcpForwardingData, int(tcpId), require)

	id, err := s.tcpforwardingRepository.AddTcpforwarding(ctx, tcpModel)
	if err != nil {
		return  err
	}
	TcpRuleModel := s.buildTcpRuleModel(&req.TcpForwardingData, require, id, cdnOriginIds)
	if _, err = s.tcpforwardingRepository.AddTcpforwardingIps(ctx, *TcpRuleModel); 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)
		}
		go s.wafformatter.PublishIpWhitelistTask(ips, "add","","white")
	}


	return  nil
}

func (s *tcpforwardingService) EditTcpForwarding(ctx context.Context, req *v1.TcpForwardingRequest) error {
	require, formData, err := s.prepareWafData(ctx, req)
	if err != nil {
		return  err
	}

	oldData, err := s.tcpforwardingRepository.GetTcpforwarding(ctx, int64(req.TcpForwardingData.Id))
	if err != nil {
		return err
	}

	//修改网站端口
	if oldData.Port != req.TcpForwardingData.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.TcpForwardingData.Comment {
		err = s.cdn.EditServerBasic(ctx, int64(oldData.CdnWebId), require.Tag)
		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, ipData.BackendList)
	if err != nil {
		return err
	}
	if len(addedIps) > 0 {
		go s.wafformatter.PublishIpWhitelistTask(addedIps, "add","","white")
	}
	if len(removedIps) > 0 {
		go s.wafformatter.PublishIpWhitelistTask(removedIps, "del","0","white")
	}





	//修改源站
	addOrigins, delOrigins := s.wafformatter.findIpDifferences(ipData.BackendList, req.TcpForwardingData.BackendList)
	addedIds := make(map[string]int64)
	for _, v := range addOrigins {
		id, err := s.wafformatter.AddOrigin(ctx,v1.WebJson{
			ApiType: "tcp",
			BackendList: v,
			Comment: req.TcpForwardingData.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)
			}
		}
	}

	tcpModel := s.buildTcpForwardingModel(&req.TcpForwardingData,oldData.CdnWebId, 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, ipData.CdnOriginIds)
	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 {
		oldData, err := s.tcpforwardingRepository.GetTcpforwarding(ctx, int64(Id))
		if err != nil {
			return err
		}

		err = s.cdn.DelServer(ctx, int64(oldData.CdnWebId))
		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)
		if err != nil {
			return err
		}
		if len(ips) > 0 {
			go s.wafformatter.PublishIpWhitelistTask(ips, "del","0","white")
		}


		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,
			Comment: r.Forwarding.Comment,

		}
		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

}