demo/examples/custom_dns_reverse_proxy/main.go

package main

import (
	"encoding/json"
	"errors"
	"flag"
	"log/slog"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"

	"github.com/darkit/goproxy/pkg/dns"
	"github.com/darkit/goproxy/pkg/reverse"
)

// 命令行参数
var configFile = flag.String("config", "config.json", "配置文件路径")

// Config 配置文件结构
type Config struct {
	DNS struct {
		Records  map[string][]string `json:"records"`  // 普通记录和泛解析记录
		Fallback bool                `json:"fallback"` // 是否回退到系统DNS
		TTL      int                 `json:"ttl"`      // 缓存TTL，单位为秒
		Strategy string              `json:"strategy"` // 负载均衡策略：round-robin, random, first-available
	} `json:"dns" yaml:"dns" toml:"dns"`
}

func main() {
	flag.Parse()

	// 创建日志记录器
	logger := slog.Default()

	// 加载配置文件
	config := &Config{}
	if *configFile != "" {
		data, err := os.ReadFile(*configFile)
		if err != nil {
			logger.Error("读取配置文件失败", "file", *configFile, "error", err)
			os.Exit(1)
		}

		if err := json.Unmarshal(data, config); err != nil {
			logger.Error("解析配置文件失败", "error", err)
			os.Exit(1)
		}

		logger.Info("成功加载配置文件", "file", *configFile)
	}

	// 选择负载均衡策略
	//var strategy dns.LoadBalanceStrategy
	//switch config.DNS.Strategy {
	//case "random":
	//	strategy = dns.Random
	//case "first-available":
	//	strategy = dns.FirstAvailable
	//default:
	//	strategy = dns.RoundRobin
	//}

	// 创建反向代理配置
	cfg := reverse.DefaultConfig()
	cfg.TargetAddr = "www.shabi.in"
	//cfg.DNSResolver = dns.NewResolver(
	//	dns.WithFallback(config.DNS.Fallback),                  // 是否回退到系统DNS
	//	dns.WithLoadBalanceStrategy(strategy),                  // 使用配置的负载均衡策略
	//	dns.WithTTL(time.Duration(config.DNS.TTL)*time.Second), // 设置DNS缓存TTL
	//) // 创建DNS解析器

	// 添加DNS记录，支持多个地址
	for domain, addrs := range config.DNS.Records {
		for _, addr := range addrs {
			host, port, err := net.SplitHostPort(addr)
			if err != nil {
				// 地址没有端口，将整个地址作为IP使用
				if !dns.IsWildcardDomain(domain) {
					cfg.DNSResolver.Add(domain, addr)
				} else {
					cfg.DNSResolver.AddWildcard(domain, addr)
				}
			} else {
				tPort, err := strconv.Atoi(port)
				if err != nil {
					logger.Error("无效的端口", "domain", domain, "addr", addr, "port", port, "error", err)
					continue
				}
				if !dns.IsWildcardDomain(domain) {
					cfg.DNSResolver.AddWithPort(domain, host, tPort)
				} else {
					cfg.DNSResolver.AddWildcardWithPort(domain, host, tPort)
				}
			}
		}
	}

	// 创建反向代理服务器
	proxy, err := reverse.New(cfg)
	if err != nil {
		logger.Error("创建反向代理服务器失败", "error", err)
		os.Exit(1)
	}

	// 创建HTTP服务器
	server := &http.Server{
		Addr:    cfg.ListenAddr,
		Handler: proxy,
	}

	// 启动服务器
	go func() {
		logger.Info("启动反向代理服务器", "addr", cfg.ListenAddr)
		if err = server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			logger.Error("服务器运行错误", "error", err)
		}
	}()

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

	// 优雅关闭
	logger.Info("正在关闭服务器...")
	if err := server.Close(); err != nil {
		logger.Error("关闭服务器失败", "error", err)
	}
}

// 示例配置文件 config.json:
/*
{
    "dns": {
        "records": {
            "api.example.com": [
                "192.168.1.1:8080",
                "192.168.1.2:8080",
                "192.168.1.3:8080"
            ],
            "*.example.com": [
                "192.168.1.1:8080",
                "192.168.1.2:8080"
            ]
        },
        "fallback": true,
        "ttl": 300,
        "strategy": "round-robin"
    }
}
*/

// 示例配置文件 config.yaml:
/*
dns:
  records:
    api.example.com:
      - 192.168.1.1:8080
      - 192.168.1.2:8080
      - 192.168.1.3:8080
    "*.example.com":
      - 192.168.1.1:8080
      - 192.168.1.2:8080
  fallback: true
  ttl: 300
  strategy: round-robin
*/

// 示例配置文件 config.toml:
/*
[dns]
  [dns.records]
    api.example.com = [
      "192.168.1.1:8080",
      "192.168.1.2:8080",
      "192.168.1.3:8080"
    ]
    "*.example.com" = [
      "192.168.1.1:8080",
      "192.168.1.2:8080"
    ]
  fallback = true
  ttl = 300
  strategy = round-robin
*/