go_gameshield/internal/service/generateip.go

package service

import (
	"context"
	"fmt"
	"github.com/go-nunu/nunu-layout-advanced/internal/repository"
	"github.com/sourcegraph/conc/pool"
	"math/rand"
	"strconv"
	"strings"
	"sync"
	"time"
)

type GenerateIpService interface {
	GenerateMultipleLoopbackIps(ctx context.Context, count ...int) ([]string, error)
}

func NewGenerateIpService(
	service *Service,
	gameShieldRepository repository.GameShieldRepository,
) GenerateIpService {
	return &generateIpService{
		Service:              service,
		gameShieldRepository: gameShieldRepository,
	}
}

type generateIpService struct {
	*Service
	gameShieldRepository repository.GameShieldRepository
}

// GenerateMultipleLoopbackIps 生成多个唯一127.0.0.0/8范围IP地址
// count 指定需要生成的IP地址数量，默认为10
// 返回生成的唯一IP地址数组和可能的错误
func (s *generateIpService) GenerateMultipleLoopbackIps(ctx context.Context, count ...int) ([]string, error) {
	requiredCount := 10
	if len(count) > 0 && count[0] > 0 {
		requiredCount = count[0]
	}

	maxAttempts := 1000 // 最大尝试次数
	attempts := 0

	// 使用互斥锁保护结果集
	var mu sync.Mutex
	collected := make([]string, 0, requiredCount)
	checkedSegments := make(map[string]struct{}, maxAttempts)

	// 使用一个固定的随机源以提高性能
	rnd := rand.New(rand.NewSource(time.Now().UnixNano()))

	for len(collected) < requiredCount && attempts < maxAttempts {
		// 随机生成多个x/z段组合，批量处理
		candidates := make([]struct{ x, z int }, 0, 5)
		for i := 0; i < 5 && len(candidates) < 5; i++ {
			x := rnd.Intn(256)
			z := rnd.Intn(256)
			segmentKey := fmt.Sprintf("%d-%d", x, z)

			// 跳过已检查过的段
			mu.Lock()
			_, exists := checkedSegments[segmentKey]
			if !exists {
				checkedSegments[segmentKey] = struct{}{}
				mu.Unlock()
				candidates = append(candidates, struct{ x, z int }{x, z})
			} else {
				mu.Unlock()
			}
		}

		// 使用conc/pool包并发处理候选IP段
		// 创建一个工作池，限制并发数为4，并收集错误
		p := pool.New().WithMaxGoroutines(4).WithErrors()

		// 存储结果
		var results [][]string

		// 对每个候选项并发处理
		for _, candidate := range candidates {
			// 创建副本避免闭包捕获循环变量
			candidate := candidate

			p.Go(func() error {
				if ctx.Err() != nil {
					return ctx.Err()
				}

				x, z := candidate.x, candidate.z
				prefix := fmt.Sprintf("127.%d.%d.", x, z)

				// 查询数据库中该前缀已存在的IP
				existingIps, err := s.gameShieldRepository.GetGameShieldExistingIps(ctx, prefix)
				if err != nil {
					// 返回错误以便收集
					return fmt.Errorf("获取IP %s 前缀已存在的记录失败: %w", prefix, err)
				}

				// 使用map标记已存在的第四段
				existingFourth := make(map[int]bool, len(existingIps))
				for _, ip := range existingIps {
					parts := strings.Split(ip, ".")
					if len(parts) == 4 {
						fourth, err := strconv.Atoi(parts[3])
						if err == nil {
							// 排除127.0.0.1的情况
							if x == 0 && z == 0 && fourth == 1 {
								continue
							}
							existingFourth[fourth] = true
						}
					}
				}

				// 快速估算可用数量
				availableCount := 256 - len(existingFourth)
				// 排除127.0.0.1的特殊情况
				if x == 0 && z == 0 && !existingFourth[1] {
					availableCount--
				}

				// 如果可用数量太少，可以跳过
				if availableCount <= 0 {
					return nil
				}

				// 预分配足够大的切片避免多次扩容
				available := make([]string, 0, availableCount)

				// 如果可用数量很多，可以只收集部分，避免不必要的工作
				collectLimit := min(availableCount, requiredCount*2)
				collectCount := 0

				// 随机选择起始点，避免每次都从0开始
				startPoint := rnd.Intn(256)
				for i := 0; i < 256 && collectCount < collectLimit; i++ {
					fourth := (startPoint + i) % 256

					if existingFourth[fourth] {
						continue
					}

					// 确保不生成127.0.0.1
					if x == 0 && z == 0 && fourth == 1 {
						continue
					}

					available = append(available, fmt.Sprintf("127.%d.%d.%d", x, z, fourth))
					collectCount++
				}

				// 随机打乱结果
				for i := len(available) - 1; i > 0; i-- {
					j := rnd.Intn(i + 1)
					available[i], available[j] = available[j], available[i]
				}

				// 将结果添加到结果集
				mu.Lock()
				results = append(results, available)
				mu.Unlock()

				return nil
			})
		}

		// 等待所有goroutine完成并检查错误
		if err := p.Wait(); err != nil {
			// 这里可以选择记录错误但继续尝试，或者直接返回错误
			// 我们选择只有当收集的IP不足时才返回错误
			if len(collected) < requiredCount {
				// 如果已经是最后一次尝试且收集的IP不足，返回错误
				if attempts >= maxAttempts-1 {
					return nil, fmt.Errorf("生成IP地址失败: %w", err)
				}
				// 否则继续尝试下一批候选IP段
			}
		}

		// 收集结果
		for _, ips := range results {
			mu.Lock()
			for _, ip := range ips {
				if len(collected) >= requiredCount {
					break
				}
				collected = append(collected, ip)
			}
			mu.Unlock()
		}

		attempts++

		// 如果已经收集足够的IP，跳出循环
		mu.Lock()
		hasEnough := len(collected) >= requiredCount
		mu.Unlock()
		if hasEnough {
			break
		}
	}

	mu.Lock()
	defer mu.Unlock()

	if len(collected) < requiredCount {
		return nil, fmt.Errorf("无法生成足够的唯一IP地址，地址池可能已耗尽，仅生成了%d个", len(collected))
	}

	// 最终随机排序
	for i := len(collected) - 1; i > 0; i-- {
		j := rand.Intn(i + 1)
		collected[i], collected[j] = collected[j], collected[i]
	}

	return collected[:requiredCount], nil
}