proxy-platform/cmd/agent/main.go

package main

import (
	"context"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"proxy-platform/internal/agent"
	"proxy-platform/internal/config"
	"proxy-platform/internal/socks5"
	"proxy-platform/internal/unlock"
	"proxy-platform/internal/warp"

	"go.uber.org/zap"
)

func main() {
	// 加载配置
	cfg, err := config.LoadAgent("configs/agent.yaml")
	if err != nil {
		log.Fatalf("加载配置失败: %v", err)
	}

	// 初始化日志
	logger, err := zap.NewProduction()
	if err != nil {
		log.Fatalf("初始化日志失败: %v", err)
	}
	defer logger.Sync()

	logger.Info("节点 Agent 启动",
		zap.String("node_id", cfg.Agent.NodeID),
		zap.String("name", cfg.Agent.Name),
	)

	// 初始化 WARP 客户端
	warpClient := warp.NewClient(
		cfg.WARP.SOCKS5Port,
		cfg.WARP.RefreshCooldown,
		cfg.WARP.MaxRefreshRetries,
		cfg.WARP.RefreshRetryDelayMin,
		cfg.WARP.RefreshRetryDelayMax,
		logger,
	)

	// 连接 WARP
	if cfg.WARP.Enabled {
		logger.Info("连接 WARP...")
		if err := warpClient.Connect(context.Background()); err != nil {
			logger.Fatal("连接 WARP 失败", zap.Error(err))
		}
		logger.Info("WARP 连接成功")

		// 获取当前 IP
		ip, err := warpClient.GetCurrentIP(context.Background())
		if err != nil {
			logger.Error("获取 IP 失败", zap.Error(err))
		} else {
			logger.Info("当前 IP", zap.String("ip", ip))
		}
	}

	// 初始化解锁检测器
	detector := unlock.NewDetector(
		"127.0.0.1",
		cfg.WARP.SOCKS5Port,
		30,
		logger,
	)

	// 初始化 Agent 客户端
	agentClient := agent.NewClient(
		cfg.Scheduler.Host,
		cfg.Scheduler.APIKey,
		cfg.Agent.NodeID,
		cfg.Agent.Name,
		cfg.Agent.Region,
		time.Duration(cfg.Scheduler.HeartbeatInterval)*time.Second,
		time.Duration(cfg.Scheduler.ReportInterval)*time.Second,
		logger,
	)

	// 初始化 SOCKS5 服务器（供调度中心连接）
	socks5Server := socks5.NewServer(
		cfg.SOCKS5.Host,
		cfg.SOCKS5.Port,
		cfg.SOCKS5.MaxConnections,
		30,
		nil, // 无需认证
		&LocalBackendSelector{warpClient: warpClient},
		logger,
	)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// 启动 SOCKS5 服务器
	go func() {
		if err := socks5Server.Start(ctx); err != nil {
			logger.Error("SOCKS5 服务器错误", zap.Error(err))
		}
	}()

	// 启动心跳上报
	go agentClient.StartHeartbeat(ctx)

	// 启动解锁检测
	go startUnlockCheck(ctx, detector, agentClient, logger, time.Duration(cfg.Unlock.CheckInterval)*time.Second)

	// 等待中断信号
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit

	logger.Info("正在关闭 Agent...")
	cancel()

	// 断开 WARP
	if cfg.WARP.Enabled {
		warpClient.Disconnect(context.Background())
	}

	logger.Info("Agent 已关闭")
}

// LocalBackendSelector 本地后端选择器
type LocalBackendSelector struct {
	warpClient *warp.Client
}

func (s *LocalBackendSelector) SelectBackend(ctx context.Context, targetHost string, targetPort int, services []string) (string, int, error) {
	// 本地处理，通过 WARP 出口
	// 返回 WARP 的 SOCKS5 端口
	return "127.0.0.1", 40000, nil // WARP SOCKS5 端口
}

func (s *LocalBackendSelector) ReleaseBackend(host string, port int, bytesIn, bytesOut int64) {
	// 本地处理，无需释放
}

// startUnlockCheck 启动解锁检测
func startUnlockCheck(ctx context.Context, detector *unlock.Detector, client *agent.Client, logger *zap.Logger, interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	// 立即执行一次
	doUnlockCheck(ctx, detector, client, logger)

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			doUnlockCheck(ctx, detector, client, logger)
		}
	}
}

func doUnlockCheck(ctx context.Context, detector *unlock.Detector, client *agent.Client, logger *zap.Logger) {
	logger.Info("开始解锁检测...")

	services := make([]unlock.ServiceConfig, len(unlock.DefaultServices))
	copy(services, unlock.DefaultServices)

	results := detector.CheckAll(ctx, services)

	unlocks := make(map[string]struct {
		Unlocked bool   `json:"unlocked"`
		Region   string `json:"region"`
	})

	for service, result := range results {
		unlocks[service] = struct {
			Unlocked bool   `json:"unlocked"`
			Region   string `json:"region"`
		}{
			Unlocked: result.Unlocked,
			Region:   result.Region,
		}
		logger.Info("解锁检测结果",
			zap.String("service", service),
			zap.Bool("unlocked", result.Unlocked),
			zap.String("region", result.Region),
		)
	}

	// 上报解锁状态
	client.ReportUnlock(ctx, unlocks)
}